@mastra/mongodb 0.12.0 → 0.12.1

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

@@ -0,0 +1,152 @@
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import { safelyParseJSON, StoreOperations, TABLE_SCHEMAS } from '@mastra/core/storage';
+import type { StorageColumn, TABLE_NAMES } from '@mastra/core/storage';
+import type { ConnectorHandler } from '../../connectors/base';
+
+export interface MongoDBOperationsConfig {
+  connector: ConnectorHandler;
+}
+export class StoreOperationsMongoDB extends StoreOperations {
+  readonly #connector: ConnectorHandler;
+
+  constructor(config: MongoDBOperationsConfig) {
+    super();
+    this.#connector = config.connector;
+  }
+
+  async getCollection(collectionName: string) {
+    return this.#connector.getCollection(collectionName);
+  }
+
+  async hasColumn(_table: string, _column: string): Promise<boolean> {
+    // MongoDB is schemaless, so we can assume any column exists
+    // We could check a sample document, but for now return true
+    return true;
+  }
+
+  async createTable(): Promise<void> {
+    // Nothing to do here, MongoDB is schemaless
+  }
+
+  async alterTable(_args: {
+    tableName: TABLE_NAMES;
+    schema: Record<string, StorageColumn>;
+    ifNotExists: string[];
+  }): Promise<void> {
+    // Nothing to do here, MongoDB is schemaless
+  }
+
+  async clearTable({ tableName }: { tableName: TABLE_NAMES }): Promise<void> {
+    try {
+      const collection = await this.getCollection(tableName);
+      await collection.deleteMany({});
+    } catch (error) {
+      if (error instanceof Error) {
+        const matstraError = new MastraError(
+          {
+            id: 'STORAGE_MONGODB_STORE_CLEAR_TABLE_FAILED',
+            domain: ErrorDomain.STORAGE,
+            category: ErrorCategory.THIRD_PARTY,
+            details: { tableName },
+          },
+          error,
+        );
+        this.logger.error(matstraError.message);
+        this.logger?.trackException(matstraError);
+      }
+    }
+  }
+
+  async dropTable({ tableName }: { tableName: TABLE_NAMES }): Promise<void> {
+    try {
+      const collection = await this.getCollection(tableName);
+      await collection.drop();
+    } catch (error) {
+      // Collection might not exist, which is fine
+      if (error instanceof Error && error.message.includes('ns not found')) {
+        return;
+      }
+      throw new MastraError(
+        {
+          id: 'MONGODB_STORE_DROP_TABLE_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { tableName },
+        },
+        error,
+      );
+    }
+  }
+
+  async insert({ tableName, record }: { tableName: TABLE_NAMES; record: Record<string, any> }): Promise<void> {
+    try {
+      const collection = await this.getCollection(tableName);
+
+      const schema = TABLE_SCHEMAS[tableName];
+
+      const recordToInsert = Object.fromEntries(
+        Object.entries(schema).map(([key, value]) => {
+          if (value.type === 'jsonb' && record[key] && typeof record[key] === 'string') {
+            return [key, safelyParseJSON(record[key])];
+          }
+          return [key, record[key]];
+        }),
+      );
+
+      await collection.insertOne(recordToInsert);
+    } catch (error) {
+      if (error instanceof Error) {
+        const matstraError = new MastraError(
+          {
+            id: 'STORAGE_MONGODB_STORE_INSERT_FAILED',
+            domain: ErrorDomain.STORAGE,
+            category: ErrorCategory.THIRD_PARTY,
+            details: { tableName },
+          },
+          error,
+        );
+        this.logger.error(matstraError.message);
+        this.logger?.trackException(matstraError);
+      }
+    }
+  }
+
+  async batchInsert({ tableName, records }: { tableName: TABLE_NAMES; records: Record<string, any>[] }): Promise<void> {
+    if (!records.length) {
+      return;
+    }
+
+    try {
+      const collection = await this.getCollection(tableName);
+      await collection.insertMany(records);
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_MONGODB_STORE_BATCH_INSERT_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { tableName },
+        },
+        error,
+      );
+    }
+  }
+
+  async load<R>({ tableName, keys }: { tableName: TABLE_NAMES; keys: Record<string, string> }): Promise<R | null> {
+    this.logger.info(`Loading ${tableName} with keys ${JSON.stringify(keys)}`);
+    try {
+      const collection = await this.getCollection(tableName);
+      return (await collection.find(keys).toArray()) as R;
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_MONGODB_STORE_LOAD_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { tableName },
+        },
+        error,
+      );
+    }
+  }
+}

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

