@mastra/mongodb 0.14.6 → 0.14.7

package/dist/index.js

@@ -3,11 +3,15 @@ import { MastraVector } from '@mastra/core/vector';
 import { MongoClient } from 'mongodb';
 import { v4 } from 'uuid';
 import { BaseFilterTranslator } from '@mastra/core/vector/filter';
-import { MastraStorage, StoreOperations, TABLE_SCHEMAS, safelyParseJSON, MemoryStorage, TABLE_MESSAGES, resolveMessageLimit, TABLE_THREADS, TABLE_RESOURCES, TracesStorage, TABLE_TRACES, LegacyEvalsStorage, TABLE_EVALS, ScoresStorage, TABLE_SCORERS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT } from '@mastra/core/storage';
+import { MastraStorage, StoreOperations, TABLE_SCHEMAS, safelyParseJSON, MemoryStorage, TABLE_MESSAGES, resolveMessageLimit, TABLE_THREADS, TABLE_RESOURCES, TracesStorage, TABLE_TRACES, LegacyEvalsStorage, TABLE_EVALS, ScoresStorage, TABLE_SCORERS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, ObservabilityStorage, TABLE_AI_SPANS } from '@mastra/core/storage';
 import { MessageList } from '@mastra/core/agent';
 import { saveScorePayloadSchema } from '@mastra/core/scores';
 
 // src/vector/index.ts
+
+// package.json
+var package_default = {
+  version: "0.14.7"};
 var MongoDBFilterTranslator = class extends BaseFilterTranslator {
   getSupportedOperators() {
     return {
@@ -112,7 +116,14 @@ var MongoDBVector = class extends MastraVector {
   };
   constructor({ uri, dbName, options }) {
     super();
-    this.client = new MongoClient(uri, options);
+    const client = new MongoClient(uri, {
+      ...options,
+      driverInfo: {
+        name: "mastra-vector",
+        version: package_default.version
+      }
+    });
+    this.client = client;
     this.db = this.client.db(dbName);
     this.collections = /* @__PURE__ */ new Map();
   }
@@ -632,8 +643,15 @@ var MongoDBConnector = class _MongoDBConnector {
         "MongoDBStore: dbName must be provided and cannot be empty. Passing an empty string may cause fallback to local MongoDB defaults."
       );
     }
+    const client = new MongoClient(config.url, {
+      ...config.options,
+      driverInfo: {
+        name: "mastra-storage",
+        version: package_default.version
+      }
+    });
     return new _MongoDBConnector({
-      client: new MongoClient(config.url, config.options),
+      client,
       dbName: config.dbName,
       handler: void 0
     });
@@ -1426,6 +1444,286 @@ var MemoryStorageMongoDB = class extends MemoryStorage {
     }
   }
 };
