@mastra/upstash 0.10.2 → 0.10.3-alpha.1

package/src/storage/upstash.test.ts

@@ -1,4 +1,5 @@
 import { randomUUID } from 'crypto';
+import { createSampleMessageV2, createSampleThread, createSampleWorkflowSnapshot } from '@internal/storage-test-utils';
 import type { MastraMessageV2 } from '@mastra/core';
 import type { TABLE_NAMES } from '@mastra/core/storage';
 import {
@@ -9,56 +10,13 @@ import {
   TABLE_TRACES,
 } from '@mastra/core/storage';
 import type { WorkflowRunState } from '@mastra/core/workflows';
-import { describe, it, expect, beforeAll, beforeEach, afterAll, vi } from 'vitest';
+import { describe, it, expect, beforeAll, beforeEach, afterAll, vi, afterEach } from 'vitest';
 
 import { UpstashStore } from './index';
 
 // Increase timeout for all tests in this file to 30 seconds
 vi.setConfig({ testTimeout: 200_000, hookTimeout: 200_000 });
 
-const createSampleThread = (date?: Date) => ({
-  id: `thread-${randomUUID()}`,
-  resourceId: `resource-${randomUUID()}`,
-  title: 'Test Thread',
-  createdAt: date || new Date(),
-  updatedAt: date || new Date(),
-  metadata: { key: 'value' },
-});
-
-const createSampleMessage = (threadId: string, content: string = 'Hello'): MastraMessageV2 => ({
-  id: `msg-${randomUUID()}`,
-  role: 'user',
-  threadId,
-  content: { format: 2, parts: [{ type: 'text', text: content }] },
-  createdAt: new Date(),
-  resourceId: `resource-${randomUUID()}`,
-});
-
-const createSampleWorkflowSnapshot = (status: string, createdAt?: Date) => {
-  const runId = `run-${randomUUID()}`;
-  const stepId = `step-${randomUUID()}`;
-  const timestamp = createdAt || new Date();
-  const snapshot: WorkflowRunState = {
-    value: {},
-    context: {
-      [stepId]: {
-        status: status,
-        payload: {},
-        error: undefined,
-        startedAt: timestamp.getTime(),
-        endedAt: new Date(timestamp.getTime() + 15000).getTime(),
-      },
-      input: {},
-    } as WorkflowRunState['context'],
-    serializedStepGraph: [],
-    activePaths: [],
-    suspendedPaths: {},
-    runId,
-    timestamp: timestamp.getTime(),
-  };
-  return { snapshot, runId, stepId };
-};
-
 const createSampleTrace = (name: string, scope?: string, attributes?: Record<string, string>) => ({
   id: `trace-${randomUUID()}`,
   parentSpanId: `span-${randomUUID()}`,
@@ -176,7 +134,7 @@ describe('UpstashStore', () => {
 
     it('should create and retrieve a thread', async () => {
       const now = new Date();
-      const thread = createSampleThread(now);
+      const thread = createSampleThread({ date: now });
 
       const savedThread = await store.saveThread({ thread });
       expect(savedThread).toEqual(thread);
@@ -236,6 +194,23 @@ describe('UpstashStore', () => {
       const retrievedThreads = await store.getThreadsByResourceId({ resourceId });
       expect(retrievedThreads).toHaveLength(total);
     });
+    it('should delete thread and its messages', async () => {
+      const thread = createSampleThread();
+      await store.saveThread({ thread });
+
+      // Add some messages
+      const messages = [createSampleMessageV2({ threadId: thread.id }), createSampleMessageV2({ threadId: thread.id })];
+      await store.saveMessages({ messages, format: 'v2' });
+
+      await store.deleteThread({ threadId: thread.id });
+
+      const retrievedThread = await store.getThreadById({ threadId: thread.id });
+      expect(retrievedThread).toBeNull();
+
+      // Verify messages were also deleted
+      const retrievedMessages = await store.getMessages({ threadId: thread.id });
+      expect(retrievedMessages).toHaveLength(0);
+    });
   });
 
   describe('Date Handling', () => {
@@ -245,7 +220,7 @@ describe('UpstashStore', () => {
 
     it('should handle Date objects in thread operations', async () => {
       const now = new Date();
-      const thread = createSampleThread(now);
+      const thread = createSampleThread({ date: now });
 
       await store.saveThread({ thread });
       const retrievedThread = await store.getThreadById({ threadId: thread.id });
@@ -257,7 +232,7 @@ describe('UpstashStore', () => {
 
     it('should handle ISO string dates in thread operations', async () => {
       const now = new Date();
-      const thread = createSampleThread(now);
+      const thread = createSampleThread({ date: now });
 
       await store.saveThread({ thread });
       const retrievedThread = await store.getThreadById({ threadId: thread.id });
@@ -269,7 +244,7 @@ describe('UpstashStore', () => {
 
     it('should handle mixed date formats in thread operations', async () => {
       const now = new Date();
-      const thread = createSampleThread(now);
+      const thread = createSampleThread({ date: now });
 
       await store.saveThread({ thread });
       const retrievedThread = await store.getThreadById({ threadId: thread.id });
@@ -281,8 +256,8 @@ describe('UpstashStore', () => {
 
     it('should handle date serialization in getThreadsByResourceId', async () => {
       const now = new Date();
-      const thread1 = createSampleThread(now);
-      const thread2 = { ...createSampleThread(now), resourceId: thread1.resourceId };
+      const thread1 = createSampleThread({ date: now });
+      const thread2 = { ...createSampleThread({ date: now }), resourceId: thread1.resourceId };
       const threads = [thread1, thread2];
 
       await Promise.all(threads.map(thread => store.saveThread({ thread })));
@@ -320,9 +295,9 @@ describe('UpstashStore', () => {
 
     it('should save and retrieve messages in order', async () => {
       const messages: MastraMessageV2[] = [
-        createSampleMessage(threadId, 'First'),
-        createSampleMessage(threadId, 'Second'),
-        createSampleMessage(threadId, 'Third'),
+        createSampleMessageV2({ threadId, content: 'First' }),
+        createSampleMessageV2({ threadId, content: 'Second' }),
+        createSampleMessageV2({ threadId, content: 'Third' }),
       ];
 
       await store.saveMessages({ messages: messages, format: 'v2' });
@@ -360,6 +335,140 @@ describe('UpstashStore', () => {
       const retrievedMessages = await store.getMessages({ threadId, format: 'v2' });
       expect(retrievedMessages[0].content).toEqual(messages[0].content);
     });
+
+    describe('getPaginatedMessages', () => {
+      it('should return paginated messages with total count', async () => {
+        const thread = createSampleThread();
+        await store.saveThread({ thread });
+
+        const messages = Array.from({ length: 15 }, (_, i) =>
+          createSampleMessageV2({ threadId: thread.id, content: `Message ${i + 1}` }),
+        );
+
+        await store.saveMessages({ messages, format: 'v2' });
+
+        const page1 = await store.getMessages({
+          threadId: thread.id,
+          page: 0,
+          perPage: 5,
+          format: 'v2',
+        });
+        expect(page1.messages).toHaveLength(5);
+        expect(page1.total).toBe(15);
+        expect(page1.page).toBe(0);
+        expect(page1.perPage).toBe(5);
+        expect(page1.hasMore).toBe(true);
+
+        const page3 = await store.getMessages({
+          threadId: thread.id,
+          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({
+          threadId: thread.id,
+          page: 3,
+          perPage: 5,
+          format: 'v2',
+        });
+        expect(page4.messages).toHaveLength(0);
+        expect(page4.total).toBe(15);
+        expect(page4.hasMore).toBe(false);
+      });
+
+      it('should maintain chronological order in pagination', async () => {
+        const thread = createSampleThread();
+        await store.saveThread({ thread });
+
+        const messages = Array.from({ length: 10 }, (_, i) => {
+          const message = createSampleMessageV2({ threadId: thread.id, content: `Message ${i + 1}` });
+          // Ensure different timestamps
+          message.createdAt = new Date(Date.now() + i * 1000);
+          return message;
+        });
+
+        await store.saveMessages({ messages, format: 'v2' });
+
+        const page1 = await store.getMessages({
+          threadId: thread.id,
+          page: 0,
+          perPage: 3,
+          format: 'v2',
+        });
+
+        // Check that messages are in chronological order
+        for (let i = 1; i < page1.messages.length; i++) {
+          const prevMessage = page1.messages[i - 1] as MastraMessageV2;
+          const currentMessage = page1.messages[i] as MastraMessageV2;
+          expect(new Date(prevMessage.createdAt).getTime()).toBeLessThanOrEqual(
+            new Date(currentMessage.createdAt).getTime(),
+          );
+        }
+      });
+
+      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 });
+
+        const now = new Date();
+        const yesterday = new Date(now.getTime() - 24 * 60 * 60 * 1000);
+        const tomorrow = new Date(now.getTime() + 24 * 60 * 60 * 1000);
+
+        const oldMessages = Array.from({ length: 3 }, (_, i) => {
+          const message = createSampleMessageV2({ threadId: thread.id, content: `Old Message ${i + 1}` });
+          message.createdAt = yesterday;
+          return message;
+        });
+
+        const newMessages = Array.from({ length: 4 }, (_, i) => {
+          const message = createSampleMessageV2({ threadId: thread.id, content: `New Message ${i + 1}` });
+          message.createdAt = tomorrow;
+          return message;
+        });
+
+        await store.saveMessages({ messages: [...oldMessages, ...newMessages], format: 'v2' });
+
+        const recentMessages = await store.getMessages({
+          threadId: thread.id,
+          page: 0,
+          perPage: 10,
+          fromDate: now,
+          format: 'v2',
+        });
+        expect(recentMessages.messages).toHaveLength(4);
+        expect(recentMessages.total).toBe(4);
+      });
+    });
   });
 
   describe('Trace Operations', () => {
@@ -864,4 +973,270 @@ describe('UpstashStore', () => {
       expect(runs.length).toBe(0);
     });
   });
+
+  describe('alterTable (no-op/schemaless)', () => {
+    const TEST_TABLE = 'test_alter_table'; // Use "table" or "collection" as appropriate
+    beforeEach(async () => {
+      await store.clearTable({ tableName: TEST_TABLE as TABLE_NAMES });
+    });
+
+    afterEach(async () => {
+      await store.clearTable({ tableName: TEST_TABLE as TABLE_NAMES });
+    });
+
+    it('allows inserting records with new fields without alterTable', async () => {
+      await store.insert({
+        tableName: TEST_TABLE as TABLE_NAMES,
+        record: { id: '1', name: 'Alice' },
+      });
+      await store.insert({
+        tableName: TEST_TABLE as TABLE_NAMES,
+        record: { id: '2', name: 'Bob', newField: 123 },
+      });
+
+      const row = await store.load<{ id: string; name: string; newField?: number }>({
+        tableName: TEST_TABLE as TABLE_NAMES,
+        keys: { id: '2' },
+      });
+      expect(row?.newField).toBe(123);
+    });
+
+    it('does not throw when calling alterTable (no-op)', async () => {
+      await expect(
+        store.alterTable({
+          tableName: TEST_TABLE as TABLE_NAMES,
+          schema: {
+            id: { type: 'text', primaryKey: true, nullable: false },
+            name: { type: 'text', nullable: true },
+            extra: { type: 'integer', nullable: true },
+          },
+          ifNotExists: [],
+        }),
+      ).resolves.not.toThrow();
+    });
+
+    it('can add multiple new fields at write time', async () => {
+      await store.insert({
+        tableName: TEST_TABLE as TABLE_NAMES,
+        record: { id: '3', name: 'Charlie', age: 30, city: 'Paris' },
+      });
+      const row = await store.load<{ id: string; name: string; age?: number; city?: string }>({
+        tableName: TEST_TABLE as TABLE_NAMES,
+        keys: { id: '3' },
+      });
+      expect(row?.age).toBe(30);
+      expect(row?.city).toBe('Paris');
+    });
+
+    it('can retrieve all fields, including dynamically added ones', async () => {
+      await store.insert({
+        tableName: TEST_TABLE as TABLE_NAMES,
+        record: { id: '4', name: 'Dana', hobby: 'skiing' },
+      });
+      const row = await store.load<{ id: string; name: string; hobby?: string }>({
+        tableName: TEST_TABLE as TABLE_NAMES,
+        keys: { id: '4' },
+      });
+      expect(row?.hobby).toBe('skiing');
+    });
+
+    it('does not restrict or error on arbitrary new fields', async () => {
+      await expect(
+        store.insert({
+          tableName: TEST_TABLE as TABLE_NAMES,
+          record: { id: '5', weirdField: { nested: true }, another: [1, 2, 3] },
+        }),
+      ).resolves.not.toThrow();
+
+      const row = await store.load<{ id: string; weirdField?: any; another?: any }>({
+        tableName: TEST_TABLE as TABLE_NAMES,
+        keys: { id: '5' },
+      });
+      expect(row?.weirdField).toEqual({ nested: true });
+      expect(row?.another).toEqual([1, 2, 3]);
+    });
+  });
+
+  describe('Pagination Features', () => {
+    beforeEach(async () => {
+      // Clear all test data
+      await store.clearTable({ tableName: TABLE_THREADS });
+      await store.clearTable({ tableName: TABLE_MESSAGES });
+      await store.clearTable({ tableName: TABLE_EVALS });
+      await store.clearTable({ tableName: TABLE_TRACES });
+    });
+
+    describe('getEvals with pagination', () => {
+      it('should return paginated evals with total count', async () => {
+        const agentName = 'test-agent';
+        const evals = Array.from({ length: 25 }, (_, i) => createSampleEval(agentName, i % 2 === 0));
+
+        // Insert all evals
+        for (const evalRecord of evals) {
+          await store.insert({
+            tableName: TABLE_EVALS,
+            record: evalRecord,
+          });
+        }
+
+        // Test page-based pagination
+        const page1 = await store.getEvals({ agentName, page: 0, perPage: 10 });
+        expect(page1.evals).toHaveLength(10);
+        expect(page1.total).toBe(25);
+        expect(page1.page).toBe(0);
+        expect(page1.perPage).toBe(10);
+        expect(page1.hasMore).toBe(true);
+
+        const page2 = await store.getEvals({ agentName, page: 1, perPage: 10 });
+        expect(page2.evals).toHaveLength(10);
+        expect(page2.total).toBe(25);
+        expect(page2.hasMore).toBe(true);
+
+        const page3 = await store.getEvals({ agentName, page: 2, perPage: 10 });
+        expect(page3.evals).toHaveLength(5);
+        expect(page3.total).toBe(25);
+        expect(page3.hasMore).toBe(false);
+      });
+
+      it('should support limit/offset pagination', async () => {
+        const agentName = 'test-agent-2';
+        const evals = Array.from({ length: 15 }, () => createSampleEval(agentName));
+
+        for (const evalRecord of evals) {
+          await store.insert({
+            tableName: TABLE_EVALS,
+            record: evalRecord,
+          });
+        }
+
+        // Test offset-based pagination
+        const result1 = await store.getEvals({ agentName, limit: 5, offset: 0 });
+        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 });
+        expect(result2.evals).toHaveLength(5);
+        expect(result2.total).toBe(15);
+        expect(result2.hasMore).toBe(false);
+      });
+
+      it('should filter by type with pagination', async () => {
+        const agentName = 'test-agent-3';
+        const testEvals = Array.from({ length: 10 }, () => createSampleEval(agentName, true));
+        const liveEvals = Array.from({ length: 8 }, () => createSampleEval(agentName, false));
+
+        for (const evalRecord of [...testEvals, ...liveEvals]) {
+          await store.insert({
+            tableName: TABLE_EVALS,
+            record: evalRecord,
+          });
+        }
+
+        const testResults = await store.getEvals({ agentName, type: 'test', page: 0, perPage: 5 });
+        expect(testResults.evals).toHaveLength(5);
+        expect(testResults.total).toBe(10);
+
+        const liveResults = await store.getEvals({ agentName, type: 'live', page: 0, perPage: 5 });
+        expect(liveResults.evals).toHaveLength(5);
+        expect(liveResults.total).toBe(8);
+      });
+    });
+
+    describe('getTracesPaginated', () => {
+      it('should return paginated traces with total count', async () => {
+        const traces = Array.from({ length: 18 }, (_, i) => createSampleTrace(`test-trace-${i}`, 'test-scope'));
+
+        for (const trace of traces) {
+          await store.insert({
+            tableName: TABLE_TRACES,
+            record: trace,
+          });
+        }
+
+        const page1 = await store.getTraces({
+          scope: 'test-scope',
+          page: 0,
+          perPage: 8,
+          returnPaginationResults: true,
+        });
+        expect(page1.traces).toHaveLength(8);
+        expect(page1.total).toBe(18);
+        expect(page1.page).toBe(0);
+        expect(page1.perPage).toBe(8);
+        expect(page1.hasMore).toBe(true);
+
+        const page3 = await store.getTraces({
+          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' }),
+        );
+
+        for (const trace of [...tracesWithAttr, ...tracesWithoutAttr]) {
+          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' },
+          page: 0,
+          perPage: 10,
+          returnPaginationResults: true,
+        });
+        expect(devTraces.traces).toHaveLength(5);
+        expect(devTraces.total).toBe(5);
+        expect(devTraces.hasMore).toBe(false);
+      });
+    });
+
+    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,
+        }));
+
+        for (const thread of threads) {
+          await store.saveThread({ thread });
+        }
+
+        const page1 = await store.getThreadsByResourceId({ resourceId, page: 0, perPage: 7 });
+        expect(page1.threads).toHaveLength(7);
+
+        const page3 = await store.getThreadsByResourceId({ resourceId, page: 2, perPage: 7 });
+        expect(page3.threads).toHaveLength(3);
+
+        const limited = await store.getThreadsByResourceId({ resourceId, page: 1, perPage: 5 });
+        expect(limited.threads).toHaveLength(5);
+      });
+    });
+  });
 });