@mastra/upstash 0.12.1 → 0.12.2

package/src/storage/domains/operations/index.ts

@@ -0,0 +1,168 @@
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import { StoreOperations } from '@mastra/core/storage';
+import type { TABLE_NAMES, StorageColumn } from '@mastra/core/storage';
+import type { Redis } from '@upstash/redis';
+import { getKey, processRecord } from '../utils';
+
+export class StoreOperationsUpstash extends StoreOperations {
+  private client: Redis;
+
+  constructor({ client }: { client: Redis }) {
+    super();
+    this.client = client;
+  }
+
+  async createTable({
+    tableName: _tableName,
+    schema: _schema,
+  }: {
+    tableName: TABLE_NAMES;
+    schema: Record<string, StorageColumn>;
+  }): Promise<void> {
+    // For Redis/Upstash, tables are created implicitly when data is inserted
+    // This method is a no-op for Redis-based storage
+  }
+
+  async alterTable({
+    tableName: _tableName,
+    schema: _schema,
+    ifNotExists: _ifNotExists,
+  }: {
+    tableName: TABLE_NAMES;
+    schema: Record<string, StorageColumn>;
+    ifNotExists: string[];
+  }): Promise<void> {
+    // For Redis/Upstash, schema changes are handled implicitly
+    // This method is a no-op for Redis-based storage
+  }
+
+  async clearTable({ tableName }: { tableName: TABLE_NAMES }): Promise<void> {
+    const pattern = `${tableName}:*`;
+    try {
+      await this.scanAndDelete(pattern);
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_UPSTASH_STORAGE_CLEAR_TABLE_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            tableName,
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  async dropTable({ tableName }: { tableName: TABLE_NAMES }): Promise<void> {
+    return this.clearTable({ tableName });
+  }
+
+  async insert({ tableName, record }: { tableName: TABLE_NAMES; record: Record<string, any> }): Promise<void> {
+    const { key, processedRecord } = processRecord(tableName, record);
+
+    try {
+      await this.client.set(key, processedRecord);
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_UPSTASH_STORAGE_INSERT_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            tableName,
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  async batchInsert(input: { tableName: TABLE_NAMES; records: Record<string, any>[] }): Promise<void> {
+    const { tableName, records } = input;
+    if (!records.length) return;
+
+    const batchSize = 1000;
+    try {
+      for (let i = 0; i < records.length; i += batchSize) {
+        const batch = records.slice(i, i + batchSize);
+        const pipeline = this.client.pipeline();
+        for (const record of batch) {
+          const { key, processedRecord } = processRecord(tableName, record);
+          pipeline.set(key, processedRecord);
+        }
+        await pipeline.exec();
+      }
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_UPSTASH_STORAGE_BATCH_INSERT_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            tableName,
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  async load<R>({ tableName, keys }: { tableName: TABLE_NAMES; keys: Record<string, string> }): Promise<R | null> {
+    const key = getKey(tableName, keys);
+    try {
+      const data = await this.client.get<R>(key);
+      return data || null;
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_UPSTASH_STORAGE_LOAD_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            tableName,
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  async hasColumn(_tableName: TABLE_NAMES, _column: string): Promise<boolean> {
+    // For Redis/Upstash, columns are dynamic and always available
+    // This method always returns true for Redis-based storage
+    return true;
+  }
+
+  async scanKeys(pattern: string, batchSize = 10000): Promise<string[]> {
+    let cursor = '0';
+    let keys: string[] = [];
+    do {
+      const [nextCursor, batch] = await this.client.scan(cursor, {
+        match: pattern,
+        count: batchSize,
+      });
+      keys.push(...batch);
+      cursor = nextCursor;
+    } while (cursor !== '0');
+    return keys;
+  }
+
+  async scanAndDelete(pattern: string, batchSize = 10000): Promise<number> {
+    let cursor = '0';
+    let totalDeleted = 0;
+    do {
+      const [nextCursor, keys] = await this.client.scan(cursor, {
+        match: pattern,
+        count: batchSize,
+      });
+      if (keys.length > 0) {
+        await this.client.del(...keys);
+        totalDeleted += keys.length;
+      }
+      cursor = nextCursor;
+    } while (cursor !== '0');
+    return totalDeleted;
+  }
+}

package/src/storage/domains/scores/index.ts

@@ -0,0 +1,216 @@
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import type { ScoreRowData } from '@mastra/core/scores';
+import { ScoresStorage, TABLE_SCORERS } from '@mastra/core/storage';
+import type { Redis } from '@upstash/redis';
+import type { StoreOperationsUpstash } from '../operations';
+import { processRecord } from '../utils';
+
+function transformScoreRow(row: Record<string, any>): ScoreRowData {
+  const parseField = (v: any) => {
+    if (typeof v === 'string') {
+      try {
+        return JSON.parse(v);
+      } catch {
+        return v;
+      }
+    }
+    return v;
+  };
+  return {
+    ...row,
+    scorer: parseField(row.scorer),
+    extractStepResult: parseField(row.extractStepResult),
+    analyzeStepResult: parseField(row.analyzeStepResult),
+    metadata: parseField(row.metadata),
+    input: parseField(row.input),
+    output: parseField(row.output),
+    additionalContext: parseField(row.additionalContext),
+    runtimeContext: parseField(row.runtimeContext),
+    entity: parseField(row.entity),
+    createdAt: row.createdAt,
+    updatedAt: row.updatedAt,
+  } as ScoreRowData;
+}
+
+export class ScoresUpstash extends ScoresStorage {
+  private client: Redis;
+  private operations: StoreOperationsUpstash;
+
+  constructor({ client, operations }: { client: Redis; operations: StoreOperationsUpstash }) {
+    super();
+    this.client = client;
+    this.operations = operations;
+  }
+
+  async getScoreById({ id }: { id: string }): Promise<ScoreRowData | null> {
+    try {
+      const data = await this.operations.load<ScoreRowData>({
+        tableName: TABLE_SCORERS,
+        keys: { id },
+      });
+      if (!data) return null;
+      return transformScoreRow(data);
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_UPSTASH_STORAGE_GET_SCORE_BY_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { id },
+        },
+        error,
+      );
+    }
+  }
+
+  async getScoresByScorerId({
+    scorerId,
+    pagination = { page: 0, perPage: 20 },
+  }: {
+    scorerId: string;
+    pagination?: { page: number; perPage: number };
+  }): Promise<{
+    scores: ScoreRowData[];
+    pagination: { total: number; page: number; perPage: number; hasMore: boolean };
+  }> {
+    const pattern = `${TABLE_SCORERS}:*`;
+    const keys = await this.operations.scanKeys(pattern);
+    if (keys.length === 0) {
+      return {
+        scores: [],
+        pagination: { total: 0, page: pagination.page, perPage: pagination.perPage, hasMore: false },
+      };
+    }
+    const pipeline = this.client.pipeline();
+    keys.forEach(key => pipeline.get(key));
+    const results = await pipeline.exec();
+    // Filter out nulls and by scorerId
+    const filtered = results
+      .map((row: any) => row as Record<string, any> | null)
+      .filter((row): row is Record<string, any> => !!row && typeof row === 'object' && row.scorerId === scorerId);
+    const total = filtered.length;
+    const { page, perPage } = pagination;
+    const start = page * perPage;
+    const end = start + perPage;
+    const paged = filtered.slice(start, end);
+    const scores = paged.map(row => transformScoreRow(row));
+    return {
+      scores,
+      pagination: {
+        total,
+        page,
+        perPage,
+        hasMore: end < total,
+      },
+    };
+  }
+
+  async saveScore(score: ScoreRowData): Promise<{ score: ScoreRowData }> {
+    const { key, processedRecord } = processRecord(TABLE_SCORERS, score);
+    try {
+      await this.client.set(key, processedRecord);
+      return { score };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_UPSTASH_STORAGE_SAVE_SCORE_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { id: score.id },
+        },
+        error,
+      );
+    }
+  }
+
+  async getScoresByRunId({
+    runId,
+    pagination = { page: 0, perPage: 20 },
+  }: {
+    runId: string;
+    pagination?: { page: number; perPage: number };
+  }): Promise<{
+    scores: ScoreRowData[];
+    pagination: { total: number; page: number; perPage: number; hasMore: boolean };
+  }> {
+    const pattern = `${TABLE_SCORERS}:*`;
+    const keys = await this.operations.scanKeys(pattern);
+    if (keys.length === 0) {
+      return {
+        scores: [],
+        pagination: { total: 0, page: pagination.page, perPage: pagination.perPage, hasMore: false },
+      };
+    }
+    const pipeline = this.client.pipeline();
+    keys.forEach(key => pipeline.get(key));
+    const results = await pipeline.exec();
+    // Filter out nulls and by runId
+    const filtered = results
+      .map((row: any) => row as Record<string, any> | null)
+      .filter((row): row is Record<string, any> => !!row && typeof row === 'object' && row.runId === runId);
+    const total = filtered.length;
+    const { page, perPage } = pagination;
+    const start = page * perPage;
+    const end = start + perPage;
+    const paged = filtered.slice(start, end);
+    const scores = paged.map(row => transformScoreRow(row));
+    return {
+      scores,
+      pagination: {
+        total,
+        page,
+        perPage,
+        hasMore: end < total,
+      },
+    };
+  }
+
+  async getScoresByEntityId({
+    entityId,
+    entityType,
+    pagination = { page: 0, perPage: 20 },
+  }: {
+    entityId: string;
+    entityType?: string;
+    pagination?: { page: number; perPage: number };
+  }): Promise<{
+    scores: ScoreRowData[];
+    pagination: { total: number; page: number; perPage: number; hasMore: boolean };
+  }> {
+    const pattern = `${TABLE_SCORERS}:*`;
+    const keys = await this.operations.scanKeys(pattern);
+    if (keys.length === 0) {
+      return {
+        scores: [],
+        pagination: { total: 0, page: pagination.page, perPage: pagination.perPage, hasMore: false },
+      };
+    }
+    const pipeline = this.client.pipeline();
+    keys.forEach(key => pipeline.get(key));
+    const results = await pipeline.exec();
+
+    const filtered = results
+      .map((row: any) => row as Record<string, any> | null)
+      .filter((row): row is Record<string, any> => {
+        if (!row || typeof row !== 'object') return false;
+        if (row.entityId !== entityId) return false;
+        if (entityType && row.entityType !== entityType) return false;
+        return true;
+      });
+    const total = filtered.length;
+    const { page, perPage } = pagination;
+    const start = page * perPage;
+    const end = start + perPage;
+    const paged = filtered.slice(start, end);
+    const scores = paged.map(row => transformScoreRow(row));
+    return {
+      scores,
+      pagination: {
+        total,
+        page,
+        perPage,
+        hasMore: end < total,
+      },
+    };
+  }
+}

package/src/storage/domains/traces/index.ts

@@ -0,0 +1,172 @@
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import type { StorageGetTracesArg, PaginationInfo, PaginationArgs } from '@mastra/core/storage';
+import { TracesStorage, TABLE_TRACES } from '@mastra/core/storage';
+import type { Redis } from '@upstash/redis';
+import type { StoreOperationsUpstash } from '../operations';
+import { ensureDate, parseJSON } from '../utils';
+
+export class TracesUpstash extends TracesStorage {
+  private client: Redis;
+  private operations: StoreOperationsUpstash;
+
+  constructor({ client, operations }: { client: Redis; operations: StoreOperationsUpstash }) {
+    super();
+    this.client = client;
+    this.operations = operations;
+  }
+
+  /**
+   * @deprecated use getTracesPaginated instead
+   */
+  public async getTraces(args: StorageGetTracesArg): Promise<any[]> {
+    if (args.fromDate || args.toDate) {
+      (args as any).dateRange = {
+        start: args.fromDate,
+        end: args.toDate,
+      };
+    }
+    try {
+      const { traces } = await this.getTracesPaginated(args);
+      return traces;
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_UPSTASH_STORAGE_GET_TRACES_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+        },
+        error,
+      );
+    }
+  }
+
+  public async getTracesPaginated(
+    args: {
+      name?: string;
+      scope?: string;
+      attributes?: Record<string, string>;
+      filters?: Record<string, any>;
+    } & PaginationArgs,
+  ): Promise<PaginationInfo & { traces: any[] }> {
+    const { name, scope, page = 0, perPage = 100, attributes, filters, dateRange } = args;
+    const fromDate = dateRange?.start;
+    const toDate = dateRange?.end;
+
+    try {
+      const pattern = `${TABLE_TRACES}:*`;
+      const keys = await this.operations.scanKeys(pattern);
+
+      if (keys.length === 0) {
+        return {
+          traces: [],
+          total: 0,
+          page,
+          perPage: perPage || 100,
+          hasMore: false,
+        };
+      }
+
+      const pipeline = this.client.pipeline();
+      keys.forEach(key => pipeline.get(key));
+      const results = await pipeline.exec();
+
+      let filteredTraces = results.filter(
+        (record): record is Record<string, any> => record !== null && typeof record === 'object',
+      );
+
+      if (name) {
+        filteredTraces = filteredTraces.filter(record => record.name?.toLowerCase().startsWith(name.toLowerCase()));
+      }
+      if (scope) {
+        filteredTraces = filteredTraces.filter(record => record.scope === scope);
+      }
+      if (attributes) {
+        filteredTraces = filteredTraces.filter(record => {
+          const recordAttributes = record.attributes;
+          if (!recordAttributes) return false;
+          const parsedAttributes =
+            typeof recordAttributes === 'string' ? JSON.parse(recordAttributes) : recordAttributes;
+          return Object.entries(attributes).every(([key, value]) => parsedAttributes[key] === value);
+        });
+      }
+      if (filters) {
+        filteredTraces = filteredTraces.filter(record =>
+          Object.entries(filters).every(([key, value]) => record[key] === value),
+        );
+      }
+      if (fromDate) {
+        filteredTraces = filteredTraces.filter(
+          record => new Date(record.createdAt).getTime() >= new Date(fromDate).getTime(),
+        );
+      }
+      if (toDate) {
+        filteredTraces = filteredTraces.filter(
+          record => new Date(record.createdAt).getTime() <= new Date(toDate).getTime(),
+        );
+      }
+
+      filteredTraces.sort((a, b) => new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime());
+
+      const transformedTraces = filteredTraces.map(record => ({
+        id: record.id,
+        parentSpanId: record.parentSpanId,
+        traceId: record.traceId,
+        name: record.name,
+        scope: record.scope,
+        kind: record.kind,
+        status: parseJSON(record.status),
+        events: parseJSON(record.events),
+        links: parseJSON(record.links),
+        attributes: parseJSON(record.attributes),
+        startTime: record.startTime,
+        endTime: record.endTime,
+        other: parseJSON(record.other),
+        createdAt: ensureDate(record.createdAt),
+      }));
+
+      const total = transformedTraces.length;
+      const resolvedPerPage = perPage || 100;
+      const start = page * resolvedPerPage;
+      const end = start + resolvedPerPage;
+      const paginatedTraces = transformedTraces.slice(start, end);
+      const hasMore = end < total;
+
+      return {
+        traces: paginatedTraces,
+        total,
+        page,
+        perPage: resolvedPerPage,
+        hasMore,
+      };
+    } catch (error) {
+      const mastraError = new MastraError(
+        {
+          id: 'STORAGE_UPSTASH_STORAGE_GET_TRACES_PAGINATED_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            name: args.name || '',
+            scope: args.scope || '',
+          },
+        },
+        error,
+      );
+      this.logger?.trackException(mastraError);
+      this.logger.error(mastraError.toString());
+      return {
+        traces: [],
+        total: 0,
+        page,
+        perPage: perPage || 100,
+        hasMore: false,
+      };
+    }
+  }
+
+  async batchTraceInsert(args: { records: Record<string, any>[] }): Promise<void> {
+    return this.operations.batchInsert({
+      tableName: TABLE_TRACES,
+      records: args.records,
+    });
+  }
+}

package/src/storage/domains/utils.ts

@@ -0,0 +1,57 @@
+import { serializeDate, TABLE_MESSAGES, TABLE_WORKFLOW_SNAPSHOT, TABLE_EVALS } from '@mastra/core/storage';
+import type { TABLE_NAMES } from '@mastra/core/storage';
+
+export function ensureDate(value: any): Date | null {
+  if (!value) return null;
+  if (value instanceof Date) return value;
+  if (typeof value === 'string') return new Date(value);
+  if (typeof value === 'number') return new Date(value);
+  return null;
+}
+
+export function parseJSON(value: any): any {
+  if (typeof value === 'string') {
+    try {
+      return JSON.parse(value);
+    } catch {
+      return value;
+    }
+  }
+  return value;
+}
+
+export function getKey(tableName: TABLE_NAMES, keys: Record<string, any>): string {
+  const keyParts = Object.entries(keys)
+    .filter(([_, value]) => value !== undefined)
+    .map(([key, value]) => `${key}:${value}`);
+  return `${tableName}:${keyParts.join(':')}`;
+}
+
+export function processRecord(tableName: TABLE_NAMES, record: Record<string, any>) {
+  let key: string;
+
+  if (tableName === TABLE_MESSAGES) {
+    // For messages, use threadId as the primary key component
+    key = getKey(tableName, { threadId: record.threadId, id: record.id });
+  } else if (tableName === TABLE_WORKFLOW_SNAPSHOT) {
+    key = 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 = getKey(tableName, { id: record.run_id });
+  } else {
+    key = getKey(tableName, { id: record.id });
+  }
+
+  // Convert dates to ISO strings before storing
+  const processedRecord = {
+    ...record,
+    createdAt: serializeDate(record.createdAt),
+    updatedAt: serializeDate(record.updatedAt),
+  };
+
+  return { key, processedRecord };
+}