@mastra/mongodb 0.10.0 → 0.10.1-alpha.0

package/dist/index.js

@@ -2,6 +2,7 @@ import { MastraVector } from '@mastra/core/vector';
 import { MongoClient } from 'mongodb';
 import { v4 } from 'uuid';
 import { BaseFilterTranslator } from '@mastra/core/vector/filter';
+import { MastraStorage, TABLE_THREADS, TABLE_MESSAGES, TABLE_TRACES, TABLE_WORKFLOW_SNAPSHOT, TABLE_EVALS } from '@mastra/core/storage';
 
 // src/vector/index.ts
 var MongoDBFilterTranslator = class extends BaseFilterTranslator {
@@ -400,6 +401,524 @@ var MongoDBVector = class extends MastraVector {
     return translator.translate(filter);
   }
 };
+function safelyParseJSON(jsonString) {
+  try {
+    return JSON.parse(jsonString);
+  } catch {
+    return {};
+  }
+}
+var MongoDBStore = class extends MastraStorage {
+  #isConnected = false;
+  #client;
+  #db;
+  #dbName;
+  constructor(config) {
+    super({ name: "MongoDBStore" });
+    this.#isConnected = false;
+    if (!config.url?.trim().length) {
+      throw new Error(
+        "MongoDBStore: url must be provided and cannot be empty. Passing an empty string may cause fallback to local MongoDB defaults."
+      );
+    }
+    if (!config.dbName?.trim().length) {
+      throw new Error(
+        "MongoDBStore: dbName must be provided and cannot be empty. Passing an empty string may cause fallback to local MongoDB defaults."
+      );
+    }
+    this.#dbName = config.dbName;
+    this.#client = new MongoClient(config.url);
+  }
+  async getConnection() {
+    if (this.#isConnected) {
+      return this.#db;
+    }
+    await this.#client.connect();
+    this.#db = this.#client.db(this.#dbName);
+    this.#isConnected = true;
+    return this.#db;
+  }
+  async getCollection(collectionName) {
+    const db = await this.getConnection();
+    return db.collection(collectionName);
+  }
+  async createTable() {
+  }
+  async clearTable({ tableName }) {
+    try {
+      const collection = await this.getCollection(tableName);
+      await collection.deleteMany({});
+    } catch (error) {
+      if (error instanceof Error) {
+        this.logger.error(error.message);
+      }
+    }
+  }
+  async insert({ tableName, record }) {
+    try {
+      const collection = await this.getCollection(tableName);
+      await collection.insertOne(record);
+    } catch (error) {
+      this.logger.error(`Error upserting into table ${tableName}: ${error}`);
+      throw error;
+    }
+  }
+  async batchInsert({ tableName, records }) {
+    if (!records.length) {
+      return;
+    }
+    try {
+      const collection = await this.getCollection(tableName);
+      await collection.insertMany(records);
+    } catch (error) {
+      this.logger.error(`Error upserting into table ${tableName}: ${error}`);
+      throw error;
+    }
+  }
+  async load({ tableName, keys }) {
+    this.logger.info(`Loading ${tableName} with keys ${JSON.stringify(keys)}`);
+    try {
+      const collection = await this.getCollection(tableName);
+      return await collection.find(keys).toArray();
+    } catch (error) {
+      this.logger.error(`Error loading ${tableName} with keys ${JSON.stringify(keys)}: ${error}`);
+      throw error;
+    }
+  }
+  async getThreadById({ threadId }) {
+    try {
+      const collection = await this.getCollection(TABLE_THREADS);
+      const result = await collection.findOne({ id: threadId });
+      if (!result) {
+        return null;
+      }
+      return {
+        ...result,
+        metadata: typeof result.metadata === "string" ? JSON.parse(result.metadata) : result.metadata
+      };
+    } catch (error) {
+      this.logger.error(`Error loading thread with ID ${threadId}: ${error}`);
+      throw error;
+    }
+  }
+  async getThreadsByResourceId({ resourceId }) {
+    try {
+      const collection = await this.getCollection(TABLE_THREADS);
+      const results = await collection.find({ resourceId }).toArray();
+      if (!results.length) {
+        return [];
+      }
+      return results.map((result) => ({
+        ...result,
+        metadata: typeof result.metadata === "string" ? JSON.parse(result.metadata) : result.metadata
+      }));
+    } catch (error) {
+      this.logger.error(`Error loading threads by resourceId ${resourceId}: ${error}`);
+      throw error;
+    }
+  }
+  async saveThread({ thread }) {
+    try {
+      const collection = await this.getCollection(TABLE_THREADS);
+      await collection.updateOne(
+        { id: thread.id },
+        {
+          $set: {
+            ...thread,
+            metadata: JSON.stringify(thread.metadata)
+          }
+        },
+        { upsert: true }
+      );
+      return thread;
+    } catch (error) {
+      this.logger.error(`Error saving thread ${thread.id}: ${error}`);
+      throw error;
+    }
+  }
+  async updateThread({
+    id,
+    title,
+    metadata
+  }) {
+    const thread = await this.getThreadById({ threadId: id });
+    if (!thread) {
+      throw new Error(`Thread ${id} not found`);
+    }
+    const updatedThread = {
+      ...thread,
+      title,
+      metadata: {
+        ...thread.metadata,
+        ...metadata
+      }
+    };
+    try {
+      const collection = await this.getCollection(TABLE_THREADS);
+      await collection.updateOne(
+        { id },
+        {
+          $set: {
+            title,
+            metadata: JSON.stringify(updatedThread.metadata)
+          }
+        }
+      );
+    } catch (error) {
+      this.logger.error(`Error updating thread ${id}:) ${error}`);
+      throw error;
+    }
+    return updatedThread;
+  }
+  async deleteThread({ threadId }) {
+    try {
+      const collectionMessages = await this.getCollection(TABLE_MESSAGES);
+      await collectionMessages.deleteMany({ thread_id: threadId });
+      const collectionThreads = await this.getCollection(TABLE_THREADS);
+      await collectionThreads.deleteOne({ id: threadId });
+    } catch (error) {
+      this.logger.error(`Error deleting thread ${threadId}: ${error}`);
+      throw error;
+    }
+  }
+  async getMessages({ threadId, selectBy }) {
+    try {
+      const limit = typeof selectBy?.last === "number" ? selectBy.last : 40;
+      const include = selectBy?.include || [];
+      let messages = [];
+      let allMessages = [];
+      const collection = await this.getCollection(TABLE_MESSAGES);
+      allMessages = (await collection.find({ thread_id: threadId }).sort({ createdAt: -1 }).toArray()).map(
+        (row) => this.parseRow(row)
+      );
+      if (include.length) {
+        const idToIndex = /* @__PURE__ */ new Map();
+        allMessages.forEach((msg, idx) => {
+          idToIndex.set(msg.id, idx);
+        });
+        const selectedIndexes = /* @__PURE__ */ new Set();
+        for (const inc of include) {
+          const idx = idToIndex.get(inc.id);
+          if (idx === void 0) continue;
+          for (let i = 1; i <= (inc.withPreviousMessages || 0); i++) {
+            if (idx + i < allMessages.length) selectedIndexes.add(idx + i);
+          }
+          selectedIndexes.add(idx);
+          for (let i = 1; i <= (inc.withNextMessages || 0); i++) {
+            if (idx - i >= 0) selectedIndexes.add(idx - i);
+          }
+        }
+        messages.push(
+          ...Array.from(selectedIndexes).map((i) => allMessages[i]).filter((m) => !!m)
+        );
+      }
+      const excludeIds = new Set(messages.map((m) => m.id));
+      for (const msg of allMessages) {
+        if (messages.length >= limit) break;
+        if (!excludeIds.has(msg.id)) {
+          messages.push(msg);
+        }
+      }
+      messages.sort((a, b) => a.createdAt.getTime() - b.createdAt.getTime());
+      return messages.slice(0, limit);
+    } catch (error) {
+      this.logger.error("Error getting messages:", error);
+      throw error;
+    }
+  }
+  async saveMessages({ messages }) {
+    if (!messages.length) {
+      return messages;
+    }
+    const threadId = messages[0]?.threadId;
+    if (!threadId) {
+      this.logger.error("Thread ID is required to save messages");
+      throw new Error("Thread ID is required");
+    }
+    try {
+      const messagesToInsert = messages.map((message) => {
+        const time = message.createdAt || /* @__PURE__ */ new Date();
+        return {
+          id: message.id,
+          thread_id: threadId,
+          content: typeof message.content === "string" ? message.content : JSON.stringify(message.content),
+          role: message.role,
+          type: message.type,
+          resourceId: message.resourceId,
+          createdAt: time instanceof Date ? time.toISOString() : time
+        };
+      });
+      const collection = await this.getCollection(TABLE_MESSAGES);
+      await collection.insertMany(messagesToInsert);
+      return messages;
+    } catch (error) {
+      this.logger.error("Failed to save messages in database: " + error?.message);
+      throw error;
+    }
+  }
+  async getTraces({
+    name,
+    scope,
+    page,
+    perPage,
+    attributes,
+    filters
+  } = {
+    page: 0,
+    perPage: 100
+  }) {
+    const limit = perPage;
+    const offset = page * perPage;
+    const query = {};
+    if (name) {
+      query["name"] = `%${name}%`;
+    }
+    if (scope) {
+      query["scope"] = scope;
+    }
+    if (attributes) {
+      Object.keys(attributes).forEach((key) => {
+        query[`attributes.${key}`] = attributes[key];
+      });
+    }
+    if (filters) {
+      Object.entries(filters).forEach(([key, value]) => {
+        query[key] = value;
+      });
+    }
+    const collection = await this.getCollection(TABLE_TRACES);
+    const result = await collection.find(query, {
+      sort: { startTime: -1 }
+    }).limit(limit).skip(offset).toArray();
+    return 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),
+      events: safelyParseJSON(row.events),
+      links: safelyParseJSON(row.links),
+      attributes: safelyParseJSON(row.attributes),
+      startTime: row.startTime,
+      endTime: row.endTime,
+      other: safelyParseJSON(row.other),
+      createdAt: row.createdAt
+    }));
+  }
+  async getWorkflowRuns({
+    workflowName,
+    fromDate,
+    toDate,
+    limit,
+    offset
+  } = {}) {
+    const query = {};
+    if (workflowName) {
+      query["workflow_name"] = workflowName;
+    }
+    if (fromDate || toDate) {
+      query["createdAt"] = {};
+      if (fromDate) {
+        query["createdAt"]["$gte"] = fromDate;
+      }
+      if (toDate) {
+        query["createdAt"]["$lte"] = toDate;
+      }
+    }
+    const collection = await this.getCollection(TABLE_WORKFLOW_SNAPSHOT);
+    let total = 0;
+    if (limit !== void 0 && offset !== void 0) {
+      total = await collection.countDocuments(query);
+    }
+    const request = collection.find(query).sort({ createdAt: "desc" });
+    if (limit) {
+      request.limit(limit);
+    }
+    if (offset) {
+      request.skip(offset);
+    }
+    const result = await request.toArray();
+    const runs = result.map((row) => {
+      let parsedSnapshot = row.snapshot;
+      if (typeof parsedSnapshot === "string") {
+        try {
+          parsedSnapshot = JSON.parse(row.snapshot);
+        } catch (e) {
+          console.warn(`Failed to parse snapshot for workflow ${row.workflow_name}: ${e}`);
+        }
+      }
+      return {
+        workflowName: row.workflow_name,
+        runId: row.run_id,
+        snapshot: parsedSnapshot,
+        createdAt: new Date(row.createdAt),
+        updatedAt: new Date(row.updatedAt)
+      };
+    });
+    return { runs, total: total || runs.length };
+  }
+  async getEvalsByAgentName(agentName, type) {
+    try {
+      const query = {
+        agent_name: agentName
+      };
+      if (type === "test") {
+        query["test_info"] = { $ne: null };
+      }
+      if (type === "live") {
+        query["test_info"] = null;
+      }
+      const collection = await this.getCollection(TABLE_EVALS);
+      const documents = await collection.find(query).sort({ created_at: "desc" }).toArray();
+      const result = documents.map((row) => this.transformEvalRow(row));
+      return result.filter((row) => {
+        if (type === "live") {
+          return !Boolean(row.testInfo?.testPath);
+        }
+        if (type === "test") {
+          return row.testInfo?.testPath !== null;
+        }
+        return true;
+      });
+    } catch (error) {
+      if (error instanceof Error && error.message.includes("no such table")) {
+        return [];
+      }
+      this.logger.error("Failed to get evals for the specified agent: " + error?.message);
+      throw error;
+    }
+  }
+  async persistWorkflowSnapshot({
+    workflowName,
+    runId,
+    snapshot
+  }) {
+    try {
+      const now = (/* @__PURE__ */ new Date()).toISOString();
+      const collection = await this.getCollection(TABLE_WORKFLOW_SNAPSHOT);
+      await collection.updateOne(
+        { workflow_name: workflowName, run_id: runId },
+        {
+          $set: {
+            snapshot: JSON.stringify(snapshot),
+            updatedAt: now
+          },
+          $setOnInsert: {
+            createdAt: now
+          }
+        },
+        { upsert: true }
+      );
+    } catch (error) {
+      this.logger.error(`Error persisting workflow snapshot: ${error}`);
+      throw error;
+    }
+  }
+  async loadWorkflowSnapshot({
+    workflowName,
+    runId
+  }) {
+    try {
+      const result = await this.load({
+        tableName: TABLE_WORKFLOW_SNAPSHOT,
+        keys: {
+          workflow_name: workflowName,
+          run_id: runId
+        }
+      });
+      if (!result?.length) {
+        return null;
+      }
+      return JSON.parse(result[0].snapshot);
+    } catch (error) {
+      console.error("Error loading workflow snapshot:", error);
+      throw error;
+    }
+  }
+  async getWorkflowRunById({
+    runId,
+    workflowName
+  }) {
+    try {
+      const query = {};
+      if (runId) {
+        query["run_id"] = runId;
+      }
+      if (workflowName) {
+        query["workflow_name"] = workflowName;
+      }
+      const collection = await this.getCollection(TABLE_WORKFLOW_SNAPSHOT);
+      const result = await collection.findOne(query);
+      if (!result) {
+        return null;
+      }
+      return this.parseWorkflowRun(result);
+    } catch (error) {
+      console.error("Error getting workflow run by ID:", error);
+      throw error;
+    }
+  }
+  parseWorkflowRun(row) {
+    let parsedSnapshot = row.snapshot;
+    if (typeof parsedSnapshot === "string") {
+      try {
+        parsedSnapshot = JSON.parse(row.snapshot);
+      } catch (e) {
+        console.warn(`Failed to parse snapshot for workflow ${row.workflow_name}: ${e}`);
+      }
+    }
+    return {
+      workflowName: row.workflow_name,
+      runId: row.run_id,
+      snapshot: parsedSnapshot,
+      createdAt: row.createdAt,
+      updatedAt: row.updatedAt,
+      resourceId: row.resourceId
+    };
+  }
+  parseRow(row) {
+    let content = row.content;
+    try {
+      content = JSON.parse(row.content);
+    } catch {
+    }
+    return {
+      id: row.id,
+      content,
+      role: row.role,
+      type: row.type,
+      createdAt: new Date(row.createdAt),
+      threadId: row.thread_id
+    };
+  }
+  transformEvalRow(row) {
+    let testInfoValue = null;
+    if (row.test_info) {
+      try {
+        testInfoValue = typeof row.test_info === "string" ? JSON.parse(row.test_info) : row.test_info;
+      } catch (e) {
+        console.warn("Failed to parse test_info:", e);
+      }
+    }
+    return {
+      input: row.input,
+      output: row.output,
+      result: row.result,
+      agentName: row.agent_name,
+      metricName: row.metric_name,
+      instructions: row.instructions,
+      testInfo: testInfoValue,
+      globalRunId: row.global_run_id,
+      runId: row.run_id,
+      createdAt: row.created_at
+    };
+  }
+  async close() {
+    await this.#client.close();
+  }
+};
 
 // src/vector/prompt.ts
 var MONGODB_PROMPT = `When querying MongoDB Vector, you can ONLY use the operators listed below. Any other operators will be rejected.
@@ -496,4 +1015,4 @@ Example Complex Query:
   ]
 }`;
 
