@mastra/mongodb 0.12.0 → 0.12.1

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

@@ -0,0 +1,196 @@
+import { ErrorDomain, ErrorCategory, MastraError } from '@mastra/core/error';
+import { WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, safelyParseJSON } from '@mastra/core/storage';
+import type { WorkflowRun, WorkflowRuns } from '@mastra/core/storage';
+import type { WorkflowRunState } from '@mastra/core/workflows';
+import type { StoreOperationsMongoDB } from '../operations';
+
+export class WorkflowsStorageMongoDB extends WorkflowsStorage {
+  private operations: StoreOperationsMongoDB;
+
+  constructor({ operations }: { operations: StoreOperationsMongoDB }) {
+    super();
+    this.operations = operations;
+  }
+
+  async persistWorkflowSnapshot({
+    workflowName,
+    runId,
+    snapshot,
+  }: {
+    workflowName: string;
+    runId: string;
+    snapshot: WorkflowRunState;
+  }): Promise<void> {
+    try {
+      const collection = await this.operations.getCollection(TABLE_WORKFLOW_SNAPSHOT);
+      await collection.updateOne(
+        { workflow_name: workflowName, run_id: runId },
+        {
+          $set: {
+            workflow_name: workflowName,
+            run_id: runId,
+            snapshot,
+            createdAt: new Date(),
+            updatedAt: new Date(),
+          },
+        },
+        { upsert: true },
+      );
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_MONGODB_STORE_PERSIST_WORKFLOW_SNAPSHOT_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { workflowName, runId },
+        },
+        error,
+      );
+    }
+  }
+
+  async loadWorkflowSnapshot({
+    workflowName,
+    runId,
+  }: {
+    workflowName: string;
+    runId: string;
+  }): Promise<WorkflowRunState | null> {
+    try {
+      const result = await this.operations.load<any[]>({
+        tableName: TABLE_WORKFLOW_SNAPSHOT,
+        keys: {
+          workflow_name: workflowName,
+          run_id: runId,
+        },
+      });
+
+      if (!result?.length) {
+        return null;
+      }
+
+      return typeof result[0].snapshot === 'string' ? safelyParseJSON(result[0].snapshot) : result[0].snapshot;
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_MONGODB_STORE_LOAD_WORKFLOW_SNAPSHOT_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { workflowName, runId },
+        },
+        error,
+      );
+    }
+  }
+
+  async getWorkflowRuns(args?: {
+    workflowName?: string;
+    fromDate?: Date;
+    toDate?: Date;
+    limit?: number;
+    offset?: number;
+    resourceId?: string;
+  }): Promise<WorkflowRuns> {
+    const options = args || {};
+    try {
+      const query: any = {};
+      if (options.workflowName) {
+        query['workflow_name'] = options.workflowName;
+      }
+      if (options.fromDate) {
+        query['createdAt'] = { $gte: options.fromDate };
+      }
+      if (options.toDate) {
+        if (query['createdAt']) {
+          query['createdAt'].$lte = options.toDate;
+        } else {
+          query['createdAt'] = { $lte: options.toDate };
+        }
+      }
+      if (options.resourceId) {
+        query['resourceId'] = options.resourceId;
+      }
+
+      const collection = await this.operations.getCollection(TABLE_WORKFLOW_SNAPSHOT);
+      const total = await collection.countDocuments(query);
+
+      let cursor = collection.find(query).sort({ createdAt: -1 });
+      if (options.offset) {
+        cursor = cursor.skip(options.offset);
+      }
+      if (options.limit) {
+        cursor = cursor.limit(options.limit);
+      }
+
+      const results = await cursor.toArray();
+
+      const runs = results.map(row => this.parseWorkflowRun(row));
+
+      return {
+        runs,
+        total,
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_MONGODB_STORE_GET_WORKFLOW_RUNS_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { workflowName: options.workflowName || 'unknown' },
+        },
+        error,
+      );
+    }
+  }
+
+  async getWorkflowRunById(args: { runId: string; workflowName?: string }): Promise<WorkflowRun | null> {
+    try {
+      const query: any = {};
+      if (args.runId) {
+        query['run_id'] = args.runId;
+      }
+      if (args.workflowName) {
+        query['workflow_name'] = args.workflowName;
+      }
+
+      const collection = await this.operations.getCollection(TABLE_WORKFLOW_SNAPSHOT);
+      const result = await collection.findOne(query);
+      if (!result) {
+        return null;
+      }
+
+      return this.parseWorkflowRun(result);
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_MONGODB_STORE_GET_WORKFLOW_RUN_BY_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { runId: args.runId },
+        },
+        error,
+      );
+    }
+  }
+
+  private parseWorkflowRun(row: any): WorkflowRun {
+    let parsedSnapshot: WorkflowRunState | string = row.snapshot as string;
+    if (typeof parsedSnapshot === 'string') {
+      try {
+        parsedSnapshot = typeof row.snapshot === 'string' ? safelyParseJSON(row.snapshot as string) : row.snapshot;
+      } 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 as string,
+      runId: row.run_id as string,
+      snapshot: parsedSnapshot,
+      createdAt: new Date(row.createdAt as string),
+      updatedAt: new Date(row.updatedAt as string),
+      resourceId: row.resourceId,
+    };
+  }
+}