@mastra/upstash 0.12.1 → 0.12.2

package/src/storage/domains/workflows/index.ts

@@ -0,0 +1,243 @@
+import type { WorkflowRun, WorkflowRuns, WorkflowRunState } from '@mastra/core';
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import { TABLE_WORKFLOW_SNAPSHOT, WorkflowsStorage } from '@mastra/core/storage';
+import type { Redis } from '@upstash/redis';
+import type { StoreOperationsUpstash } from '../operations';
+import { ensureDate, getKey } from '../utils';
+
+function parseWorkflowRun(row: any): WorkflowRun {
+  let parsedSnapshot: WorkflowRunState | string = row.snapshot as string;
+  if (typeof parsedSnapshot === 'string') {
+    try {
+      parsedSnapshot = JSON.parse(row.snapshot as string) as WorkflowRunState;
+    } catch (e) {
+      // If parsing fails, return the raw snapshot string
+      console.warn(`Failed to parse snapshot for workflow ${row.workflow_name}: ${e}`);
+    }
+  }
+
+  return {
+    workflowName: row.workflow_name,
+    runId: row.run_id,
+    snapshot: parsedSnapshot,
+    createdAt: ensureDate(row.createdAt)!,
+    updatedAt: ensureDate(row.updatedAt)!,
+    resourceId: row.resourceId,
+  };
+}
+export class WorkflowsUpstash extends WorkflowsStorage {
+  private client: Redis;
+  private operations: StoreOperationsUpstash;
+
+  constructor({ client, operations }: { client: Redis; operations: StoreOperationsUpstash }) {
+    super();
+    this.client = client;
+    this.operations = operations;
+  }
+
+  async persistWorkflowSnapshot(params: {
+    namespace: string;
+    workflowName: string;
+    runId: string;
+    snapshot: WorkflowRunState;
+  }): Promise<void> {
+    const { namespace = 'workflows', workflowName, runId, snapshot } = params;
+    try {
+      await this.operations.insert({
+        tableName: TABLE_WORKFLOW_SNAPSHOT,
+        record: {
+          namespace,
+          workflow_name: workflowName,
+          run_id: runId,
+          snapshot,
+          createdAt: new Date(),
+          updatedAt: new Date(),
+        },
+      });
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_UPSTASH_STORAGE_PERSIST_WORKFLOW_SNAPSHOT_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            namespace,
+            workflowName,
+            runId,
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  async loadWorkflowSnapshot(params: {
+    namespace: string;
+    workflowName: string;
+    runId: string;
+  }): Promise<WorkflowRunState | null> {
+    const { namespace = 'workflows', workflowName, runId } = params;
+    const key = getKey(TABLE_WORKFLOW_SNAPSHOT, {
+      namespace,
+      workflow_name: workflowName,
+      run_id: runId,
+    });
+    try {
+      const data = await this.client.get<{
+        namespace: string;
+        workflow_name: string;
+        run_id: string;
+        snapshot: WorkflowRunState;
+      }>(key);
+      if (!data) return null;
+      return data.snapshot;
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_UPSTASH_STORAGE_LOAD_WORKFLOW_SNAPSHOT_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            namespace,
+            workflowName,
+            runId,
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  async getWorkflowRunById({
+    runId,
+    workflowName,
+  }: {
+    runId: string;
+    workflowName?: string;
+  }): Promise<WorkflowRun | null> {
+    try {
+      const key =
+        getKey(TABLE_WORKFLOW_SNAPSHOT, { namespace: 'workflows', workflow_name: workflowName, run_id: runId }) + '*';
+      const keys = await this.operations.scanKeys(key);
+      const workflows = await Promise.all(
+        keys.map(async key => {
+          const data = await this.client.get<{
+            workflow_name: string;
+            run_id: string;
+            snapshot: WorkflowRunState | string;
+            createdAt: string | Date;
+            updatedAt: string | Date;
+            resourceId: string;
+          }>(key);
+          return data;
+        }),
+      );
+      const data = workflows.find(w => w?.run_id === runId && w?.workflow_name === workflowName) as WorkflowRun | null;
+      if (!data) return null;
+      return parseWorkflowRun(data);
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_UPSTASH_STORAGE_GET_WORKFLOW_RUN_BY_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            namespace: 'workflows',
+            runId,
+            workflowName: workflowName || '',
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  async getWorkflowRuns({
+    workflowName,
+    fromDate,
+    toDate,
+    limit,
+    offset,
+    resourceId,
+  }: {
+    workflowName?: string;
+    fromDate?: Date;
+    toDate?: Date;
+    limit?: number;
+    offset?: number;
+    resourceId?: string;
+  }): Promise<WorkflowRuns> {
+    try {
+      // Get all workflow keys
+      let pattern = getKey(TABLE_WORKFLOW_SNAPSHOT, { namespace: 'workflows' }) + ':*';
+      if (workflowName && resourceId) {
+        pattern = getKey(TABLE_WORKFLOW_SNAPSHOT, {
+          namespace: 'workflows',
+          workflow_name: workflowName,
+          run_id: '*',
+          resourceId,
+        });
+      } else if (workflowName) {
+        pattern = getKey(TABLE_WORKFLOW_SNAPSHOT, { namespace: 'workflows', workflow_name: workflowName }) + ':*';
+      } else if (resourceId) {
+        pattern = getKey(TABLE_WORKFLOW_SNAPSHOT, {
+          namespace: 'workflows',
+          workflow_name: '*',
+          run_id: '*',
+          resourceId,
+        });
+      }
+      const keys = await this.operations.scanKeys(pattern);
+
+      // Check if we have any keys before using pipeline
+      if (keys.length === 0) {
+        return { runs: [], total: 0 };
+      }
+
+      // Use pipeline for batch fetching to improve performance
+      const pipeline = this.client.pipeline();
+      keys.forEach(key => pipeline.get(key));
+      const results = await pipeline.exec();
+
+      // Filter and transform results - handle undefined results
+      let runs = results
+        .map((result: any) => result as Record<string, any> | null)
+        .filter(
+          (record): record is Record<string, any> =>
+            record !== null && record !== undefined && typeof record === 'object' && 'workflow_name' in record,
+        )
+        // Only filter by workflowName if it was specifically requested
+        .filter(record => !workflowName || record.workflow_name === workflowName)
+        .map(w => parseWorkflowRun(w!))
+        .filter(w => {
+          if (fromDate && w.createdAt < fromDate) return false;
+          if (toDate && w.createdAt > toDate) return false;
+          return true;
+        })
+        .sort((a, b) => b.createdAt.getTime() - a.createdAt.getTime());
+
+      const total = runs.length;
+
+      // Apply pagination if requested
+      if (limit !== undefined && offset !== undefined) {
+        runs = runs.slice(offset, offset + limit);
+      }
+
+      return { runs, total };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_UPSTASH_STORAGE_GET_WORKFLOW_RUNS_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            namespace: 'workflows',
+            workflowName: workflowName || '',
+            resourceId: resourceId || '',
+          },
+        },
+        error,
+      );
+    }
+  }
+}

package/src/storage/index.test.ts

@@ -0,0 +1,13 @@
+import { createTestSuite } from '@internal/storage-test-utils';
+import { vi } 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 });
+
+createTestSuite(
+  new UpstashStore({
+    url: 'http://localhost:8079',
+    token: 'test_token',
+  }),
+);