@mastra/upstash 1.0.4 → 1.0.5-alpha.0

package/dist/index.js

@@ -1,7 +1,7 @@
-import { MemoryStorage, TABLE_THREADS, TABLE_MESSAGES, TABLE_RESOURCES, ensureDate, createStorageErrorId, normalizePerPage, calculatePagination, filterByDateRange, ScoresStorage, TABLE_SCORERS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, MastraCompositeStore, createVectorErrorId, serializeDate, transformScoreRow as transformScoreRow$1 } from '@mastra/core/storage';
+import { BackgroundTasksStorage, TABLE_BACKGROUND_TASKS, MemoryStorage, TABLE_THREADS, TABLE_MESSAGES, TABLE_RESOURCES, ensureDate, createStorageErrorId, normalizePerPage, calculatePagination, filterByDateRange, ScoresStorage, TABLE_SCORERS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, MastraCompositeStore, createVectorErrorId, serializeDate, transformScoreRow as transformScoreRow$1 } from '@mastra/core/storage';
 import { Redis } from '@upstash/redis';
-import { MessageList } from '@mastra/core/agent';
 import { MastraError, ErrorCategory, ErrorDomain } from '@mastra/core/error';
+import { MessageList } from '@mastra/core/agent';
 import { saveScorePayloadSchema } from '@mastra/core/evals';
 import { randomUUID } from 'crypto';
 import { MastraVector } from '@mastra/core/vector';
@@ -138,7 +138,164 @@ var UpstashDB = class {
   }
 };
 
