@mastra/mssql 1.2.0 → 1.2.1-alpha.1

package/dist/index.js

@@ -1,9 +1,9 @@
 import { MastraError, ErrorCategory, ErrorDomain } from '@mastra/core/error';
-import { MemoryStorage, TABLE_THREADS, TABLE_MESSAGES, TABLE_RESOURCES, TABLE_SCHEMAS, createStorageErrorId, normalizePerPage, calculatePagination, ObservabilityStorage, TABLE_SPANS, SPAN_SCHEMA, listTracesArgsSchema, toTraceSpans, ScoresStorage, TABLE_SCORERS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, MastraCompositeStore, getDefaultValue, transformScoreRow as transformScoreRow$1, TraceStatus } from '@mastra/core/storage';
+import { BackgroundTasksStorage, TABLE_BACKGROUND_TASKS, TABLE_SCHEMAS, MemoryStorage, TABLE_THREADS, TABLE_MESSAGES, TABLE_RESOURCES, createStorageErrorId, normalizePerPage, calculatePagination, ObservabilityStorage, TABLE_SPANS, SPAN_SCHEMA, listTracesArgsSchema, toTraceSpans, ScoresStorage, TABLE_SCORERS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, MastraCompositeStore, getDefaultValue, transformScoreRow as transformScoreRow$1, TraceStatus } from '@mastra/core/storage';
 import sql from 'mssql';
-import { MessageList } from '@mastra/core/agent';
 import { MastraBase } from '@mastra/core/base';
 import { parseSqlIdentifier } from '@mastra/core/utils';
+import { MessageList } from '@mastra/core/agent';
 import { randomUUID } from 'crypto';
 import { saveScorePayloadSchema } from '@mastra/core/evals';
 