+var ObservabilityMongoDB = class extends ObservabilityStorage {
+  operations;
+  constructor({ operations }) {
+    super();
+    this.operations = operations;
+  }
+  get aiTracingStrategy() {
+    return {
+      preferred: "batch-with-updates",
+      supported: ["batch-with-updates", "insert-only"]
+    };
+  }
+  async createAISpan(span) {
+    try {
+      const startedAt = span.startedAt instanceof Date ? span.startedAt.toISOString() : span.startedAt;
+      const endedAt = span.endedAt instanceof Date ? span.endedAt.toISOString() : span.endedAt;
+      const record = {
+        ...span,
+        startedAt,
+        endedAt,
+        createdAt: (/* @__PURE__ */ new Date()).toISOString(),
+        updatedAt: (/* @__PURE__ */ new Date()).toISOString()
+      };
+      return this.operations.insert({ tableName: TABLE_AI_SPANS, record });
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: "MONGODB_STORE_CREATE_AI_SPAN_FAILED",
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER,
+          details: {
+            spanId: span.spanId,
+            traceId: span.traceId,
+            spanType: span.spanType,
+            spanName: span.name
+          }
+        },
+        error
+      );
+    }
+  }
+  async getAITrace(traceId) {
+    try {
+      const collection = await this.operations.getCollection(TABLE_AI_SPANS);
+      const spans = await collection.find({ traceId }).sort({ startedAt: -1 }).toArray();
+      if (!spans || spans.length === 0) {
+        return null;
+      }
+      return {
+        traceId,
+        spans: spans.map((span) => this.transformSpanFromMongo(span))
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: "MONGODB_STORE_GET_AI_TRACE_FAILED",
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER,
+          details: {
+            traceId
+          }
+        },
+        error
+      );
+    }
+  }
+  async updateAISpan({
+    spanId,
+    traceId,
+    updates
+  }) {
+    try {
+      const data = { ...updates };
+      if (data.endedAt instanceof Date) {
+        data.endedAt = data.endedAt.toISOString();
+      }
+      if (data.startedAt instanceof Date) {
+        data.startedAt = data.startedAt.toISOString();
+      }
+      const updateData = {
+        ...data,
+        updatedAt: (/* @__PURE__ */ new Date()).toISOString()
+      };
+      await this.operations.update({
+        tableName: TABLE_AI_SPANS,
+        keys: { spanId, traceId },
+        data: updateData
+      });
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: "MONGODB_STORE_UPDATE_AI_SPAN_FAILED",
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER,
+          details: {
+            spanId,
+            traceId
+          }
+        },
+        error
+      );
+    }
+  }
+  async getAITracesPaginated({
+    filters,
+    pagination
+  }) {
+    const page = pagination?.page ?? 0;
+    const perPage = pagination?.perPage ?? 10;
+    const { entityId, entityType, ...actualFilters } = filters || {};
+    try {
+      const collection = await this.operations.getCollection(TABLE_AI_SPANS);
+      const mongoFilter = {
+        parentSpanId: null,
+        // Only get root spans for traces
+        ...actualFilters
+      };
+      if (pagination?.dateRange) {
+        const dateFilter = {};
+        if (pagination.dateRange.start) {
+          dateFilter.$gte = pagination.dateRange.start instanceof Date ? pagination.dateRange.start.toISOString() : pagination.dateRange.start;
+        }
+        if (pagination.dateRange.end) {
+          dateFilter.$lte = pagination.dateRange.end instanceof Date ? pagination.dateRange.end.toISOString() : pagination.dateRange.end;
+        }
+        if (Object.keys(dateFilter).length > 0) {
+          mongoFilter.startedAt = dateFilter;
+        }
+      }
+      if (entityId && entityType) {
+        let name = "";
+        if (entityType === "workflow") {
+          name = `workflow run: '${entityId}'`;
+        } else if (entityType === "agent") {
+          name = `agent run: '${entityId}'`;
+        } else {
+          const error = new MastraError({
+            id: "MONGODB_STORE_GET_AI_TRACES_PAGINATED_FAILED",
+            domain: ErrorDomain.STORAGE,
+            category: ErrorCategory.USER,
+            details: {
+              entityType
+            },
+            text: `Cannot filter by entity type: ${entityType}`
+          });
+          throw error;
+        }
+        mongoFilter.name = name;
+      }
+      const count = await collection.countDocuments(mongoFilter);
+      if (count === 0) {
+        return {
+          pagination: {
+            total: 0,
+            page,
+            perPage,
+            hasMore: false
+          },
+          spans: []
+        };
+      }
+      const spans = await collection.find(mongoFilter).sort({ startedAt: -1 }).skip(page * perPage).limit(perPage).toArray();
+      return {
+        pagination: {
+          total: count,
+          page,
+          perPage,
+          hasMore: spans.length === perPage
+        },
+        spans: spans.map((span) => this.transformSpanFromMongo(span))
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: "MONGODB_STORE_GET_AI_TRACES_PAGINATED_FAILED",
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER
+        },
+        error
+      );
+    }
+  }
+  async batchCreateAISpans(args) {
+    try {
+      const records = args.records.map((record) => {
+        const startedAt = record.startedAt instanceof Date ? record.startedAt.toISOString() : record.startedAt;
+        const endedAt = record.endedAt instanceof Date ? record.endedAt.toISOString() : record.endedAt;
+        return {
+          ...record,
+          startedAt,
+          endedAt,
+          createdAt: (/* @__PURE__ */ new Date()).toISOString(),
+          updatedAt: (/* @__PURE__ */ new Date()).toISOString()
+        };
+      });
+      return this.operations.batchInsert({
+        tableName: TABLE_AI_SPANS,
+        records
+      });
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: "MONGODB_STORE_BATCH_CREATE_AI_SPANS_FAILED",
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER
+        },
+        error
+      );
+    }
+  }
+  async batchUpdateAISpans(args) {
+    try {
+      return this.operations.batchUpdate({
+        tableName: TABLE_AI_SPANS,
+        updates: args.records.map((record) => {
+          const data = { ...record.updates };
+          if (data.endedAt instanceof Date) {
+            data.endedAt = data.endedAt.toISOString();
+          }
+          if (data.startedAt instanceof Date) {
+            data.startedAt = data.startedAt.toISOString();
+          }
+          const updateData = {
+            ...data,
+            updatedAt: (/* @__PURE__ */ new Date()).toISOString()
+          };
+          return {
+            keys: { spanId: record.spanId, traceId: record.traceId },
+            data: updateData
+          };
+        })
+      });
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: "MONGODB_STORE_BATCH_UPDATE_AI_SPANS_FAILED",
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER
+        },
+        error
+      );
+    }
+  }
+  async batchDeleteAITraces(args) {
+    try {
+      const collection = await this.operations.getCollection(TABLE_AI_SPANS);
+      await collection.deleteMany({
+        traceId: { $in: args.traceIds }
+      });
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: "MONGODB_STORE_BATCH_DELETE_AI_TRACES_FAILED",
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER
+        },
+        error
+      );
+    }
+  }
+  /**
+   * Transform MongoDB document to AISpanRecord format
+   */
+  transformSpanFromMongo(doc) {
+    const { _id, ...span } = doc;
+    if (span.startedAt && typeof span.startedAt === "string") {
+      span.startedAt = span.startedAt;
+    }
+    if (span.endedAt && typeof span.endedAt === "string") {
+      span.endedAt = span.endedAt;
+    }
+    if (span.createdAt && typeof span.createdAt === "string") {
+      span.createdAt = new Date(span.createdAt);
+    }
+    if (span.updatedAt && typeof span.updatedAt === "string") {
+      span.updatedAt = new Date(span.updatedAt);
+    }
+    return span;
+  }
+};
 var StoreOperationsMongoDB = class extends StoreOperations {
   #connector;
   constructor(config) {
@@ -1447,19 +1745,18 @@ var StoreOperationsMongoDB = class extends StoreOperations {
       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);
-      }
+      const mastraError = new MastraError(
+        {
+          id: "STORAGE_MONGODB_STORE_CLEAR_TABLE_FAILED",
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { tableName }
+        },
+        error
+      );
+      this.logger.error(mastraError.message);
+      this.logger?.trackException(mastraError);
+      throw mastraError;
     }
   }
   async dropTable({ tableName }) {
@@ -1483,6 +1780,9 @@ var StoreOperationsMongoDB = class extends StoreOperations {
   }
   processJsonbFields(tableName, record) {
     const schema = TABLE_SCHEMAS[tableName];
+    if (!schema) {
+      return record;
+    }
     return Object.fromEntries(
       Object.entries(schema).map(([key, value]) => {
         if (value.type === "jsonb" && record[key] && typeof record[key] === "string") {
@@ -1498,19 +1798,18 @@ var StoreOperationsMongoDB = class extends StoreOperations {
       const recordToInsert = this.processJsonbFields(tableName, record);
       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);
-      }
+      const mastraError = new MastraError(
+        {
+          id: "STORAGE_MONGODB_STORE_INSERT_FAILED",
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { tableName }
+        },
+        error
+      );
+      this.logger.error(mastraError.message);
+      this.logger?.trackException(mastraError);
+      throw mastraError;
     }
   }
   async batchInsert({ tableName, records }) {
@@ -1550,6 +1849,60 @@ var StoreOperationsMongoDB = class extends StoreOperations {
       );
     }
   }
+  async update({
+    tableName,
+    keys,
+    data
+  }) {
+    try {
+      const collection = await this.getCollection(tableName);
+      const processedData = this.processJsonbFields(tableName, data);
+      const cleanData = Object.fromEntries(Object.entries(processedData).filter(([_, value]) => value !== void 0));
+      await collection.updateOne(keys, { $set: cleanData });
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: "STORAGE_MONGODB_STORE_UPDATE_FAILED",
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { tableName }
+        },
+        error
+      );
+    }
+  }
+  async batchUpdate({
+    tableName,
+    updates
+  }) {
+    if (!updates.length) {
+      return;
+    }
+    try {
+      const collection = await this.getCollection(tableName);
+      const bulkOps = updates.map(({ keys, data }) => {
+        const processedData = this.processJsonbFields(tableName, data);
+        const cleanData = Object.fromEntries(Object.entries(processedData).filter(([_, value]) => value !== void 0));
+        return {
+          updateOne: {
+            filter: keys,
+            update: { $set: cleanData }
+          }
+        };
+      });
+      await collection.bulkWrite(bulkOps);
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: "STORAGE_MONGODB_STORE_BATCH_UPDATE_FAILED",
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { tableName }
+        },
+        error
+      );
+    }
+  }
 };
 function transformScoreRow(row) {
   let scorerValue = null;
@@ -2289,13 +2642,17 @@ var MongoDBStore = class extends MastraStorage {
     const workflows = new WorkflowsStorageMongoDB({
       operations
     });
+    const observability = new ObservabilityMongoDB({
+      operations
+    });
     this.stores = {
       operations,
       memory,
       traces,
       legacyEvals,
       scores,
-      workflows
+      workflows,
+      observability
     };
   }
   async createTable({
@@ -2328,6 +2685,9 @@ var MongoDBStore = class extends MastraStorage {
   async getThreadsByResourceId({ resourceId }) {
     return this.stores.memory.getThreadsByResourceId({ resourceId });
   }
+  async getThreadsByResourceIdPaginated(_args) {
+    return this.stores.memory.getThreadsByResourceIdPaginated(_args);
+  }
   async saveThread({ thread }) {
     return this.stores.memory.saveThread({ thread });
   }
@@ -2348,6 +2708,9 @@ var MongoDBStore = class extends MastraStorage {
   }) {
     return this.stores.memory.getMessages({ threadId, selectBy, format });
   }
+  async getMessagesPaginated(_args) {
+    return this.stores.memory.getMessagesPaginated(_args);
+  }
   async getMessagesById({
     messageIds,
     format
@@ -2357,12 +2720,6 @@ var MongoDBStore = class extends MastraStorage {
   async saveMessages(args) {
     return this.stores.memory.saveMessages(args);
   }
-  async getThreadsByResourceIdPaginated(_args) {
-    return this.stores.memory.getThreadsByResourceIdPaginated(_args);
-  }
-  async getMessagesPaginated(_args) {
-    return this.stores.memory.getMessagesPaginated(_args);
-  }
   async updateMessages(_args) {
     return this.stores.memory.updateMessages(_args);
   }
@@ -2372,6 +2729,9 @@ var MongoDBStore = class extends MastraStorage {
   async getTracesPaginated(args) {
     return this.stores.traces.getTracesPaginated(args);
   }
+  async batchTraceInsert({ records }) {
+    return this.stores.traces.batchTraceInsert({ records });
+  }
   async getWorkflowRuns(args) {
     return this.stores.workflows.getWorkflowRuns(args);
   }
@@ -2489,6 +2849,90 @@ var MongoDBStore = class extends MastraStorage {
       metadata
     });
   }
+  /**
+   * AI Tracing/Observability
+   */
+  async createAISpan(span) {
+    if (!this.stores.observability) {
+      throw new MastraError({
+        id: "MONGODB_STORE_OBSERVABILITY_NOT_INITIALIZED",
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.SYSTEM,
+        text: "Observability storage is not initialized"
+      });
+    }
+    return this.stores.observability.createAISpan(span);
+  }
+  async updateAISpan({
+    spanId,
+    traceId,
+    updates
+  }) {
+    if (!this.stores.observability) {
+      throw new MastraError({
+        id: "MONGODB_STORE_OBSERVABILITY_NOT_INITIALIZED",
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.SYSTEM,
+        text: "Observability storage is not initialized"
+      });
+    }
+    return this.stores.observability.updateAISpan({ spanId, traceId, updates });
+  }
+  async getAITrace(traceId) {
+    if (!this.stores.observability) {
+      throw new MastraError({
+        id: "MONGODB_STORE_OBSERVABILITY_NOT_INITIALIZED",
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.SYSTEM,
+        text: "Observability storage is not initialized"
+      });
+    }
+    return this.stores.observability.getAITrace(traceId);
+  }
+  async getAITracesPaginated(args) {
+    if (!this.stores.observability) {
+      throw new MastraError({
+        id: "MONGODB_STORE_OBSERVABILITY_NOT_INITIALIZED",
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.SYSTEM,
+        text: "Observability storage is not initialized"
+      });
+    }
+    return this.stores.observability.getAITracesPaginated(args);
+  }
+  async batchCreateAISpans(args) {
+    if (!this.stores.observability) {
+      throw new MastraError({
+        id: "MONGODB_STORE_OBSERVABILITY_NOT_INITIALIZED",
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.SYSTEM,
+        text: "Observability storage is not initialized"
+      });
+    }
+    return this.stores.observability.batchCreateAISpans(args);
+  }
+  async batchUpdateAISpans(args) {
+    if (!this.stores.observability) {
+      throw new MastraError({
+        id: "MONGODB_STORE_OBSERVABILITY_NOT_INITIALIZED",
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.SYSTEM,
+        text: "Observability storage is not initialized"
+      });
+    }
+    return this.stores.observability.batchUpdateAISpans(args);
+  }
+  async batchDeleteAITraces(args) {
+    if (!this.stores.observability) {
+      throw new MastraError({
+        id: "MONGODB_STORE_OBSERVABILITY_NOT_INITIALIZED",
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.SYSTEM,
+        text: "Observability storage is not initialized"
+      });
+    }
+    return this.stores.observability.batchDeleteAITraces(args);
+  }
 };
 
 // src/vector/prompt.ts