-export { MONGODB_PROMPT, MongoDBVector };
+export { MONGODB_PROMPT, MongoDBStore, MongoDBVector };

package/docker-compose.yaml

@@ -0,0 +1,30 @@
+services:
+  mongodb-storage:
+    image: mongo:6.0.16
+    container_name: 'mongodb-storage-test-db'
+    ports:
+      - '27017:27017'
+    volumes:
+      - mongodbdata:/data/db
+    healthcheck:
+      test: ['CMD', 'mongosh', '--eval', "db.adminCommand('ping')"]
+      interval: 2s
+      timeout: 2s
+      retries: 15
+      start_period: 3s
+  mongodb-vector:
+    image: mongodb/mongodb-atlas-local
+    container_name: 'mongodb-vector-test-db'
+    environment:
+      MONGODB_INITDB_ROOT_USERNAME: mongodb
+      MONGODB_INITDB_ROOT_PASSWORD: mongodb
+    ports:
+      - 27018:27017
+    healthcheck:
+      test: ['CMD', 'mongosh', '-u', 'mongodb', '-p', 'mongodb', '--eval', "db.adminCommand('ping')"]
+      interval: 2s
+      timeout: 2s
+      retries: 15
+      start_period: 3s
+volumes:
+  mongodbdata:

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mastra/mongodb",
-  "version": "0.10.0",
+  "version": "0.10.1-alpha.0",
   "description": "MongoDB provider for Mastra - includes vector store capabilities",
   "type": "module",
   "main": "dist/index.js",
@@ -32,7 +32,7 @@
     "typescript": "^5.8.2",
     "vitest": "^3.1.2",
     "@internal/lint": "0.0.6",
-    "@mastra/core": "0.10.0"
+    "@mastra/core": "0.10.1-alpha.0"
   },
   "peerDependencies": {
     "@mastra/core": "^0.10.0"
@@ -40,7 +40,7 @@
   "scripts": {
     "build": "tsup src/index.ts --format esm,cjs --experimental-dts --clean --treeshake=smallest --splitting",
     "build:watch": "pnpm build --watch",
-    "pretest": "docker compose up -d && (for i in $(seq 1 30); do docker compose exec -T mongodb mongosh --eval 'db.adminCommand(\"ping\")' --quiet && break || (sleep 1; [ $i -eq 30 ] && exit 1); done)",
+    "pretest": "docker compose up -d --wait",
     "test": "vitest run",
     "posttest": "docker compose down -v",
     "pretest:watch": "docker compose up -d",

package/src/index.ts

@@ -1,2 +1,3 @@
 export * from './vector';
+export * from './storage';
 export { MONGODB_PROMPT } from './vector/prompt';