-// src/storage/domains/memory/index.ts
+// src/storage/domains/background-tasks/index.ts
+function toStorageRecord(task) {
+  return {
+    id: task.id,
+    tool_call_id: task.toolCallId,
+    tool_name: task.toolName,
+    agent_id: task.agentId,
+    thread_id: task.threadId ?? null,
+    resource_id: task.resourceId ?? null,
+    run_id: task.runId,
+    status: task.status,
+    args: task.args,
+    result: task.result ?? null,
+    error: task.error ?? null,
+    retry_count: task.retryCount,
+    max_retries: task.maxRetries,
+    timeout_ms: task.timeoutMs,
+    createdAt: task.createdAt.toISOString(),
+    startedAt: task.startedAt?.toISOString() ?? null,
+    completedAt: task.completedAt?.toISOString() ?? null
+  };
+}
+function fromStorageRecord(record) {
+  return {
+    id: record.id,
+    status: record.status,
+    toolName: record.tool_name,
+    toolCallId: record.tool_call_id,
+    args: record.args ?? {},
+    agentId: record.agent_id,
+    threadId: record.thread_id ?? void 0,
+    resourceId: record.resource_id ?? void 0,
+    runId: record.run_id ?? "",
+    result: record.result ?? void 0,
+    error: record.error ?? void 0,
+    retryCount: Number(record.retry_count ?? 0),
+    maxRetries: Number(record.max_retries ?? 0),
+    timeoutMs: Number(record.timeout_ms ?? 3e5),
+    createdAt: new Date(record.createdAt),
+    startedAt: record.startedAt ? new Date(record.startedAt) : void 0,
+    completedAt: record.completedAt ? new Date(record.completedAt) : void 0
+  };
+}
+var BackgroundTasksUpstash = class extends BackgroundTasksStorage {
+  client;
+  #db;
+  constructor(config) {
+    super();
+    const client = resolveUpstashConfig(config);
+    this.client = client;
+    this.#db = new UpstashDB({ client });
+  }
+  async dangerouslyClearAll() {
+    await this.#db.deleteData({ tableName: TABLE_BACKGROUND_TASKS });
+  }
+  async createTask(task) {
+    const record = toStorageRecord(task);
+    const { key, processedRecord } = processRecord(TABLE_BACKGROUND_TASKS, record);
+    await this.client.set(key, processedRecord);
+  }
+  async updateTask(taskId, update) {
+    const existing = await this.getTask(taskId);
+    if (!existing) return;
+    const merged = { ...existing };
+    if ("status" in update) merged.status = update.status;
+    if ("result" in update) merged.result = update.result;
+    if ("error" in update) merged.error = update.error;
+    if ("retryCount" in update) merged.retryCount = update.retryCount;
+    if ("startedAt" in update) merged.startedAt = update.startedAt;
+    if ("completedAt" in update) merged.completedAt = update.completedAt;
+    const record = toStorageRecord(merged);
+    const { key, processedRecord } = processRecord(TABLE_BACKGROUND_TASKS, record);
+    await this.client.set(key, processedRecord);
+  }
+  async getTask(taskId) {
+    const key = getKey(TABLE_BACKGROUND_TASKS, { id: taskId });
+    const data = await this.client.get(key);
+    if (!data) return null;
+    return fromStorageRecord(data);
+  }
+  async listTasks(filter) {
+    const keys = await this.#db.scanKeys(`${TABLE_BACKGROUND_TASKS}:*`);
+    if (keys.length === 0) return { tasks: [], total: 0 };
+    const pipeline = this.client.pipeline();
+    for (const key of keys) {
+      pipeline.get(key);
+    }
+    const results = await pipeline.exec();
+    let tasks = results.filter(Boolean).map((r) => fromStorageRecord(r));
+    if (filter.status) {
+      const statuses = Array.isArray(filter.status) ? filter.status : [filter.status];
+      tasks = tasks.filter((t) => statuses.includes(t.status));
+    }
+    if (filter.agentId) {
+      tasks = tasks.filter((t) => t.agentId === filter.agentId);
+    }
+    if (filter.threadId) {
+      tasks = tasks.filter((t) => t.threadId === filter.threadId);
+    }
+    if (filter.resourceId) {
+      tasks = tasks.filter((t) => t.resourceId === filter.resourceId);
+    }
+    if (filter.runId) {
+      tasks = tasks.filter((t) => t.runId === filter.runId);
+    }
+    if (filter.toolName) {
+      tasks = tasks.filter((t) => t.toolName === filter.toolName);
+    }
+    const dateCol = filter.dateFilterBy ?? "createdAt";
+    if (filter.fromDate) {
+      tasks = tasks.filter((t) => {
+        const val = t[dateCol];
+        return val != null && val >= filter.fromDate;
+      });
+    }
+    if (filter.toDate) {
+      tasks = tasks.filter((t) => {
+        const val = t[dateCol];
+        return val != null && val < filter.toDate;
+      });
+    }
+    const orderBy = filter.orderBy ?? "createdAt";
+    const direction = filter.orderDirection ?? "asc";
+    tasks.sort((a, b) => {
+      const aVal = a[orderBy]?.getTime() ?? 0;
+      const bVal = b[orderBy]?.getTime() ?? 0;
+      return direction === "asc" ? aVal - bVal : bVal - aVal;
+    });
+    const total = tasks.length;
+    if (filter.page != null && filter.perPage != null) {
+      const start = filter.page * filter.perPage;
+      tasks = tasks.slice(start, start + filter.perPage);
+    } else if (filter.perPage != null) {
+      tasks = tasks.slice(0, filter.perPage);
+    }
+    return { tasks, total };
+  }
+  async deleteTask(taskId) {
+    const key = getKey(TABLE_BACKGROUND_TASKS, { id: taskId });
+    await this.client.del(key);
+  }
+  async deleteTasks(filter) {
+    const { tasks } = await this.listTasks(filter);
+    if (tasks.length === 0) return;
+    const keys = tasks.map((t) => getKey(TABLE_BACKGROUND_TASKS, { id: t.id }));
+    if (keys.length > 0) {
+      await this.client.del(...keys);
+    }
+  }
+  async getRunningCount() {
+    const { total } = await this.listTasks({ status: "running" });
+    return total;
+  }
+  async getRunningCountByAgent(agentId) {
+    const { total } = await this.listTasks({ status: "running", agentId });
+    return total;
+  }
+};
 function getThreadMessagesKey(threadId) {
   return `thread:${threadId}:messages`;
 }
