@mastra/upstash 0.10.0 → 0.10.1

package/dist/index.js

@@ -1,4 +1,4 @@
-import { MastraStorage, TABLE_EVALS, TABLE_TRACES, TABLE_MESSAGES, TABLE_WORKFLOW_SNAPSHOT, TABLE_THREADS } from '@mastra/core/storage';
+import { MastraStorage, TABLE_MESSAGES, TABLE_WORKFLOW_SNAPSHOT, TABLE_EVALS, TABLE_TRACES, TABLE_THREADS } from '@mastra/core/storage';
 import { Redis } from '@upstash/redis';
 import { MastraVector } from '@mastra/core/vector';
 import { Index } from '@upstash/vector';
@@ -6,13 +6,157 @@ import { BaseFilterTranslator } from '@mastra/core/vector/filter';
 
 // src/storage/index.ts
 var UpstashStore = class extends MastraStorage {
-  batchInsert(_input) {
-    throw new Error("Method not implemented.");
+  redis;
+  constructor(config) {
+    super({ name: "Upstash" });
+    this.redis = new Redis({
+      url: config.url,
+      token: config.token
+    });
+  }
+  transformEvalRecord(record) {
+    let result = record.result;
+    if (typeof result === "string") {
+      try {
+        result = JSON.parse(result);
+      } catch {
+        console.warn("Failed to parse result JSON:");
+      }
+    }
+    let testInfo = record.test_info;
+    if (typeof testInfo === "string") {
+      try {
+        testInfo = JSON.parse(testInfo);
+      } catch {
+        console.warn("Failed to parse test_info JSON:");
+      }
+    }
+    return {
+      agentName: record.agent_name,
+      input: record.input,
+      output: record.output,
+      result,
+      metricName: record.metric_name,
+      instructions: record.instructions,
+      testInfo,
+      globalRunId: record.global_run_id,
+      runId: record.run_id,
+      createdAt: typeof record.created_at === "string" ? record.created_at : record.created_at instanceof Date ? record.created_at.toISOString() : (/* @__PURE__ */ new Date()).toISOString()
+    };
+  }
+  parseJSON(value) {
+    if (typeof value === "string") {
+      try {
+        return JSON.parse(value);
+      } catch {
+        return value;
+      }
+    }
+    return value;
+  }
+  getKey(tableName, keys) {
+    const keyParts = Object.entries(keys).filter(([_, value]) => value !== void 0).map(([key, value]) => `${key}:${value}`);
+    return `${tableName}:${keyParts.join(":")}`;
+  }
+  ensureDate(date) {
+    if (!date) return void 0;
+    return date instanceof Date ? date : new Date(date);
+  }
+  serializeDate(date) {
+    if (!date) return void 0;
+    const dateObj = this.ensureDate(date);
+    return dateObj?.toISOString();
+  }
+  /**
+   * Scans for keys matching the given pattern using SCAN and returns them as an array.
+   * @param pattern Redis key pattern, e.g. "table:*"
+   * @param batchSize Number of keys to scan per batch (default: 1000)
+   */
+  async scanKeys(pattern, batchSize = 1e4) {
+    let cursor = "0";
+    let keys = [];
+    do {
+      const [nextCursor, batch] = await this.redis.scan(cursor, {
+        match: pattern,
+        count: batchSize
+      });
+      keys.push(...batch);
+      cursor = nextCursor;
+    } while (cursor !== "0");
+    return keys;
+  }
+  /**
+   * Deletes all keys matching the given pattern using SCAN and DEL in batches.
+   * @param pattern Redis key pattern, e.g. "table:*"
+   * @param batchSize Number of keys to delete per batch (default: 1000)
+   */
+  async scanAndDelete(pattern, batchSize = 1e4) {
+    let cursor = "0";
+    let totalDeleted = 0;
+    do {
+      const [nextCursor, keys] = await this.redis.scan(cursor, {
+        match: pattern,
+        count: batchSize
+      });
+      if (keys.length > 0) {
+        await this.redis.del(...keys);
+        totalDeleted += keys.length;
+      }
+      cursor = nextCursor;
+    } while (cursor !== "0");
+    return totalDeleted;
+  }
+  getMessageKey(threadId, messageId) {
+    return this.getKey(TABLE_MESSAGES, { threadId, id: messageId });
+  }
+  getThreadMessagesKey(threadId) {
+    return `thread:${threadId}:messages`;
+  }
+  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: this.ensureDate(row.createdAt),
+      updatedAt: this.ensureDate(row.updatedAt),
+      resourceId: row.resourceId
+    };
+  }
+  processRecord(tableName, record) {
+    let key;
+    if (tableName === TABLE_MESSAGES) {
+      key = this.getKey(tableName, { threadId: record.threadId, id: record.id });
+    } else if (tableName === TABLE_WORKFLOW_SNAPSHOT) {
+      key = this.getKey(tableName, {
+        namespace: record.namespace || "workflows",
+        workflow_name: record.workflow_name,
+        run_id: record.run_id,
+        ...record.resourceId ? { resourceId: record.resourceId } : {}
+      });
+    } else if (tableName === TABLE_EVALS) {
+      key = this.getKey(tableName, { id: record.run_id });
+    } else {
+      key = this.getKey(tableName, { id: record.id });
+    }
+    const processedRecord = {
+      ...record,
+      createdAt: this.serializeDate(record.createdAt),
+      updatedAt: this.serializeDate(record.updatedAt)
+    };
+    return { key, processedRecord };
   }
   async getEvalsByAgentName(agentName, type) {
     try {
       const pattern = `${TABLE_EVALS}:*`;
-      const keys = await this.redis.keys(pattern);
+      const keys = await this.scanKeys(pattern);
       const evalRecords = await Promise.all(
         keys.map(async (key) => {
           const data = await this.redis.get(key);
@@ -56,36 +200,6 @@ var UpstashStore = class extends MastraStorage {
       return [];
     }
   }
-  transformEvalRecord(record) {
-    let result = record.result;
-    if (typeof result === "string") {
-      try {
-        result = JSON.parse(result);
-      } catch {
-        console.warn("Failed to parse result JSON:");
-      }
-    }
-    let testInfo = record.test_info;
-    if (typeof testInfo === "string") {
-      try {
-        testInfo = JSON.parse(testInfo);
-      } catch {
-        console.warn("Failed to parse test_info JSON:");
-      }
-    }
-    return {
-      agentName: record.agent_name,
-      input: record.input,
-      output: record.output,
-      result,
-      metricName: record.metric_name,
-      instructions: record.instructions,
-      testInfo,
-      globalRunId: record.global_run_id,
-      runId: record.run_id,
-      createdAt: typeof record.created_at === "string" ? record.created_at : record.created_at instanceof Date ? record.created_at.toISOString() : (/* @__PURE__ */ new Date()).toISOString()
-    };
-  }
   async getTraces({
     name,
     scope,
@@ -101,7 +215,7 @@ var UpstashStore = class extends MastraStorage {
   }) {
     try {
       const pattern = `${TABLE_TRACES}:*`;
-      const keys = await this.redis.keys(pattern);
+      const keys = await this.scanKeys(pattern);
       const traceRecords = await Promise.all(
         keys.map(async (key) => {
           const data = await this.redis.get(key);
@@ -165,37 +279,6 @@ var UpstashStore = class extends MastraStorage {
       return [];
     }
   }
-  parseJSON(value) {
-    if (typeof value === "string") {
-      try {
-        return JSON.parse(value);
-      } catch {
-        return value;
-      }
-    }
-    return value;
-  }
-  redis;
-  constructor(config) {
-    super({ name: "Upstash" });
-    this.redis = new Redis({
-      url: config.url,
-      token: config.token
-    });
-  }
-  getKey(tableName, keys) {
-    const keyParts = Object.entries(keys).filter(([_, value]) => value !== void 0).map(([key, value]) => `${key}:${value}`);
-    return `${tableName}:${keyParts.join(":")}`;
-  }
-  ensureDate(date) {
-    if (!date) return void 0;
-    return date instanceof Date ? date : new Date(date);
-  }
-  serializeDate(date) {
-    if (!date) return void 0;
-    const dateObj = this.ensureDate(date);
-    return dateObj?.toISOString();
-  }
   async createTable({
     tableName,
     schema
@@ -204,34 +287,26 @@ var UpstashStore = class extends MastraStorage {
   }
   async clearTable({ tableName }) {
     const pattern = `${tableName}:*`;
-    const keys = await this.redis.keys(pattern);
-    if (keys.length > 0) {
-      await this.redis.del(...keys);
-    }
+    await this.scanAndDelete(pattern);
   }
   async insert({ tableName, record }) {
-    let key;
-    if (tableName === TABLE_MESSAGES) {
-      key = this.getKey(tableName, { threadId: record.threadId, id: record.id });
-    } else if (tableName === TABLE_WORKFLOW_SNAPSHOT) {
-      key = this.getKey(tableName, {
-        namespace: record.namespace || "workflows",
-        workflow_name: record.workflow_name,
-        run_id: record.run_id,
-        ...record.resourceId ? { resourceId: record.resourceId } : {}
-      });
-    } else if (tableName === TABLE_EVALS) {
-      key = this.getKey(tableName, { id: record.run_id });
-    } else {
-      key = this.getKey(tableName, { id: record.id });
-    }
-    const processedRecord = {
-      ...record,
-      createdAt: this.serializeDate(record.createdAt),
-      updatedAt: this.serializeDate(record.updatedAt)
-    };
+    const { key, processedRecord } = this.processRecord(tableName, record);
     await this.redis.set(key, processedRecord);
   }
+  async batchInsert(input) {
+    const { tableName, records } = input;
+    if (!records.length) return;
+    const batchSize = 1e3;
+    for (let i = 0; i < records.length; i += batchSize) {
+      const batch = records.slice(i, i + batchSize);
+      const pipeline = this.redis.pipeline();
+      for (const record of batch) {
+        const { key, processedRecord } = this.processRecord(tableName, record);
+        pipeline.set(key, processedRecord);
+      }
+      await pipeline.exec();
+    }
+  }
   async load({ tableName, keys }) {
     const key = this.getKey(tableName, keys);
     const data = await this.redis.get(key);
@@ -252,7 +327,7 @@ var UpstashStore = class extends MastraStorage {
   }
   async getThreadsByResourceId({ resourceId }) {
     const pattern = `${TABLE_THREADS}:*`;
-    const keys = await this.redis.keys(pattern);
+    const keys = await this.scanKeys(pattern);
     const threads = await Promise.all(
       keys.map(async (key) => {
         const data = await this.redis.get(key);
@@ -297,29 +372,27 @@ var UpstashStore = class extends MastraStorage {
     const key = this.getKey(TABLE_THREADS, { id: threadId });
     await this.redis.del(key);
   }
-  getMessageKey(threadId, messageId) {
-    return this.getKey(TABLE_MESSAGES, { threadId, id: messageId });
-  }
-  getThreadMessagesKey(threadId) {
-    return `thread:${threadId}:messages`;
-  }
   async saveMessages({ messages }) {
     if (messages.length === 0) return [];
-    const pipeline = this.redis.pipeline();
     const messagesWithIndex = messages.map((message, index) => ({
       ...message,
       _index: index
     }));
-    for (const message of messagesWithIndex) {
-      const key = this.getMessageKey(message.threadId, message.id);
-      const score = message._index !== void 0 ? message._index : new Date(message.createdAt).getTime();
-      pipeline.set(key, message);
-      pipeline.zadd(this.getThreadMessagesKey(message.threadId), {
-        score,
-        member: message.id
-      });
+    const batchSize = 1e3;
+    for (let i = 0; i < messagesWithIndex.length; i += batchSize) {
+      const batch = messagesWithIndex.slice(i, i + batchSize);
+      const pipeline = this.redis.pipeline();
+      for (const message of batch) {
+        const key = this.getMessageKey(message.threadId, message.id);
+        const score = message._index !== void 0 ? message._index : new Date(message.createdAt).getTime();
+        pipeline.set(key, message);
+        pipeline.zadd(this.getThreadMessagesKey(message.threadId), {
+          score,
+          member: message.id
+        });
+      }
+      await pipeline.exec();
     }
-    await pipeline.exec();
     return messages;
   }
   async getMessages({ threadId, selectBy }) {
@@ -383,24 +456,6 @@ var UpstashStore = class extends MastraStorage {
     if (!data) return null;
     return data.snapshot;
   }
-  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: this.ensureDate(row.createdAt),
-      updatedAt: this.ensureDate(row.updatedAt),
-      resourceId: row.resourceId
-    };
-  }
   async getWorkflowRuns({
     namespace,
     workflowName,
@@ -424,7 +479,7 @@ var UpstashStore = class extends MastraStorage {
       } else if (resourceId) {
         pattern = this.getKey(TABLE_WORKFLOW_SNAPSHOT, { namespace, workflow_name: "*", run_id: "*", resourceId });
       }
-      const keys = await this.redis.keys(pattern);
+      const keys = await this.scanKeys(pattern);
       const workflows = await Promise.all(
         keys.map(async (key) => {
           const data = await this.redis.get(key);
@@ -453,7 +508,7 @@ var UpstashStore = class extends MastraStorage {
   }) {
     try {
       const key = this.getKey(TABLE_WORKFLOW_SNAPSHOT, { namespace, workflow_name: workflowName, run_id: runId }) + "*";
-      const keys = await this.redis.keys(key);
+      const keys = await this.scanKeys(key);
       const workflows = await Promise.all(
         keys.map(async (key2) => {
           const data2 = await this.redis.get(key2);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mastra/upstash",
-  "version": "0.10.0",
+  "version": "0.10.1",
   "description": "Upstash provider for Mastra - includes both vector and db storage capabilities",
   "type": "module",
   "main": "dist/index.js",
@@ -31,8 +31,8 @@
     "tsup": "^8.4.0",
     "typescript": "^5.8.2",
     "vitest": "^3.1.2",
-    "@mastra/core": "0.10.0",
-    "@internal/lint": "0.0.6"
+    "@internal/lint": "0.0.7",
+    "@mastra/core": "0.10.1"
   },
   "peerDependencies": {
     "@mastra/core": "^0.10.0"