@@ -0,0 +1,379 @@
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import type { ScoreRowData, ScoringEntityType, ScoringSource } from '@mastra/core/scores';
+import { ScoresStorage, TABLE_SCORERS, safelyParseJSON } from '@mastra/core/storage';
+import type { PaginationInfo, StoragePagination } from '@mastra/core/storage';
+import type { StoreOperationsMongoDB } from '../operations';
+
+function transformScoreRow(row: Record<string, any>): ScoreRowData {
+  let scorerValue: any = null;
+  if (row.scorer) {
+    try {
+      scorerValue = typeof row.scorer === 'string' ? safelyParseJSON(row.scorer) : row.scorer;
+    } catch (e) {
+      console.warn('Failed to parse scorer:', e);
+    }
+  }
+
+  let extractStepResultValue: any = null;
+  if (row.extractStepResult) {
+    try {
+      extractStepResultValue =
+        typeof row.extractStepResult === 'string' ? safelyParseJSON(row.extractStepResult) : row.extractStepResult;
+    } catch (e) {
+      console.warn('Failed to parse extractStepResult:', e);
+    }
+  }
+
+  let analyzeStepResultValue: any = null;
+  if (row.analyzeStepResult) {
+    try {
+      analyzeStepResultValue =
+        typeof row.analyzeStepResult === 'string' ? safelyParseJSON(row.analyzeStepResult) : row.analyzeStepResult;
+    } catch (e) {
+      console.warn('Failed to parse analyzeStepResult:', e);
+    }
+  }
+
+  let inputValue: any = null;
+  if (row.input) {
+    try {
+      inputValue = typeof row.input === 'string' ? safelyParseJSON(row.input) : row.input;
+    } catch (e) {
+      console.warn('Failed to parse input:', e);
+    }
+  }
+
+  let outputValue: any = null;
+  if (row.output) {
+    try {
+      outputValue = typeof row.output === 'string' ? safelyParseJSON(row.output) : row.output;
+    } catch (e) {
+      console.warn('Failed to parse output:', e);
+    }
+  }
+
+  let entityValue: any = null;
+  if (row.entity) {
+    try {
+      entityValue = typeof row.entity === 'string' ? safelyParseJSON(row.entity) : row.entity;
+    } catch (e) {
+      console.warn('Failed to parse entity:', e);
+    }
+  }
+
+  let runtimeContextValue: any = null;
+  if (row.runtimeContext) {
+    try {
+      runtimeContextValue =
+        typeof row.runtimeContext === 'string' ? safelyParseJSON(row.runtimeContext) : row.runtimeContext;
+    } catch (e) {
+      console.warn('Failed to parse runtimeContext:', e);
+    }
+  }
+
+  return {
+    id: row.id as string,
+    entityId: row.entityId as string,
+    entityType: row.entityType as ScoringEntityType,
+    scorerId: row.scorerId as string,
+    traceId: row.traceId as string,
+    runId: row.runId as string,
+    scorer: scorerValue,
+    extractStepResult: extractStepResultValue,
+    analyzeStepResult: analyzeStepResultValue,
+    score: row.score as number,
+    reason: row.reason as string,
+    extractPrompt: row.extractPrompt as string,
+    analyzePrompt: row.analyzePrompt as string,
+    reasonPrompt: row.reasonPrompt as string,
+    input: inputValue,
+    output: outputValue,
+    additionalContext: row.additionalContext,
+    runtimeContext: runtimeContextValue,
+    entity: entityValue,
+    source: row.source as ScoringSource,
+    resourceId: row.resourceId as string,
+    threadId: row.threadId as string,
+    createdAt: new Date(row.createdAt),
+    updatedAt: new Date(row.updatedAt),
+  };
+}
+
+export class ScoresStorageMongoDB extends ScoresStorage {
+  private operations: StoreOperationsMongoDB;
+
+  constructor({ operations }: { operations: StoreOperationsMongoDB }) {
+    super();
+    this.operations = operations;
+  }
+
+  async getScoreById({ id }: { id: string }): Promise<ScoreRowData | null> {
+    try {
+      const collection = await this.operations.getCollection(TABLE_SCORERS);
+      const document = await collection.findOne({ id });
+
+      if (!document) {
+        return null;
+      }
+
+      return transformScoreRow(document);
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_MONGODB_STORE_GET_SCORE_BY_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { id },
+        },
+        error,
+      );
+    }
+  }
+
+  async saveScore(score: Omit<ScoreRowData, 'id' | 'createdAt' | 'updatedAt'>): Promise<{ score: ScoreRowData }> {
+    try {
+      const now = new Date();
+      const scoreId = `score-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
+
+      const scoreData = {
+        id: scoreId,
+        entityId: score.entityId,
+        entityType: score.entityType,
+        scorerId: score.scorerId,
+        traceId: score.traceId || '',
+        runId: score.runId,
+        scorer: typeof score.scorer === 'string' ? score.scorer : JSON.stringify(score.scorer),
+        extractStepResult:
+          typeof score.extractStepResult === 'string'
+            ? score.extractStepResult
+            : JSON.stringify(score.extractStepResult),
+        analyzeStepResult:
+          typeof score.analyzeStepResult === 'string'
+            ? score.analyzeStepResult
+            : JSON.stringify(score.analyzeStepResult),
+        score: score.score,
+        reason: score.reason,
+        extractPrompt: score.extractPrompt,
+        analyzePrompt: score.analyzePrompt,
+        reasonPrompt: score.reasonPrompt,
+        input: typeof score.input === 'string' ? score.input : JSON.stringify(score.input),
+        output: typeof score.output === 'string' ? score.output : JSON.stringify(score.output),
+        additionalContext: score.additionalContext,
+        runtimeContext:
+          typeof score.runtimeContext === 'string' ? score.runtimeContext : JSON.stringify(score.runtimeContext),
+        entity: typeof score.entity === 'string' ? score.entity : JSON.stringify(score.entity),
+        source: score.source,
+        resourceId: score.resourceId || '',
+        threadId: score.threadId || '',
+        createdAt: now,
+        updatedAt: now,
+      };
+
+      const collection = await this.operations.getCollection(TABLE_SCORERS);
+      await collection.insertOne(scoreData);
+
+      const savedScore: ScoreRowData = {
+        ...score,
+        id: scoreId,
+        createdAt: now,
+        updatedAt: now,
+      };
+
+      return { score: savedScore };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_MONGODB_STORE_SAVE_SCORE_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { scorerId: score.scorerId, runId: score.runId },
+        },
+        error,
+      );
+    }
+  }
+
+  async getScoresByScorerId({
+    scorerId,
+    pagination,
+    entityId,
+    entityType,
+  }: {
+    scorerId: string;
+    pagination: StoragePagination;
+    entityId?: string;
+    entityType?: string;
+  }): Promise<{ pagination: PaginationInfo; scores: ScoreRowData[] }> {
+    try {
+      const query: any = { scorerId };
+
+      if (entityId) {
+        query.entityId = entityId;
+      }
+
+      if (entityType) {
+        query.entityType = entityType;
+      }
+
+      const collection = await this.operations.getCollection(TABLE_SCORERS);
+      const total = await collection.countDocuments(query);
+      const currentOffset = pagination.page * pagination.perPage;
+
+      if (total === 0) {
+        return {
+          scores: [],
+          pagination: {
+            total: 0,
+            page: pagination.page,
+            perPage: pagination.perPage,
+            hasMore: false,
+          },
+        };
+      }
+
+      const documents = await collection
+        .find(query)
+        .sort({ createdAt: 'desc' })
+        .skip(currentOffset)
+        .limit(pagination.perPage)
+        .toArray();
+
+      const scores = documents.map(row => transformScoreRow(row));
+      const hasMore = currentOffset + scores.length < total;
+
+      return {
+        scores,
+        pagination: {
+          total,
+          page: pagination.page,
+          perPage: pagination.perPage,
+          hasMore,
+        },
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_MONGODB_STORE_GET_SCORES_BY_SCORER_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { scorerId, page: pagination.page, perPage: pagination.perPage },
+        },
+        error,
+      );
+    }
+  }
+
+  async getScoresByRunId({
+    runId,
+    pagination,
+  }: {
+    runId: string;
+    pagination: StoragePagination;
+  }): Promise<{ pagination: PaginationInfo; scores: ScoreRowData[] }> {
+    try {
+      const collection = await this.operations.getCollection(TABLE_SCORERS);
+      const total = await collection.countDocuments({ runId });
+      const currentOffset = pagination.page * pagination.perPage;
+
+      if (total === 0) {
+        return {
+          scores: [],
+          pagination: {
+            total: 0,
+            page: pagination.page,
+            perPage: pagination.perPage,
+            hasMore: false,
+          },
+        };
+      }
+
+      const documents = await collection
+        .find({ runId })
+        .sort({ createdAt: 'desc' })
+        .skip(currentOffset)
+        .limit(pagination.perPage)
+        .toArray();
+
+      const scores = documents.map(row => transformScoreRow(row));
+      const hasMore = currentOffset + scores.length < total;
+
+      return {
+        scores,
+        pagination: {
+          total,
+          page: pagination.page,
+          perPage: pagination.perPage,
+          hasMore,
+        },
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_MONGODB_STORE_GET_SCORES_BY_RUN_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { runId, page: pagination.page, perPage: pagination.perPage },
+        },
+        error,
+      );
+    }
+  }
+
+  async getScoresByEntityId({
+    entityId,
+    entityType,
+    pagination,
+  }: {
+    pagination: StoragePagination;
+    entityId: string;
+    entityType: string;
+  }): Promise<{ pagination: PaginationInfo; scores: ScoreRowData[] }> {
+    try {
+      const collection = await this.operations.getCollection(TABLE_SCORERS);
+      const total = await collection.countDocuments({ entityId, entityType });
+      const currentOffset = pagination.page * pagination.perPage;
+
+      if (total === 0) {
+        return {
+          scores: [],
+          pagination: {
+            total: 0,
+            page: pagination.page,
+            perPage: pagination.perPage,
+            hasMore: false,
+          },
+        };
+      }
+
+      const documents = await collection
+        .find({ entityId, entityType })
+        .sort({ createdAt: 'desc' })
+        .skip(currentOffset)
+        .limit(pagination.perPage)
+        .toArray();
+
+      const scores = documents.map(row => transformScoreRow(row));
+      const hasMore = currentOffset + scores.length < total;
+
+      return {
+        scores,
+        pagination: {
+          total,
+          page: pagination.page,
+          perPage: pagination.perPage,
+          hasMore,
+        },
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_MONGODB_STORE_GET_SCORES_BY_ENTITY_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { entityId, entityType, page: pagination.page, perPage: pagination.perPage },
+        },
+        error,
+      );
+    }
+  }
+}

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

@@ -0,0 +1,142 @@
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import type { PaginationInfo, StorageGetTracesArg, StorageGetTracesPaginatedArg } from '@mastra/core/storage';
+import { TABLE_TRACES, TracesStorage, safelyParseJSON } from '@mastra/core/storage';
+import type { Trace } from '@mastra/core/telemetry';
+import type { StoreOperationsMongoDB } from '../operations';
+
+export class TracesStorageMongoDB extends TracesStorage {
+  private operations: StoreOperationsMongoDB;
+
+  constructor({ operations }: { operations: StoreOperationsMongoDB }) {
+    super();
+    this.operations = operations;
+  }
+
+  async getTraces(args: StorageGetTracesArg): Promise<Trace[]> {
+    if (args.fromDate || args.toDate) {
+      (args as any).dateRange = {
+        start: args.fromDate,
+        end: args.toDate,
+      };
+    }
+    try {
+      const result = await this.getTracesPaginated(args);
+      return result.traces;
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_MONGODB_STORE_GET_TRACES_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+        },
+        error,
+      );
+    }
+  }
+
+  async getTracesPaginated(args: StorageGetTracesPaginatedArg): Promise<PaginationInfo & { traces: Trace[] }> {
+    const { name, scope, page = 0, perPage = 100, attributes, filters, dateRange } = args;
+    const fromDate = dateRange?.start;
+    const toDate = dateRange?.end;
+    const currentOffset = page * perPage;
+
+    const query: any = {};
+    if (name) {
+      query['name'] = new RegExp(name);
+    }
+
+    if (scope) {
+      query['scope'] = scope;
+    }
+
+    if (attributes) {
+      query['$and'] = Object.entries(attributes).map(([key, value]) => ({
+        [`attributes.${key}`]: value,
+      }));
+    }
+
+    if (filters) {
+      Object.entries(filters).forEach(([key, value]) => {
+        query[key] = value;
+      });
+    }
+
+    if (fromDate || toDate) {
+      query['createdAt'] = {};
+      if (fromDate) {
+        query['createdAt']['$gte'] = fromDate;
+      }
+      if (toDate) {
+        query['createdAt']['$lte'] = toDate;
+      }
+    }
+
+    try {
+      const collection = await this.operations.getCollection(TABLE_TRACES);
+
+      // Get total count
+      const total = await collection.countDocuments(query);
+
+      if (total === 0) {
+        return {
+          traces: [],
+          total: 0,
+          page,
+          perPage,
+          hasMore: false,
+        };
+      }
+
+      // Get results with pagination
+      const result = await collection
+        .find(query, {
+          sort: { startTime: -1 },
+        })
+        .limit(perPage)
+        .skip(currentOffset)
+        .toArray();
+
+      const traces = result.map(row => ({
+        id: row.id,
+        parentSpanId: row.parentSpanId,
+        traceId: row.traceId,
+        name: row.name,
+        scope: row.scope,
+        kind: row.kind,
+        status: safelyParseJSON(row.status as string),
+        events: safelyParseJSON(row.events as string),
+        links: safelyParseJSON(row.links as string),
+        attributes: safelyParseJSON(row.attributes as string),
+        startTime: row.startTime,
+        endTime: row.endTime,
+        other: safelyParseJSON(row.other as string),
+        createdAt: row.createdAt,
+      })) as Trace[];
+
+      return {
+        traces,
+        total,
+        page,
+        perPage,
+        hasMore: currentOffset + traces.length < total,
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_MONGODB_STORE_GET_TRACES_PAGINATED_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+        },
+        error,
+      );
+    }
+  }
+
+  async batchTraceInsert({ records }: { records: Record<string, any>[] }): Promise<void> {
+    this.logger.debug('Batch inserting traces', { count: records.length });
+    await this.operations.batchInsert({
+      tableName: TABLE_TRACES,
+      records,
+    });
+  }
+}

package/src/storage/domains/utils.ts

@@ -0,0 +1,43 @@
+import type { IMastraLogger } from '@mastra/core/logger';
+
+export function formatDateForMongoDB(date: Date | string): Date {
+  return typeof date === 'string' ? new Date(date) : date;
+}
+
+export function createExecuteOperationWithRetry({
+  logger,
+  maxRetries = 3,
+  initialBackoffMs = 100,
+}: {
+  logger: IMastraLogger;
+  maxRetries?: number;
+  initialBackoffMs?: number;
+}) {
+  return async function executeOperationWithRetry<T>(
+    operationFn: () => Promise<T>,
+    operationDescription: string,
+  ): Promise<T> {
+    let retries = 0;
+
+    while (true) {
+      try {
+        return await operationFn();
+      } catch (error: any) {
+        if (
+          ((error.message && error.message.includes('connection')) || error.code === 'ECONNRESET') &&
+          retries < maxRetries
+        ) {
+          retries++;
+          const backoffTime = initialBackoffMs * Math.pow(2, retries - 1);
+          logger.warn(
+            `MongoDBStore: Encountered connection error during ${operationDescription}. Retrying (${retries}/${maxRetries}) in ${backoffTime}ms...`,
+          );
+          await new Promise(resolve => setTimeout(resolve, backoffTime));
+        } else {
+          logger.error(`MongoDBStore: Error during ${operationDescription} after ${retries} retries: ${error}`);
+          throw error;
+        }
+      }
+    }
+  };
+}