@mastra/upstash 0.0.0-vnextWorkflows-20250422142014 → 0.0.0-workflow-deno-20250616130925

package/src/storage/upstash.test.ts

@@ -1,5 +1,11 @@
 import { randomUUID } from 'crypto';
-import type { MessageType } from '@mastra/core/memory';
+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 {
   TABLE_MESSAGES,
@@ -9,58 +15,19 @@ 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: 60_000, hookTimeout: 60_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') =>
-  ({
-    id: `msg-${randomUUID()}`,
-    role: 'user',
-    type: 'text',
-    threadId,
-    content: [{ type: 'text', text: content }],
-    createdAt: new Date(),
-  }) as any;
-
-const createSampleWorkflowSnapshot = (status: string, createdAt?: Date) => {
-  const runId = `run-${randomUUID()}`;
-  const stepId = `step-${randomUUID()}`;
-  const timestamp = createdAt || new Date();
-  const snapshot = {
-    result: { success: true },
-    value: {},
-    context: {
-      steps: {
-        [stepId]: {
-          status,
-          payload: {},
-          error: undefined,
-        },
-      },
-      triggerData: {},
-      attempts: {},
-    },
-    activePaths: [],
-    runId,
-    timestamp: timestamp.getTime(),
-  } as WorkflowRunState;
-  return { snapshot, runId, stepId };
-};
-
-const createSampleTrace = (name: string, scope?: string, attributes?: Record<string, string>) => ({
+vi.setConfig({ testTimeout: 200_000, hookTimeout: 200_000 });
+
+const createSampleTrace = (
+  name: string,
+  scope?: string,
+  attributes?: Record<string, string>,
+  createdAt: Date = new Date(),
+) => ({
   id: `trace-${randomUUID()}`,
   parentSpanId: `span-${randomUUID()}`,
   traceId: `trace-${randomUUID()}`,
@@ -68,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 {
@@ -90,7 +57,7 @@ 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(),
   };
 };
 
@@ -170,7 +137,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);
@@ -190,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;
@@ -220,6 +187,56 @@ describe('UpstashStore', () => {
         updated: 'value',
       });
     });
+
+    it('should update thread updatedAt when a message is saved to it', async () => {
+      const thread = createSampleThread();
+      await store.saveThread({ thread });
+
+      // Get the initial thread to capture the original updatedAt
+      const initialThread = await store.getThreadById({ threadId: thread.id });
+      expect(initialThread).toBeDefined();
+      const originalUpdatedAt = initialThread!.updatedAt;
+
+      // Wait a small amount to ensure different timestamp
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      // Create and save a message to the thread
+      const message = createSampleMessageV2({ threadId: thread.id });
+      await store.saveMessages({ messages: [message], format: 'v2' });
+
+      // Retrieve the thread again and check that updatedAt was updated
+      const updatedThread = await store.getThreadById({ threadId: thread.id });
+      expect(updatedThread).toBeDefined();
+      expect(updatedThread!.updatedAt.getTime()).toBeGreaterThan(originalUpdatedAt.getTime());
+    });
+
+    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 }));
+
+      await store.batchInsert({ tableName: TABLE_THREADS, records: threads });
+
+      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', () => {
@@ -229,7 +246,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 });
@@ -241,7 +258,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 });
@@ -253,7 +270,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 });
@@ -265,8 +282,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 })));
@@ -303,17 +320,121 @@ describe('UpstashStore', () => {
     });
 
     it('should save and retrieve messages in order', async () => {
-      const messages = [
-        createSampleMessage(threadId, 'First'),
-        createSampleMessage(threadId, 'Second'),
-        createSampleMessage(threadId, 'Third'),
+      const messages: MastraMessageV2[] = [
+        createSampleMessageV2({ threadId, content: 'First' }),
+        createSampleMessageV2({ threadId, content: 'Second' }),
+        createSampleMessageV2({ threadId, content: 'Third' }),
+      ];
+
+      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 as MessageType[] });
+      await store.saveMessages({ messages: messages, format: 'v2' });
 
-      const retrievedMessages = await store.getMessages({ threadId });
+      const retrievedMessages = await store.getMessages({ threadId: 'thread-one', format: 'v2' });
       expect(retrievedMessages).toHaveLength(3);
-      expect(retrievedMessages.map(m => m.content[0].text)).toEqual(['First', 'Second', 'Third']);
+      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 () => {
@@ -327,21 +448,131 @@ describe('UpstashStore', () => {
           id: 'msg-1',
           threadId,
           role: 'user',
-          type: 'text',
-          content: [
-            { type: 'text', text: 'Message with' },
-            { type: 'code', text: 'code block', language: 'typescript' },
-            { type: 'text', text: 'and more text' },
-          ],
+          content: {
+            format: 2,
+            parts: [
+              { type: 'text', text: 'Message with' },
+              { type: 'code', text: 'code block', language: 'typescript' },
+              { type: 'text', text: 'and more text' },
+            ],
+          },
           createdAt: new Date(),
         },
-      ];
+      ] as MastraMessageV2[];
 
-      await store.saveMessages({ messages: messages as MessageType[] });
+      await store.saveMessages({ messages, format: 'v2' });
 
-      const retrievedMessages = await store.getMessages({ threadId });
+      const retrievedMessages = await store.getMessages({ threadId, format: 'v2' });
       expect(retrievedMessages[0].content).toEqual(messages[0].content);
     });