@@ -1350,7 +1350,319 @@ function transformFromSqlRow({
   return result;
 }
 
-// src/storage/domains/memory/index.ts
+// src/storage/domains/background-tasks/index.ts
+function serializeJson(v) {
+  if (typeof v === "object" && v != null) return JSON.stringify(v);
+  return v ?? null;
+}
+function rowToTask(row) {
+  const parseJson = (val) => {
+    if (val == null) return void 0;
+    if (typeof val === "string") {
+      try {
+        return JSON.parse(val);
+      } catch {
+        return val;
+      }
+    }
+    return val;
+  };
+  return {
+    id: row.id,
+    status: row.status,
+    toolName: row.tool_name,
+    toolCallId: row.tool_call_id,
+    args: parseJson(row.args) ?? {},
+    agentId: row.agent_id,
+    threadId: row.thread_id ?? void 0,
+    resourceId: row.resource_id ?? void 0,
+    runId: row.run_id ?? "",
+    result: parseJson(row.result),
+    error: parseJson(row.error),
+    retryCount: Number(row.retry_count),
+    maxRetries: Number(row.max_retries),
+    timeoutMs: Number(row.timeout_ms),
+    createdAt: row.createdAt instanceof Date ? row.createdAt : new Date(row.createdAt),
+    startedAt: row.startedAt ? row.startedAt instanceof Date ? row.startedAt : new Date(row.startedAt) : void 0,
+    completedAt: row.completedAt ? row.completedAt instanceof Date ? row.completedAt : new Date(row.completedAt) : void 0
+  };
+}
+var BackgroundTasksMSSQL = class _BackgroundTasksMSSQL extends BackgroundTasksStorage {
+  pool;
+  db;
+  schema;
+  needsConnect;
+  skipDefaultIndexes;
+  indexes;
+  static MANAGED_TABLES = [TABLE_BACKGROUND_TASKS];
+  constructor(config) {
+    super();
+    const { pool, schemaName, skipDefaultIndexes, indexes, needsConnect } = resolveMssqlConfig(config);
+    this.pool = pool;
+    this.schema = schemaName;
+    this.db = new MssqlDB({ pool, schemaName, skipDefaultIndexes });
+    this.needsConnect = needsConnect;
+    this.skipDefaultIndexes = skipDefaultIndexes;
+    this.indexes = indexes?.filter(
+      (idx) => _BackgroundTasksMSSQL.MANAGED_TABLES.includes(idx.table)
+    );
+  }
+  async init() {
+    if (this.needsConnect) {
+      await this.pool.connect();
+      this.needsConnect = false;
+    }
+    await this.db.createTable({
+      tableName: TABLE_BACKGROUND_TASKS,
+      schema: TABLE_SCHEMAS[TABLE_BACKGROUND_TASKS]
+    });
+    await this.createDefaultIndexes();
+    await this.createCustomIndexes();
+  }
+  getDefaultIndexDefinitions() {
+    const schemaPrefix = this.schema ? `${this.schema}_` : "";
+    return [
+      {
+        name: `${schemaPrefix}mastra_bg_tasks_status_created_at_idx`,
+        table: TABLE_BACKGROUND_TASKS,
+        columns: ["status", "createdAt"]
+      },
+      {
+        name: `${schemaPrefix}mastra_bg_tasks_agent_status_idx`,
+        table: TABLE_BACKGROUND_TASKS,
+        columns: ["agent_id", "status"]
+      },
+      {
+        name: `${schemaPrefix}mastra_bg_tasks_thread_idx`,
+        table: TABLE_BACKGROUND_TASKS,
+        columns: ["thread_id", "createdAt"]
+      },
+      {
+        name: `${schemaPrefix}mastra_bg_tasks_tool_call_idx`,
+        table: TABLE_BACKGROUND_TASKS,
+        columns: ["tool_call_id"]
+      }
+    ];
+  }
+  async createDefaultIndexes() {
+    if (this.skipDefaultIndexes) return;
+    for (const indexDef of this.getDefaultIndexDefinitions()) {
+      try {
+        await this.db.createIndex(indexDef);
+      } catch (error) {
+        this.logger?.warn?.(`Failed to create index ${indexDef.name}:`, error);
+      }
+    }
+  }
+  async createCustomIndexes() {
+    if (!this.indexes || this.indexes.length === 0) return;
+    for (const indexDef of this.indexes) {
+      try {
+        await this.db.createIndex(indexDef);
+      } catch (error) {
+        this.logger?.warn?.(`Failed to create custom index ${indexDef.name}:`, error);
+      }
+    }
+  }
+  async dangerouslyClearAll() {
+    await this.db.clearTable({ tableName: TABLE_BACKGROUND_TASKS });
+  }
+  tableName() {
+    return getTableName2({ indexName: TABLE_BACKGROUND_TASKS, schemaName: getSchemaName2(this.schema) });
+  }
+  async createTask(task) {
+    await this.db.insert({
+      tableName: TABLE_BACKGROUND_TASKS,
+      record: {
+        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: serializeJson(task.args),
+        result: serializeJson(task.result),
+        error: serializeJson(task.error),
+        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
+      }
+    });
+  }
+  async updateTask(taskId, update) {
+    const setClauses = [];
+    const params = {};
+    let idx = 1;
+    if ("status" in update) {
+      setClauses.push(`[status] = @p${idx}`);
+      params[`p${idx++}`] = update.status;
+    }
+    if ("result" in update) {
+      setClauses.push(`[result] = @p${idx}`);
+      params[`p${idx++}`] = serializeJson(update.result);
+    }
+    if ("error" in update) {
+      setClauses.push(`[error] = @p${idx}`);
+      params[`p${idx++}`] = serializeJson(update.error);
+    }
+    if ("retryCount" in update) {
+      setClauses.push(`[retry_count] = @p${idx}`);
+      params[`p${idx++}`] = update.retryCount;
+    }
+    if ("startedAt" in update) {
+      setClauses.push(`[startedAt] = @p${idx}`);
+      params[`p${idx++}`] = update.startedAt?.toISOString() ?? null;
+    }
+    if ("completedAt" in update) {
+      setClauses.push(`[completedAt] = @p${idx}`);
+      params[`p${idx++}`] = update.completedAt?.toISOString() ?? null;
+    }
+    if (setClauses.length === 0) return;
+    setClauses.push(`[id] = [id]`);
+    params[`p${idx}`] = taskId;
+    const request = this.pool.request();
+    for (const [name, value] of Object.entries(params)) {
+      request.input(name, value);
+    }
+    await request.query(`UPDATE ${this.tableName()} SET ${setClauses.join(", ")} WHERE [id] = @p${idx}`);
+  }
+  async getTask(taskId) {
+    const request = this.pool.request();
+    request.input("p1", taskId);
+    const result = await request.query(`SELECT * FROM ${this.tableName()} WHERE [id] = @p1`);
+    if (result.recordset.length === 0) return null;
+    return rowToTask(result.recordset[0]);
+  }
+  async listTasks(filter) {
+    const conditions = [];
+    const params = {};
+    let idx = 1;
+    if (filter.status) {
+      const statuses = Array.isArray(filter.status) ? filter.status : [filter.status];
+      const placeholders = statuses.map(() => {
+        const name = `p${idx++}`;
+        return `@${name}`;
+      });
+      statuses.forEach((s, i) => {
+        params[`p${idx - statuses.length + i}`] = s;
+      });
+      conditions.push(`[status] IN (${placeholders.join(", ")})`);
+    }
+    if (filter.agentId) {
+      params[`p${idx}`] = filter.agentId;
+      conditions.push(`[agent_id] = @p${idx++}`);
+    }
+    if (filter.threadId) {
+      params[`p${idx}`] = filter.threadId;
+      conditions.push(`[thread_id] = @p${idx++}`);
+    }
+    if (filter.resourceId) {
+      params[`p${idx}`] = filter.resourceId;
+      conditions.push(`[resource_id] = @p${idx++}`);
+    }
+    if (filter.runId) {
+      params[`p${idx}`] = filter.runId;
+      conditions.push(`[run_id] = @p${idx++}`);
+    }
+    if (filter.toolName) {
+      params[`p${idx}`] = filter.toolName;
+      conditions.push(`[tool_name] = @p${idx++}`);
+    }
+    const dateCol = filter.dateFilterBy === "startedAt" ? "[startedAt]" : filter.dateFilterBy === "completedAt" ? "[completedAt]" : "[createdAt]";
+    if (filter.fromDate) {
+      params[`p${idx}`] = filter.fromDate.toISOString();
+      conditions.push(`${dateCol} >= @p${idx++}`);
+    }
+    if (filter.toDate) {
+      params[`p${idx}`] = filter.toDate.toISOString();
+      conditions.push(`${dateCol} < @p${idx++}`);
+    }
+    const where = conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : "";
+    const countRequest = this.pool.request();
+    for (const [name, value] of Object.entries(params)) {
+      countRequest.input(name, value);
+    }
+    const countResult = await countRequest.query(`SELECT COUNT(*) as count FROM ${this.tableName()} ${where}`);
+    const total = Number(countResult.recordset[0]?.count ?? 0);
+    const orderCol = filter.orderBy === "startedAt" ? "[startedAt]" : filter.orderBy === "completedAt" ? "[completedAt]" : "[createdAt]";
+    const direction = filter.orderDirection === "desc" ? "DESC" : "ASC";
+    let sql5 = `SELECT * FROM ${this.tableName()} ${where} ORDER BY ${orderCol} ${direction}`;
+    if (filter.perPage != null) {
+      const offset = filter.page != null ? filter.page * filter.perPage : 0;
+      params[`p${idx}`] = offset;
+      params[`p${idx + 1}`] = filter.perPage;
+      sql5 += ` OFFSET @p${idx} ROWS FETCH NEXT @p${idx + 1} ROWS ONLY`;
+      idx += 2;
+    }
+    const request = this.pool.request();
+    for (const [name, value] of Object.entries(params)) {
+      request.input(name, value);
+    }
+    const result = await request.query(sql5);
+    return { tasks: result.recordset.map(rowToTask), total };
+  }
+  async deleteTask(taskId) {
+    const request = this.pool.request();
+    request.input("p1", taskId);
+    await request.query(`DELETE FROM ${this.tableName()} WHERE [id] = @p1`);
+  }
+  async deleteTasks(filter) {
+    const conditions = [];
+    const params = {};
+    let idx = 1;
+    if (filter.status) {
+      const statuses = Array.isArray(filter.status) ? filter.status : [filter.status];
+      const placeholders = statuses.map(() => {
+        const name = `p${idx++}`;
+        return `@${name}`;
+      });
+      statuses.forEach((s, i) => {
+        params[`p${idx - statuses.length + i}`] = s;
+      });
+      conditions.push(`[status] IN (${placeholders.join(", ")})`);
+    }
+    const dateCol = filter.dateFilterBy === "startedAt" ? "[startedAt]" : filter.dateFilterBy === "completedAt" ? "[completedAt]" : "[createdAt]";
+    if (filter.fromDate) {
+      params[`p${idx}`] = filter.fromDate.toISOString();
+      conditions.push(`${dateCol} >= @p${idx++}`);
+    }
+    if (filter.toDate) {
+      params[`p${idx}`] = filter.toDate.toISOString();
+      conditions.push(`${dateCol} < @p${idx++}`);
+    }
+    if (filter.agentId) {
+      params[`p${idx}`] = filter.agentId;
+      conditions.push(`[agent_id] = @p${idx++}`);
+    }
+    if (filter.runId) {
+      params[`p${idx}`] = filter.runId;
+      conditions.push(`[run_id] = @p${idx++}`);
+    }
+    if (conditions.length === 0) return;
+    const request = this.pool.request();
+    for (const [name, value] of Object.entries(params)) {
+      request.input(name, value);
+    }
+    await request.query(`DELETE FROM ${this.tableName()} WHERE ${conditions.join(" AND ")}`);
+  }
+  async getRunningCount() {
+    const result = await this.pool.request().query(`SELECT COUNT(*) as count FROM ${this.tableName()} WHERE [status] = 'running'`);
+    return Number(result.recordset[0]?.count ?? 0);
+  }
+  async getRunningCountByAgent(agentId) {
+    const request = this.pool.request();
+    request.input("p1", agentId);
+    const result = await request.query(
+      `SELECT COUNT(*) as count FROM ${this.tableName()} WHERE [status] = 'running' AND [agent_id] = @p1`
+    );
+    return Number(result.recordset[0]?.count ?? 0);
+  }
+};
 var MemoryMSSQL = class _MemoryMSSQL extends MemoryStorage {
   pool;
   schema;
@@ -2624,6 +2936,51 @@ var ObservabilityMSSQL = class _ObservabilityMSSQL extends ObservabilityStorage
       );
     }
   }
