@mastra/dynamodb 1.0.3 → 1.0.4-alpha.1

package/dist/index.js

@@ -1,7 +1,7 @@
 import { DynamoDBClient, DescribeTableCommand } from '@aws-sdk/client-dynamodb';
 import { DynamoDBDocumentClient } from '@aws-sdk/lib-dynamodb';
 import { MastraError, ErrorCategory, ErrorDomain } from '@mastra/core/error';
-import { MemoryStorage, createStorageErrorId, normalizePerPage, calculatePagination, filterByDateRange, ScoresStorage, SCORERS_SCHEMA, WorkflowsStorage, MastraCompositeStore, TABLE_SCORERS, TABLE_WORKFLOW_SNAPSHOT, TABLE_RESOURCES, TABLE_MESSAGES, TABLE_THREADS } from '@mastra/core/storage';
+import { BackgroundTasksStorage, createStorageErrorId, MemoryStorage, normalizePerPage, calculatePagination, filterByDateRange, ScoresStorage, SCORERS_SCHEMA, WorkflowsStorage, MastraCompositeStore, TABLE_BACKGROUND_TASKS, TABLE_SCORERS, TABLE_WORKFLOW_SNAPSHOT, TABLE_RESOURCES, TABLE_MESSAGES, TABLE_THREADS } from '@mastra/core/storage';
 import { Entity, Service } from 'electrodb';
 import { MessageList } from '@mastra/core/agent';
 import { saveScorePayloadSchema } from '@mastra/core/evals';
@@ -80,7 +80,128 @@ var baseAttributes = {
   }
 };
 