@@ -391,14 +548,16 @@ var StoreMemoryUpstash = class extends MemoryStorage {
   async saveMessages(args) {
     const { messages } = args;
     if (messages.length === 0) return { messages: [] };
-    const threadId = messages[0]?.threadId;
     try {
-      if (!threadId) {
-        throw new Error("Thread ID is required");
-      }
-      const thread = await this.getThreadById({ threadId });
-      if (!thread) {
-        throw new Error(`Thread ${threadId} not found`);
+      for (const message of messages) {
+        if (!message.threadId) {
+          throw new Error("Thread ID is required");
+        }
+        if (!message.resourceId) {
+          throw new Error(
+            `Expected to find a resourceId for message, but couldn't find one. An unexpected error has occurred.`
+          );
+        }
       }
     } catch (error) {
       throw new MastraError(
@@ -411,46 +570,70 @@ var StoreMemoryUpstash = class extends MemoryStorage {
       );
     }
     const messagesWithIndex = messages.map((message, index) => {
-      if (!message.threadId) {
-        throw new Error(
-          `Expected to find a threadId for message, but couldn't find one. An unexpected error has occurred.`
-        );
-      }
-      if (!message.resourceId) {
-        throw new Error(
-          `Expected to find a resourceId for message, but couldn't find one. An unexpected error has occurred.`
-        );
-      }
       return {
         ...message,
         _index: index
       };
     });
-    const threadKey = getKey(TABLE_THREADS, { id: threadId });
-    const existingThread = await this.client.get(threadKey);
     try {
       const batchSize = 1e3;
+      const targetThreadIds = new Set(messagesWithIndex.map((message) => message.threadId));
+      const existingThreadIds = [];
+      const touchedThreadIds = new Set(targetThreadIds);
+      for (let i = 0; i < messagesWithIndex.length; i += batchSize) {
+        const batch = messagesWithIndex.slice(i, i + batchSize);
+        const indexLookupPipeline = this.client.pipeline();
+        batch.forEach((message) => {
+          indexLookupPipeline.get(getMessageIndexKey(message.id));
+        });
+        const batchExistingThreadIds = await indexLookupPipeline.exec();
+        existingThreadIds.push(...batchExistingThreadIds);
+        batchExistingThreadIds.forEach((existingThreadId) => {
+          if (existingThreadId) {
+            touchedThreadIds.add(existingThreadId);
+          }
+        });
+      }
+      const touchedThreadIdList = Array.from(touchedThreadIds);
+      const threadLookupPipeline = this.client.pipeline();
+      touchedThreadIdList.forEach((touchedThreadId) => {
+        threadLookupPipeline.get(getKey(TABLE_THREADS, { id: touchedThreadId }));
+      });
+      const threadLookupResults = await threadLookupPipeline.exec();
+      const threadRecordsById = /* @__PURE__ */ new Map();
+      touchedThreadIdList.forEach((touchedThreadId, index) => {
+        const threadRecord = threadLookupResults[index];
+        if (threadRecord) {
+          threadRecordsById.set(touchedThreadId, threadRecord);
+        }
+      });
+      for (const targetThreadId of targetThreadIds) {
+        if (!threadRecordsById.has(targetThreadId)) {
+          throw new MastraError(
+            {
+              id: createStorageErrorId("UPSTASH", "SAVE_MESSAGES", "INVALID_ARGS"),
+              domain: ErrorDomain.STORAGE,
+              category: ErrorCategory.USER
+            },
+            new Error(`Thread ${targetThreadId} not found`)
+          );
+        }
+      }
       for (let i = 0; i < messagesWithIndex.length; i += batchSize) {
         const batch = messagesWithIndex.slice(i, i + batchSize);
         const pipeline = this.client.pipeline();
-        for (const message of batch) {
+        const batchTouchedThreadIds = /* @__PURE__ */ new Set();
+        const batchExistingThreadIds = existingThreadIds.slice(i, i + batch.length);
+        for (const [batchIndex, message] of batch.entries()) {
           const key = getMessageKey(message.threadId, message.id);
           const createdAtScore = new Date(message.createdAt).getTime();
           const score = message._index !== void 0 ? message._index : createdAtScore;
-          const existingKeyPattern = getMessageKey("*", message.id);
-          const keys = await this.#db.scanKeys(existingKeyPattern);
-          if (keys.length > 0) {
-            const pipeline2 = this.client.pipeline();
-            keys.forEach((key2) => pipeline2.get(key2));
-            const results = await pipeline2.exec();
-            const existingMessages = results.filter((msg) => msg !== null);
-            for (const existingMessage of existingMessages) {
-              const existingMessageKey = getMessageKey(existingMessage.threadId, existingMessage.id);
-              if (existingMessage && existingMessage.threadId !== message.threadId) {
-                pipeline.del(existingMessageKey);
-                pipeline.zrem(getThreadMessagesKey(existingMessage.threadId), existingMessage.id);
-              }
-            }
+          batchTouchedThreadIds.add(message.threadId);
+          const existingThreadId = batchExistingThreadIds[batchIndex];
+          if (existingThreadId && existingThreadId !== message.threadId) {
+            pipeline.del(getMessageKey(existingThreadId, message.id));
+            pipeline.zrem(getThreadMessagesKey(existingThreadId), message.id);
+            batchTouchedThreadIds.add(existingThreadId);
           }
           pipeline.set(key, message);
           pipeline.set(getMessageIndexKey(message.id), message.threadId);
@@ -459,25 +642,35 @@ var StoreMemoryUpstash = class extends MemoryStorage {
             member: message.id
           });
         }
-        if (i === 0 && existingThread) {
+        const now = /* @__PURE__ */ new Date();
+        for (const touchedThreadId of batchTouchedThreadIds) {
+          const existingThread = threadRecordsById.get(touchedThreadId);
+          if (!existingThread) {
+            continue;
+          }
           const updatedThread = {
             ...existingThread,
-            updatedAt: /* @__PURE__ */ new Date()
+            updatedAt: now
           };
+          const threadKey = getKey(TABLE_THREADS, { id: touchedThreadId });
           pipeline.set(threadKey, processRecord(TABLE_THREADS, updatedThread).processedRecord);
+          threadRecordsById.set(touchedThreadId, updatedThread);
         }
         await pipeline.exec();
       }
       const list = new MessageList().add(messages, "memory");
       return { messages: list.get.all.db() };
     } catch (error) {
+      if (error instanceof MastraError) {
+        throw error;
+      }
       throw new MastraError(
         {
           id: createStorageErrorId("UPSTASH", "SAVE_MESSAGES", "FAILED"),
           domain: ErrorDomain.STORAGE,
           category: ErrorCategory.THIRD_PARTY,
           details: {
-            threadId
+            threadIds: Array.from(new Set(messages.map((message) => message.threadId).filter(Boolean))).join(",")
           }
         },
         error
@@ -834,9 +1027,41 @@ var StoreMemoryUpstash = class extends MemoryStorage {
     }
     try {
       const messageIds = messages.map((m) => m.id);
+      const updatesById = new Map(messages.map((message) => [message.id, message]));
       const existingMessages = [];
       const messageIdToKey = {};
-      for (const messageId of messageIds) {
+      const backfillIndexValues = {};
+      const indexPipeline = this.client.pipeline();
+      messageIds.forEach((messageId) => indexPipeline.get(getMessageIndexKey(messageId)));
+      const indexResults = await indexPipeline.exec();
+      const indexedLookups = [];
+      const fallbackMessageIds = [];
+      messageIds.forEach((messageId, index) => {
+        const indexedThreadId = indexResults[index];
+        if (indexedThreadId) {
+          indexedLookups.push({ messageId, threadId: indexedThreadId });
+        } else {
+          fallbackMessageIds.push(messageId);
+        }
+      });
+      if (indexedLookups.length > 0) {
+        const messagePipeline = this.client.pipeline();
+        indexedLookups.forEach(({ messageId, threadId }) => {
+          messagePipeline.get(getMessageKey(threadId, messageId));
+        });
+        const indexedMessages = await messagePipeline.exec();
+        indexedLookups.forEach(({ messageId, threadId }, index) => {
+          const key = getMessageKey(threadId, messageId);
+          const message = indexedMessages[index];
+          if (message && message.id === messageId) {
+            existingMessages.push(message);
+            messageIdToKey[messageId] = key;
+          } else {
+            fallbackMessageIds.push(messageId);
+          }
+        });
+      }
+      for (const messageId of fallbackMessageIds) {
         const pattern = getMessageKey("*", messageId);
         const keys = await this.#db.scanKeys(pattern);
         for (const key of keys) {
@@ -844,24 +1069,69 @@ var StoreMemoryUpstash = class extends MemoryStorage {
           if (message && message.id === messageId) {
             existingMessages.push(message);
             messageIdToKey[messageId] = key;
+            if (message.threadId) {
+              backfillIndexValues[messageId] = message.threadId;
+            }
             break;
           }
         }
       }
+      if (Object.keys(backfillIndexValues).length > 0) {
+        const backfillPipeline = this.client.pipeline();
+        for (const [messageId, threadId] of Object.entries(backfillIndexValues)) {
+          backfillPipeline.set(getMessageIndexKey(messageId), threadId);
+        }
+        await backfillPipeline.exec();
+      }
       if (existingMessages.length === 0) {
         return [];
       }
       const threadIdsToUpdate = /* @__PURE__ */ new Set();
-      const pipeline = this.client.pipeline();
+      const destinationThreadIds = /* @__PURE__ */ new Set();
       for (const existingMessage of existingMessages) {
-        const updatePayload = messages.find((m) => m.id === existingMessage.id);
+        const updatePayload = updatesById.get(existingMessage.id);
         if (!updatePayload) continue;
-        const { id, ...fieldsToUpdate } = updatePayload;
+        const { id: _id, ...fieldsToUpdate } = updatePayload;
         if (Object.keys(fieldsToUpdate).length === 0) continue;
         threadIdsToUpdate.add(existingMessage.threadId);
         if (updatePayload.threadId && updatePayload.threadId !== existingMessage.threadId) {
           threadIdsToUpdate.add(updatePayload.threadId);
+          destinationThreadIds.add(updatePayload.threadId);
         }
+      }
+      const threadRecordsById = /* @__PURE__ */ new Map();
+      if (threadIdsToUpdate.size > 0) {
+        const threadIdList = Array.from(threadIdsToUpdate);
+        const threadLookupPipeline = this.client.pipeline();
+        threadIdList.forEach((threadId) => {
+          threadLookupPipeline.get(getKey(TABLE_THREADS, { id: threadId }));
+        });
+        const threadLookupResults = await threadLookupPipeline.exec();
+        threadIdList.forEach((threadId, index) => {
+          const threadRecord = threadLookupResults[index];
+          if (threadRecord) {
+            threadRecordsById.set(threadId, threadRecord);
+          }
+        });
+      }
+      for (const destinationThreadId of destinationThreadIds) {
+        if (!threadRecordsById.has(destinationThreadId)) {
+          throw new MastraError(
+            {
+              id: createStorageErrorId("UPSTASH", "UPDATE_MESSAGES", "INVALID_ARGS"),
+              domain: ErrorDomain.STORAGE,
+              category: ErrorCategory.USER
+            },
+            new Error(`Thread ${destinationThreadId} not found`)
+          );
+        }
+      }
+      const pipeline = this.client.pipeline();
+      for (const existingMessage of existingMessages) {
+        const updatePayload = updatesById.get(existingMessage.id);
+        if (!updatePayload) continue;
+        const { id, ...fieldsToUpdate } = updatePayload;
+        if (Object.keys(fieldsToUpdate).length === 0) continue;
         const updatedMessage = { ...existingMessage };
         if (fieldsToUpdate.content) {
           const existingContent = existingMessage.content;
@@ -886,12 +1156,15 @@ var StoreMemoryUpstash = class extends MemoryStorage {
         const key = messageIdToKey[id];
         if (key) {
           if (updatePayload.threadId && updatePayload.threadId !== existingMessage.threadId) {
+            const newThreadId = updatedMessage.threadId;
             const oldThreadMessagesKey = getThreadMessagesKey(existingMessage.threadId);
             pipeline.zrem(oldThreadMessagesKey, id);
             pipeline.del(key);
-            const newKey = getMessageKey(updatePayload.threadId, id);
+            const newKey = getMessageKey(newThreadId, id);
             pipeline.set(newKey, updatedMessage);
-            const newThreadMessagesKey = getThreadMessagesKey(updatePayload.threadId);
+            pipeline.set(getMessageIndexKey(id), newThreadId);
+            messageIdToKey[id] = newKey;
+            const newThreadMessagesKey = getThreadMessagesKey(newThreadId);
             const score = updatedMessage._index !== void 0 ? updatedMessage._index : new Date(updatedMessage.createdAt).getTime();
             pipeline.zadd(newThreadMessagesKey, { score, member: id });
           } else {
@@ -902,14 +1175,15 @@ var StoreMemoryUpstash = class extends MemoryStorage {
       const now = /* @__PURE__ */ new Date();
       for (const threadId of threadIdsToUpdate) {
         if (threadId) {
-          const threadKey = getKey(TABLE_THREADS, { id: threadId });
-          const existingThread = await this.client.get(threadKey);
+          const existingThread = threadRecordsById.get(threadId);
           if (existingThread) {
             const updatedThread = {
               ...existingThread,
               updatedAt: now
             };
+            const threadKey = getKey(TABLE_THREADS, { id: threadId });
             pipeline.set(threadKey, processRecord(TABLE_THREADS, updatedThread).processedRecord);
+            threadRecordsById.set(threadId, updatedThread);
           }
         }
       }
@@ -926,6 +1200,9 @@ var StoreMemoryUpstash = class extends MemoryStorage {
       }
       return updatedMessages;
     } catch (error) {
+      if (error instanceof MastraError) {
+        throw error;
+      }
       throw new MastraError(
         {
           id: createStorageErrorId("UPSTASH", "UPDATE_MESSAGES", "FAILED"),
@@ -1882,10 +2159,12 @@ var UpstashStore = class extends MastraCompositeStore {
     const scores = new ScoresUpstash({ client: this.redis });
     const workflows = new WorkflowsUpstash({ client: this.redis });
     const memory = new StoreMemoryUpstash({ client: this.redis });
+    const backgroundTasks = new BackgroundTasksUpstash({ client: this.redis });
     this.stores = {
       scores,
       workflows,
-      memory
+      memory,
+      backgroundTasks
     };
   }
   async close() {
@@ -2591,6 +2870,6 @@ Example Complex Query:
   ]
 }`;
 
-export { ScoresUpstash, StoreMemoryUpstash, UPSTASH_PROMPT, UpstashStore, UpstashVector, WorkflowsUpstash };
+export { BackgroundTasksUpstash, ScoresUpstash, StoreMemoryUpstash, UPSTASH_PROMPT, UpstashStore, UpstashVector, WorkflowsUpstash };
 //# sourceMappingURL=index.js.map
 //# sourceMappingURL=index.js.map