+
+    describe('getMessagesPaginated', () => {
+      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.getMessagesPaginated({
+          threadId: thread.id,
+          selectBy: { pagination: { 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.getMessagesPaginated({
+          threadId: thread.id,
+          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.getMessagesPaginated({
+          threadId: thread.id,
+          selectBy: { pagination: { 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.getMessagesPaginated({
+          threadId: thread.id,
+          selectBy: { pagination: { 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 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.getMessagesPaginated({
+          threadId: thread.id,
+          selectBy: {
+            pagination: {
+              page: 0,
+              perPage: 10,
+              dateRange: { start: now },
+            },
+          },
+          format: 'v2',
+        });
+        expect(recentMessages.messages).toHaveLength(4);
+        expect(recentMessages.total).toBe(4);
+      });
+    });
   });
 
   describe('Trace Operations', () => {
@@ -448,22 +679,27 @@ describe('UpstashStore', () => {
       const mockSnapshot = {
         value: { step1: 'completed' },
         context: {
-          stepResults: {
-            step1: { status: 'success', payload: { result: 'done' } },
+          step1: {
+            status: 'success',
+            output: { result: 'done' },
+            payload: {},
+            startedAt: new Date().getTime(),
+            endedAt: new Date(Date.now() + 15000).getTime(),
           },
-          attempts: {},
-          triggerData: {},
-        },
+        } as WorkflowRunState['context'],
+        serializedStepGraph: [],
         runId: testRunId,
         activePaths: [],
+        suspendedPaths: {},
         timestamp: Date.now(),
-      } as unknown as WorkflowRunState;
+        status: 'success',
+      };
 
       await store.persistWorkflowSnapshot({
         namespace: testNamespace,
         workflowName: testWorkflow,
         runId: testRunId,
-        snapshot: mockSnapshot,
+        snapshot: mockSnapshot as WorkflowRunState,
       });
 
       const loadedSnapshot = await store.loadWorkflowSnapshot({
@@ -556,8 +792,8 @@ describe('UpstashStore', () => {
       const workflowName1 = 'default_test_1';
       const workflowName2 = 'default_test_2';
 
-      const { snapshot: workflow1, runId: runId1, stepId: stepId1 } = createSampleWorkflowSnapshot('completed');
-      const { snapshot: workflow2, runId: runId2, stepId: stepId2 } = createSampleWorkflowSnapshot('running');
+      const { snapshot: workflow1, runId: runId1, stepId: stepId1 } = createSampleWorkflowSnapshot('success');
+      const { snapshot: workflow2, runId: runId2, stepId: stepId2 } = createSampleWorkflowSnapshot('waiting');
 
       await store.persistWorkflowSnapshot({
         namespace: testNamespace,
@@ -578,17 +814,17 @@ describe('UpstashStore', () => {
       expect(total).toBe(2);
       expect(runs[0]!.workflowName).toBe(workflowName2); // Most recent first
       expect(runs[1]!.workflowName).toBe(workflowName1);
-      const firstSnapshot = runs[0]!.snapshot as WorkflowRunState;
-      const secondSnapshot = runs[1]!.snapshot as WorkflowRunState;
-      expect(firstSnapshot.context?.steps[stepId2]?.status).toBe('running');
-      expect(secondSnapshot.context?.steps[stepId1]?.status).toBe('completed');
+      const firstSnapshot = runs[0]!.snapshot;
+      const secondSnapshot = runs[1]!.snapshot;
+      checkWorkflowSnapshot(firstSnapshot, stepId2, 'waiting');
+      checkWorkflowSnapshot(secondSnapshot, stepId1, 'success');
     });
 
     it('filters by workflow name', async () => {
       const workflowName1 = 'filter_test_1';
       const workflowName2 = 'filter_test_2';
 
-      const { snapshot: workflow1, runId: runId1, stepId: stepId1 } = createSampleWorkflowSnapshot('completed');
+      const { snapshot: workflow1, runId: runId1, stepId: stepId1 } = createSampleWorkflowSnapshot('success');
       const { snapshot: workflow2, runId: runId2 } = createSampleWorkflowSnapshot('failed');
 
       await store.persistWorkflowSnapshot({
@@ -609,8 +845,8 @@ describe('UpstashStore', () => {
       expect(runs).toHaveLength(1);
       expect(total).toBe(1);
       expect(runs[0]!.workflowName).toBe(workflowName1);
-      const snapshot = runs[0]!.snapshot as WorkflowRunState;
-      expect(snapshot.context?.steps[stepId1]?.status).toBe('completed');
+      const snapshot = runs[0]!.snapshot;
+      checkWorkflowSnapshot(snapshot, stepId1, 'success');
     });
 
     it('filters by date range', async () => {
@@ -621,9 +857,9 @@ describe('UpstashStore', () => {
       const workflowName2 = 'date_test_2';
       const workflowName3 = 'date_test_3';
 
-      const { snapshot: workflow1, runId: runId1 } = createSampleWorkflowSnapshot('completed');
-      const { snapshot: workflow2, runId: runId2, stepId: stepId2 } = createSampleWorkflowSnapshot('running');
-      const { snapshot: workflow3, runId: runId3, stepId: stepId3 } = createSampleWorkflowSnapshot('waiting');
+      const { snapshot: workflow1, runId: runId1 } = createSampleWorkflowSnapshot('success');
+      const { snapshot: workflow2, runId: runId2, stepId: stepId2 } = createSampleWorkflowSnapshot('waiting');
+      const { snapshot: workflow3, runId: runId3, stepId: stepId3 } = createSampleWorkflowSnapshot('skipped');
 
       await store.insert({
         tableName: TABLE_WORKFLOW_SNAPSHOT,
@@ -668,10 +904,10 @@ describe('UpstashStore', () => {
       expect(runs).toHaveLength(2);
       expect(runs[0]!.workflowName).toBe(workflowName3);
       expect(runs[1]!.workflowName).toBe(workflowName2);
-      const firstSnapshot = runs[0]!.snapshot as WorkflowRunState;
-      const secondSnapshot = runs[1]!.snapshot as WorkflowRunState;
-      expect(firstSnapshot.context?.steps[stepId3]?.status).toBe('waiting');
-      expect(secondSnapshot.context?.steps[stepId2]?.status).toBe('running');
+      const firstSnapshot = runs[0]!.snapshot;
+      const secondSnapshot = runs[1]!.snapshot;
+      checkWorkflowSnapshot(firstSnapshot, stepId3, 'skipped');
+      checkWorkflowSnapshot(secondSnapshot, stepId2, 'waiting');
     });
 
     it('handles pagination', async () => {
@@ -679,9 +915,9 @@ describe('UpstashStore', () => {
       const workflowName2 = 'page_test_2';
       const workflowName3 = 'page_test_3';
 
-      const { snapshot: workflow1, runId: runId1, stepId: stepId1 } = createSampleWorkflowSnapshot('completed');
-      const { snapshot: workflow2, runId: runId2, stepId: stepId2 } = createSampleWorkflowSnapshot('running');
-      const { snapshot: workflow3, runId: runId3, stepId: stepId3 } = createSampleWorkflowSnapshot('waiting');
+      const { snapshot: workflow1, runId: runId1, stepId: stepId1 } = createSampleWorkflowSnapshot('success');
+      const { snapshot: workflow2, runId: runId2, stepId: stepId2 } = createSampleWorkflowSnapshot('waiting');
+      const { snapshot: workflow3, runId: runId3, stepId: stepId3 } = createSampleWorkflowSnapshot('skipped');
 
       await store.persistWorkflowSnapshot({
         namespace: testNamespace,
@@ -714,10 +950,10 @@ describe('UpstashStore', () => {
       expect(page1.total).toBe(3); // Total count of all records
       expect(page1.runs[0]!.workflowName).toBe(workflowName3);
       expect(page1.runs[1]!.workflowName).toBe(workflowName2);
-      const firstSnapshot = page1.runs[0]!.snapshot as WorkflowRunState;
-      const secondSnapshot = page1.runs[1]!.snapshot as WorkflowRunState;
-      expect(firstSnapshot.context?.steps[stepId3]?.status).toBe('waiting');
-      expect(secondSnapshot.context?.steps[stepId2]?.status).toBe('running');
+      const firstSnapshot = page1.runs[0]!.snapshot;
+      const secondSnapshot = page1.runs[1]!.snapshot;
+      checkWorkflowSnapshot(firstSnapshot, stepId3, 'skipped');
+      checkWorkflowSnapshot(secondSnapshot, stepId2, 'waiting');
 
       // Get second page
       const page2 = await store.getWorkflowRuns({
@@ -728,8 +964,391 @@ describe('UpstashStore', () => {
       expect(page2.runs).toHaveLength(1);
       expect(page2.total).toBe(3);
       expect(page2.runs[0]!.workflowName).toBe(workflowName1);
-      const snapshot = page2.runs[0]!.snapshot as WorkflowRunState;
-      expect(snapshot.context?.steps[stepId1]?.status).toBe('completed');
+      const snapshot = page2.runs[0]!.snapshot;
+      checkWorkflowSnapshot(snapshot, stepId1, 'success');
+    });
+  });
+  describe('getWorkflowRunById', () => {
+    const testNamespace = 'test-workflows-id';
+    const workflowName = 'workflow-id-test';
+    let runId: string;
+    let stepId: string;
+
+    beforeAll(async () => {
+      // Insert a workflow run for positive test
+      const sample = createSampleWorkflowSnapshot('success');
+      runId = sample.runId;
+      stepId = sample.stepId;
+      await store.insert({
+        tableName: TABLE_WORKFLOW_SNAPSHOT,
+        record: {
+          namespace: testNamespace,
+          workflow_name: workflowName,
+          run_id: runId,
+          resourceId: 'resource-abc',
+          snapshot: sample.snapshot,
+          createdAt: new Date(),
+          updatedAt: new Date(),
+        },
+      });
+    });
+
+    it('should retrieve a workflow run by ID', async () => {
+      const found = await store.getWorkflowRunById({
+        namespace: testNamespace,
+        runId,
+        workflowName,
+      });
+      expect(found).not.toBeNull();
+      expect(found?.runId).toBe(runId);
+      const snapshot = found?.snapshot;
+      checkWorkflowSnapshot(snapshot!, stepId, 'success');
+    });
+
+    it('should return null for non-existent workflow run ID', async () => {
+      const notFound = await store.getWorkflowRunById({
+        namespace: testNamespace,
+        runId: 'non-existent-id',
+        workflowName,
+      });
+      expect(notFound).toBeNull();
+    });
+  });
+  describe('getWorkflowRuns with resourceId', () => {
+    const testNamespace = 'test-workflows-id';
+    const workflowName = 'workflow-id-test';
+    let resourceId: string;
+    let runIds: string[] = [];
+
+    beforeAll(async () => {
+      // Insert multiple workflow runs for the same resourceId
+      resourceId = 'resource-shared';
+      for (const status of ['success', 'waiting']) {
+        const sample = createSampleWorkflowSnapshot(status);
+        runIds.push(sample.runId);
+        await store.insert({
+          tableName: TABLE_WORKFLOW_SNAPSHOT,
+          record: {
+            namespace: testNamespace,
+            workflow_name: workflowName,
+            run_id: sample.runId,
+            resourceId,
+            snapshot: sample.snapshot,
+            createdAt: new Date(),
+            updatedAt: new Date(),
+          },
+        });
+      }
+      // Insert a run with a different resourceId
+      const other = createSampleWorkflowSnapshot('waiting');
+      await store.insert({
+        tableName: TABLE_WORKFLOW_SNAPSHOT,
+        record: {
+          namespace: testNamespace,
+          workflow_name: workflowName,
+          run_id: other.runId,
+          resourceId: 'resource-other',
+          snapshot: other.snapshot,
+          createdAt: new Date(),
+          updatedAt: new Date(),
+        },
+      });
+    });
+
+    it('should retrieve all workflow runs by resourceId', async () => {
+      const { runs } = await store.getWorkflowRuns({
+        namespace: testNamespace,
+        resourceId,
+        workflowName,
+      });
+      expect(Array.isArray(runs)).toBe(true);
+      expect(runs.length).toBeGreaterThanOrEqual(2);
+      for (const run of runs) {
+        expect(run.resourceId).toBe(resourceId);
+      }
+    });
+
+    it('should return an empty array if no workflow runs match resourceId', async () => {
+      const { runs } = await store.getWorkflowRuns({
+        namespace: testNamespace,
+        resourceId: 'non-existent-resource',
+        workflowName,
+      });
+      expect(Array.isArray(runs)).toBe(true);
+      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 page/perPage 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, 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, page: 2, perPage: 5 });
+        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);
+      });
+
+      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('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'));
+
+        for (const trace of traces) {
+          await store.insert({
+            tableName: TABLE_TRACES,
+            record: trace,
+          });
+        }
+
+        const page1 = await store.getTracesPaginated({
+          scope: 'test-scope',
+          page: 0,
+          perPage: 8,
+        });
+        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.getTracesPaginated({
+          scope: 'test-scope',
+          page: 2,
+          perPage: 8,
+        });
+        expect(page3.traces).toHaveLength(2);
+        expect(page3.total).toBe(18);
+        expect(page3.hasMore).toBe(false);
+      });
+
+      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 traces) {
+          await store.insert({
+            tableName: TABLE_TRACES,
+            record: trace,
+          });
+        }
+
+        const result = await store.getTracesPaginated({
+          scope,
+          page: 0,
+          perPage: 10,
+          dateRange: { start: now },
+        });
+
+        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 }));
+
+        for (const thread of threads) {
+          await store.saveThread({ thread });
+        }
+
+        const page1 = await store.getThreadsByResourceIdPaginated({ resourceId, page: 0, perPage: 7 });
+        expect(page1.threads).toHaveLength(7);
+
+        const page3 = await store.getThreadsByResourceIdPaginated({ resourceId, page: 2, perPage: 7 });
+        expect(page3.threads).toHaveLength(3);
+
+        const limited = await store.getThreadsByResourceIdPaginated({ resourceId, page: 1, perPage: 5 });
+        expect(limited.threads).toHaveLength(5);
+      });
     });
   });
 });