-// src/entities/eval.ts
+// src/entities/background-task.ts
+var jsonSetGet = {
+  set: (value) => {
+    if (value == null) return value;
+    if (typeof value === "string") return value;
+    return JSON.stringify(value);
+  },
+  get: (value) => {
+    if (!value) return value;
+    if (typeof value !== "string") return value;
+    try {
+      if (value.startsWith("{") || value.startsWith("[")) {
+        return JSON.parse(value);
+      }
+    } catch {
+    }
+    return value;
+  }
+};
+var backgroundTaskEntity = new Entity({
+  model: {
+    entity: "background_task",
+    version: "1",
+    service: "mastra"
+  },
+  attributes: {
+    entity: {
+      type: "string",
+      required: true
+    },
+    ...baseAttributes,
+    id: {
+      type: "string",
+      required: true
+    },
+    status: {
+      type: "string",
+      required: true
+    },
+    toolName: {
+      type: "string",
+      required: true
+    },
+    toolCallId: {
+      type: "string",
+      required: true
+    },
+    agentId: {
+      type: "string",
+      required: true
+    },
+    runId: {
+      type: "string",
+      required: true
+    },
+    threadId: {
+      type: "string",
+      required: false
+    },
+    resourceId: {
+      type: "string",
+      required: false
+    },
+    args: {
+      type: "string",
+      required: true,
+      ...jsonSetGet
+    },
+    result: {
+      type: "string",
+      required: false,
+      ...jsonSetGet
+    },
+    error: {
+      type: "string",
+      required: false,
+      ...jsonSetGet
+    },
+    retryCount: {
+      type: "number",
+      required: true
+    },
+    maxRetries: {
+      type: "number",
+      required: true
+    },
+    timeoutMs: {
+      type: "number",
+      required: true
+    },
+    // Separate ISO strings for date columns — use these for GSI sort keys
+    startedAtIso: {
+      type: "string",
+      required: false
+    },
+    completedAtIso: {
+      type: "string",
+      required: false
+    }
+  },
+  indexes: {
+    primary: {
+      pk: { field: "pk", composite: ["entity", "id"] },
+      sk: { field: "sk", composite: ["entity"] }
+    },
+    byAgent: {
+      index: "gsi1",
+      pk: { field: "gsi1pk", composite: ["entity", "agentId"] },
+      sk: { field: "gsi1sk", composite: ["createdAt"] }
+    },
+    byRun: {
+      index: "gsi2",
+      pk: { field: "gsi2pk", composite: ["entity", "runId"] },
+      sk: { field: "gsi2sk", composite: ["createdAt"] }
+    },
+    byStatus: {
+      index: "gsi3",
+      pk: { field: "gsi3pk", composite: ["entity", "status"] },
+      sk: { field: "gsi3sk", composite: ["createdAt"] }
+    }
+  }
+});
 var evalEntity = new Entity({
   model: {
     entity: "eval",
@@ -970,7 +1091,8 @@ function getElectroDbService(client, tableName) {
       trace: traceEntity,
       workflow_snapshot: workflowSnapshotEntity,
       resource: resourceEntity,
-      score: scoreEntity
+      score: scoreEntity,
+      background_task: backgroundTaskEntity
     },
     {
       client,
@@ -993,38 +1115,13 @@ function resolveDynamoDBConfig(config) {
     ttl: config.ttl
   };
 }
-
-// src/storage/ttl.ts
-function calculateTtl(entityName, ttlConfig, customTtlSeconds) {
-  const entityConfig = ttlConfig?.[entityName];
-  if (!entityConfig?.enabled) {
-    return void 0;
-  }
-  const ttlSeconds = customTtlSeconds ?? entityConfig.defaultTtlSeconds;
-  if (ttlSeconds === void 0 || ttlSeconds <= 0) {
-    return void 0;
-  }
-  return Math.floor(Date.now() / 1e3) + ttlSeconds;
-}
-function getTtlAttributeName(entityName, ttlConfig) {
-  const entityConfig = ttlConfig?.[entityName];
-  return entityConfig?.attributeName ?? "ttl";
-}
-function isTtlEnabled(entityName, ttlConfig) {
-  return ttlConfig?.[entityName]?.enabled === true;
-}
-function getTtlProps(entityName, ttlConfig, customTtlSeconds) {
-  const ttlValue = calculateTtl(entityName, ttlConfig, customTtlSeconds);
-  if (ttlValue === void 0) return {};
-  const attributeName = getTtlAttributeName(entityName, ttlConfig);
-  return { [attributeName]: ttlValue };
-}
 var ENTITY_MAP = {
   [TABLE_THREADS]: "thread",
   [TABLE_MESSAGES]: "message",
   [TABLE_RESOURCES]: "resource",
   [TABLE_WORKFLOW_SNAPSHOT]: "workflow_snapshot",
-  [TABLE_SCORERS]: "score"
+  [TABLE_SCORERS]: "score",
+  [TABLE_BACKGROUND_TASKS]: "background_task"
 };
 function getDeleteKey(entityName, item) {
   const key = { entity: entityName };
@@ -1033,6 +1130,7 @@ function getDeleteKey(entityName, item) {
     case "message":
     case "resource":
     case "score":
+    case "background_task":
       key.id = item.id;
       break;
     case "workflow_snapshot":
@@ -1062,6 +1160,296 @@ async function deleteTableData(service, tableName) {
   }
 }
 
+// src/storage/domains/background-tasks/index.ts
+var ENTITY = "background_task";
+function serializeJson(v) {
+  if (typeof v === "object" && v != null) return JSON.stringify(v);
+  return v ?? void 0;
+}
+function toElectroRecord(task) {
+  return {
+    entity: ENTITY,
+    id: task.id,
+    status: task.status,
+    toolName: task.toolName,
+    toolCallId: task.toolCallId,
+    agentId: task.agentId,
+    runId: task.runId,
+    threadId: task.threadId ?? void 0,
+    resourceId: task.resourceId ?? void 0,
+    args: serializeJson(task.args),
+    result: serializeJson(task.result),
+    error: serializeJson(task.error),
+    retryCount: task.retryCount,
+    maxRetries: task.maxRetries,
+    timeoutMs: task.timeoutMs,
+    createdAt: task.createdAt.toISOString(),
+    startedAtIso: task.startedAt?.toISOString(),
+    completedAtIso: task.completedAt?.toISOString()
+  };
+}
+function fromElectroRecord(data) {
+  const parseJson = (v) => {
+    if (v == null || v === "") return void 0;
+    if (typeof v === "string") {
+      try {
+        if (v.startsWith("{") || v.startsWith("[")) return JSON.parse(v);
+      } catch {
+        return v;
+      }
+    }
+    return v;
+  };
+  const asDate = (v) => v ? new Date(String(v)) : void 0;
+  return {
+    id: String(data.id),
+    status: String(data.status),
+    toolName: String(data.toolName),
+    toolCallId: String(data.toolCallId),
+    args: parseJson(data.args) ?? {},
+    agentId: String(data.agentId),
+    threadId: data.threadId != null && data.threadId !== "" ? String(data.threadId) : void 0,
+    resourceId: data.resourceId != null && data.resourceId !== "" ? String(data.resourceId) : void 0,
+    runId: String(data.runId),
+    result: parseJson(data.result),
+    error: parseJson(data.error),
+    retryCount: Number(data.retryCount ?? 0),
+    maxRetries: Number(data.maxRetries ?? 0),
+    timeoutMs: Number(data.timeoutMs ?? 3e5),
+    createdAt: asDate(data.createdAt) ?? /* @__PURE__ */ new Date(),
+    startedAt: asDate(data.startedAtIso),
+    completedAt: asDate(data.completedAtIso)
+  };
+}
+var BackgroundTasksStorageDynamoDB = class extends BackgroundTasksStorage {
+  service;
+  constructor(config) {
+    super();
+    const resolved = resolveDynamoDBConfig(config);
+    this.service = resolved.service;
+  }
+  async dangerouslyClearAll() {
+    await deleteTableData(this.service, TABLE_BACKGROUND_TASKS);
+  }
+  async createTask(task) {
+    try {
+      await this.service.entities.background_task.create(toElectroRecord(task)).go();
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: createStorageErrorId("DYNAMODB", "BACKGROUND_TASKS_CREATE", "FAILED"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { taskId: task.id }
+        },
+        error
+      );
+    }
+  }
+  async updateTask(taskId, update) {
+    try {
+      const setFields = {};
+      const removeFields = [];
+      if ("status" in update && update.status !== void 0) {
+        setFields.status = update.status;
+      }
+      if ("retryCount" in update && update.retryCount !== void 0) {
+        setFields.retryCount = update.retryCount;
+      }
+      if ("result" in update) {
+        if (update.result === void 0 || update.result === null) {
+          removeFields.push("result");
+        } else {
+          setFields.result = serializeJson(update.result);
+        }
+      }
+      if ("error" in update) {
+        if (update.error === void 0 || update.error === null) {
+          removeFields.push("error");
+        } else {
+          setFields.error = serializeJson(update.error);
+        }
+      }
+      if ("startedAt" in update) {
+        if (update.startedAt === void 0 || update.startedAt === null) {
+          removeFields.push("startedAtIso");
+        } else {
+          setFields.startedAtIso = update.startedAt.toISOString();
+        }
+      }
+      if ("completedAt" in update) {
+        if (update.completedAt === void 0 || update.completedAt === null) {
+          removeFields.push("completedAtIso");
+        } else {
+          setFields.completedAtIso = update.completedAt.toISOString();
+        }
+      }
+      if (Object.keys(setFields).length === 0 && removeFields.length === 0) return;
+      let op = this.service.entities.background_task.patch({ entity: ENTITY, id: taskId });
+      if (Object.keys(setFields).length > 0) op = op.set(setFields);
+      if (removeFields.length > 0) op = op.remove(removeFields);
+      await op.go();
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: createStorageErrorId("DYNAMODB", "BACKGROUND_TASKS_UPDATE", "FAILED"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { taskId }
+        },
+        error
+      );
+    }
+  }
+  async getTask(taskId) {
+    try {
+      const result = await this.service.entities.background_task.get({ entity: ENTITY, id: taskId }).go();
+      if (!result.data) return null;
+      return fromElectroRecord(result.data);
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: createStorageErrorId("DYNAMODB", "BACKGROUND_TASKS_GET", "FAILED"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { taskId }
+        },
+        error
+      );
+    }
+  }
+  async listTasks(filter) {
+    try {
+      let rawResults = [];
+      if (filter.runId) {
+        const res = await this.service.entities.background_task.query.byRun({ entity: ENTITY, runId: filter.runId }).go({ pages: "all" });
+        rawResults = res.data;
+      } else if (filter.agentId) {
+        const res = await this.service.entities.background_task.query.byAgent({ entity: ENTITY, agentId: filter.agentId }).go({ pages: "all" });
+        rawResults = res.data;
+      } else if (filter.status && !Array.isArray(filter.status)) {
+        const res = await this.service.entities.background_task.query.byStatus({ entity: ENTITY, status: filter.status }).go({ pages: "all" });
+        rawResults = res.data;
+      } else {
+        const res = await this.service.entities.background_task.scan.go({ pages: "all" });
+        rawResults = res.data;
+      }
+      let tasks = rawResults.map(fromElectroRecord);
+      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.toolName) tasks = tasks.filter((t) => t.toolName === filter.toolName);
+      if (filter.runId) tasks = tasks.filter((t) => t.runId === filter.runId);
+      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 dir = filter.orderDirection === "desc" ? -1 : 1;
+      tasks.sort((a, b) => ((a[orderBy]?.getTime() ?? 0) - (b[orderBy]?.getTime() ?? 0)) * dir);
+      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 };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: createStorageErrorId("DYNAMODB", "BACKGROUND_TASKS_LIST", "FAILED"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY
+        },
+        error
+      );
+    }
+  }
+  async deleteTask(taskId) {
+    try {
+      await this.service.entities.background_task.delete({ entity: ENTITY, id: taskId }).go();
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: createStorageErrorId("DYNAMODB", "BACKGROUND_TASKS_DELETE", "FAILED"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { taskId }
+        },
+        error
+      );
+    }
+  }
+  async deleteTasks(filter) {
+    const { tasks } = await this.listTasks(filter);
+    if (tasks.length === 0) return;
+    try {
+      const batchSize = 25;
+      for (let i = 0; i < tasks.length; i += batchSize) {
+        const batch = tasks.slice(i, i + batchSize).map((t) => ({ entity: ENTITY, id: t.id }));
+        await this.service.entities.background_task.delete(batch).go();
+      }
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: createStorageErrorId("DYNAMODB", "BACKGROUND_TASKS_DELETE_MANY", "FAILED"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY
+        },
+        error
+      );
+    }
+  }
+  async getRunningCount() {
+    const { total } = await this.listTasks({ status: "running" });
+    return total;
+  }
+  async getRunningCountByAgent(agentId) {
+    const { total } = await this.listTasks({ status: "running", agentId });
+    return total;
+  }
+};
+
+// src/storage/ttl.ts
+function calculateTtl(entityName, ttlConfig, customTtlSeconds) {
+  const entityConfig = ttlConfig?.[entityName];
+  if (!entityConfig?.enabled) {
+    return void 0;
+  }
+  const ttlSeconds = customTtlSeconds ?? entityConfig.defaultTtlSeconds;
+  if (ttlSeconds === void 0 || ttlSeconds <= 0) {
+    return void 0;
+  }
+  return Math.floor(Date.now() / 1e3) + ttlSeconds;
+}
+function getTtlAttributeName(entityName, ttlConfig) {
+  const entityConfig = ttlConfig?.[entityName];
+  return entityConfig?.attributeName ?? "ttl";
+}
+function isTtlEnabled(entityName, ttlConfig) {
+  return ttlConfig?.[entityName]?.enabled === true;
+}
+function getTtlProps(entityName, ttlConfig, customTtlSeconds) {
+  const ttlValue = calculateTtl(entityName, ttlConfig, customTtlSeconds);
+  if (ttlValue === void 0) return {};
+  const attributeName = getTtlAttributeName(entityName, ttlConfig);
+  return { [attributeName]: ttlValue };
+}
+
 // src/storage/domains/memory/index.ts
 var MemoryStorageDynamoDB = class extends MemoryStorage {
   service;
@@ -2614,10 +3002,12 @@ var DynamoDBStore = class extends MastraCompositeStore {
       const workflows = new WorkflowStorageDynamoDB(domainConfig);
       const memory = new MemoryStorageDynamoDB(domainConfig);
       const scores = new ScoresStorageDynamoDB(domainConfig);
+      const backgroundTasks = new BackgroundTasksStorageDynamoDB(domainConfig);
       this.stores = {
         workflows,
         memory,
-        scores
+        scores,
+        backgroundTasks
       };
     } catch (error) {
       throw new MastraError(
@@ -2720,6 +3110,6 @@ var DynamoDBStore = class extends MastraCompositeStore {
   }
 };
 
-export { DynamoDBStore, MemoryStorageDynamoDB, ScoresStorageDynamoDB, WorkflowStorageDynamoDB, calculateTtl, getTtlAttributeName, getTtlProps, isTtlEnabled };
+export { BackgroundTasksStorageDynamoDB, DynamoDBStore, MemoryStorageDynamoDB, ScoresStorageDynamoDB, WorkflowStorageDynamoDB, calculateTtl, getTtlAttributeName, getTtlProps, isTtlEnabled };
 //# sourceMappingURL=index.js.map
 //# sourceMappingURL=index.js.map