+  async getTraceLight(args) {
+    const { traceId } = args;
+    try {
+      const tableName = getTableName2({
+        indexName: TABLE_SPANS,
+        schemaName: getSchemaName2(this.schema)
+      });
+      const request = this.pool.request();
+      request.input("traceId", traceId);
+      const result = await request.query(
+        `SELECT
+          [traceId], [spanId], [parentSpanId], [name],
+          [entityType], [entityId], [entityName],
+          [spanType], [error], [isEvent],
+          [startedAt], [endedAt], [createdAt], [updatedAt]
+        FROM ${tableName}
+        WHERE [traceId] = @traceId
+        ORDER BY [startedAt] ASC`
+      );
+      if (!result.recordset || result.recordset.length === 0) {
+        return null;
+      }
+      return {
+        traceId,
+        spans: result.recordset.map(
+          (span) => transformFromSqlRow({
+            tableName: TABLE_SPANS,
+            sqlRow: span
+          })
+        )
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: createStorageErrorId("MSSQL", "GET_TRACE_LIGHT", "FAILED"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER,
+          details: {
+            traceId
+          }
+        },
+        error
+      );
+    }
+  }
   async getSpan(args) {
     const { traceId, spanId } = args;
     try {
@@ -3988,11 +4345,13 @@ var MSSQLStore = class extends MastraCompositeStore {
       const workflows = new WorkflowsMSSQL(domainConfig);
       const memory = new MemoryMSSQL(domainConfig);
       const observability = new ObservabilityMSSQL(domainConfig);
+      const backgroundTasks = new BackgroundTasksMSSQL(domainConfig);
       this.stores = {
         scores,
         workflows,
         memory,
-        observability
+        observability,
+        backgroundTasks
       };
     } catch (e) {
       throw new MastraError(
@@ -4045,6 +4404,6 @@ var MSSQLStore = class extends MastraCompositeStore {
   }
 };
 
-export { MSSQLStore, MemoryMSSQL, ObservabilityMSSQL, ScoresMSSQL, WorkflowsMSSQL };
+export { BackgroundTasksMSSQL, MSSQLStore, MemoryMSSQL, ObservabilityMSSQL, ScoresMSSQL, WorkflowsMSSQL };
 //# sourceMappingURL=index.js.map
 //# sourceMappingURL=index.js.map