@mastra/libsql 0.13.4 → 0.13.5

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

@@ -1,9 +1,14 @@
-import type { Client } from '@libsql/client';
+import type { Client, InValue } from '@libsql/client';
 import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
-import { TABLE_WORKFLOW_SNAPSHOT, StoreOperations } from '@mastra/core/storage';
+import { TABLE_WORKFLOW_SNAPSHOT, StoreOperations, TABLE_AI_SPANS } from '@mastra/core/storage';
 import type { StorageColumn, TABLE_NAMES } from '@mastra/core/storage';
 import { parseSqlIdentifier } from '@mastra/core/utils';
-import { createExecuteWriteOperationWithRetry, prepareStatement } from '../utils';
+import {
+  createExecuteWriteOperationWithRetry,
+  prepareDeleteStatement,
+  prepareStatement,
+  prepareUpdateStatement,
+} from '../utils';
 
 export class StoreOperationsLibSQL extends StoreOperations {
   private client: Client;
@@ -66,6 +71,14 @@ export class StoreOperationsLibSQL extends StoreOperations {
       return stmnt;
     }
 
+    if (tableName === TABLE_AI_SPANS) {
+      const stmnt = `CREATE TABLE IF NOT EXISTS ${parsedTableName} (
+                    ${columns.join(',\n')},
+                    PRIMARY KEY (traceId, spanId)
+                )`;
+      return stmnt;
+    }
+
     return `CREATE TABLE IF NOT EXISTS ${parsedTableName} (${columns.join(', ')})`;
   }
 
@@ -164,6 +177,93 @@ export class StoreOperationsLibSQL extends StoreOperations {
     return parsed as R;
   }
 
+  async loadMany<R>({
+    tableName,
+    whereClause,
+    orderBy,
+    offset,
+    limit,
+    args,
+  }: {
+    tableName: TABLE_NAMES;
+    whereClause?: { sql: string; args: InValue[] };
+    orderBy?: string;
+    offset?: number;
+    limit?: number;
+    args?: any[];
+  }): Promise<R[]> {
+    const parsedTableName = parseSqlIdentifier(tableName, 'table name');
+
+    let statement = `SELECT * FROM ${parsedTableName}`;
+
+    if (whereClause?.sql) {
+      statement += `${whereClause.sql}`;
+    }
+
+    if (orderBy) {
+      statement += ` ORDER BY ${orderBy}`;
+    }
+
+    if (limit) {
+      statement += ` LIMIT ${limit}`;
+    }
+
+    if (offset) {
+      statement += ` OFFSET ${offset}`;
+    }
+
+    const result = await this.client.execute({
+      sql: statement,
+      args: [...(whereClause?.args ?? []), ...(args ?? [])],
+    });
+
+    return result.rows as R[];
+  }
+
+  async loadTotalCount({
+    tableName,
+    whereClause,
+  }: {
+    tableName: TABLE_NAMES;
+    whereClause?: { sql: string; args: InValue[] };
+  }): Promise<number> {
+    const parsedTableName = parseSqlIdentifier(tableName, 'table name');
+
+    const statement = `SELECT COUNT(*) as count FROM ${parsedTableName} ${whereClause ? `${whereClause.sql}` : ''}`;
+
+    const result = await this.client.execute({
+      sql: statement,
+      args: whereClause?.args ?? [],
+    });
+
+    if (!result.rows || result.rows.length === 0) {
+      return 0;
+    }
+
+    return (result.rows[0]?.count as number) ?? 0;
+  }
+
+  public update(args: { tableName: TABLE_NAMES; keys: Record<string, any>; data: Record<string, any> }): Promise<void> {
+    const executeWriteOperationWithRetry = createExecuteWriteOperationWithRetry({
+      logger: this.logger,
+      maxRetries: this.maxRetries,
+      initialBackoffMs: this.initialBackoffMs,
+    });
+    return executeWriteOperationWithRetry(() => this.executeUpdate(args), `update table ${args.tableName}`);
+  }
+
+  private async executeUpdate({
+    tableName,
+    keys,
+    data,
+  }: {
+    tableName: TABLE_NAMES;
+    keys: Record<string, any>;
+    data: Record<string, any>;
+  }): Promise<void> {
+    await this.client.execute(prepareUpdateStatement({ tableName, updates: data, keys }));
+  }
+
   private async doBatchInsert({
     tableName,
     records,
@@ -201,6 +301,116 @@ export class StoreOperationsLibSQL extends StoreOperations {
     });
   }
 
+  /**
+   * Public batch update method with retry logic
+   */
+  public batchUpdate(args: {
+    tableName: TABLE_NAMES;
+    updates: Array<{
+      keys: Record<string, any>;
+      data: Record<string, any>;
+    }>;
+  }): Promise<void> {
+    const executeWriteOperationWithRetry = createExecuteWriteOperationWithRetry({
+      logger: this.logger,
+      maxRetries: this.maxRetries,
+      initialBackoffMs: this.initialBackoffMs,
+    });
+
+    return executeWriteOperationWithRetry(
+      () => this.executeBatchUpdate(args),
+      `batch update in table ${args.tableName}`,
+    ).catch(error => {
+      throw new MastraError(
+        {
+          id: 'LIBSQL_STORE_BATCH_UPDATE_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            tableName: args.tableName,
+          },
+        },
+        error,
+      );
+    });
+  }
+
+  /**
+   * Updates multiple records in batch. Each record can be updated based on single or composite keys.
+   */
+  private async executeBatchUpdate({
+    tableName,
+    updates,
+  }: {
+    tableName: TABLE_NAMES;
+    updates: Array<{
+      keys: Record<string, any>;
+      data: Record<string, any>;
+    }>;
+  }): Promise<void> {
+    if (updates.length === 0) return;
+
+    const batchStatements = updates.map(({ keys, data }) =>
+      prepareUpdateStatement({
+        tableName,
+        updates: data,
+        keys,
+      }),
+    );
+
+    await this.client.batch(batchStatements, 'write');
+  }
+
+  /**
+   * Public batch delete method with retry logic
+   */
+  public batchDelete({ tableName, keys }: { tableName: TABLE_NAMES; keys: Array<Record<string, any>> }): Promise<void> {
+    const executeWriteOperationWithRetry = createExecuteWriteOperationWithRetry({
+      logger: this.logger,
+      maxRetries: this.maxRetries,
+      initialBackoffMs: this.initialBackoffMs,
+    });
+
+    return executeWriteOperationWithRetry(
+      () => this.executeBatchDelete({ tableName, keys }),
+      `batch delete from table ${tableName}`,
+    ).catch(error => {
+      throw new MastraError(
+        {
+          id: 'LIBSQL_STORE_BATCH_DELETE_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            tableName,
+          },
+        },
+        error,
+      );
+    });
+  }
+
+  /**
+   * Deletes multiple records in batch. Each record can be deleted based on single or composite keys.
+   */
+  private async executeBatchDelete({
+    tableName,
+    keys,
+  }: {
+    tableName: TABLE_NAMES;
+    keys: Array<Record<string, any>>;
+  }): Promise<void> {
+    if (keys.length === 0) return;
+
+    const batchStatements = keys.map(keyObj =>
+      prepareDeleteStatement({
+        tableName,
+        keys: keyObj,
+      }),
+    );
+
+    await this.client.batch(batchStatements, 'write');
+  }
+
   /**
    * Alters table schema to add columns if they don't exist
    * @param tableName Name of the table

package/src/storage/domains/utils.ts

@@ -1,6 +1,7 @@
 import type { InValue } from '@libsql/client';
 import type { IMastraLogger } from '@mastra/core/logger';
-import type { TABLE_NAMES } from '@mastra/core/storage';
+import { safelyParseJSON, TABLE_SCHEMAS } from '@mastra/core/storage';
+import type { PaginationArgs, StorageColumn, TABLE_NAMES } from '@mastra/core/storage';
 import { parseSqlIdentifier } from '@mastra/core/utils';
 
 export function createExecuteWriteOperationWithRetry({
@@ -49,7 +50,7 @@ export function prepareStatement({ tableName, record }: { tableName: TABLE_NAMES
   const parsedTableName = parseSqlIdentifier(tableName, 'table name');
   const columns = Object.keys(record).map(col => parseSqlIdentifier(col, 'column name'));
   const values = Object.values(record).map(v => {
-    if (typeof v === `undefined`) {
+    if (typeof v === `undefined` || v === null) {
       // returning an undefined value will cause libsql to throw
       return null;
     }
@@ -65,3 +66,207 @@ export function prepareStatement({ tableName, record }: { tableName: TABLE_NAMES
     args: values,
   };
 }
+
+export function prepareUpdateStatement({
+  tableName,
+  updates,
+  keys,
+}: {
+  tableName: TABLE_NAMES;
+  updates: Record<string, any>;
+  keys: Record<string, any>;
+}): {
+  sql: string;
+  args: InValue[];
+} {
+  const parsedTableName = parseSqlIdentifier(tableName, 'table name');
+  const schema = TABLE_SCHEMAS[tableName];
+
+  // Prepare SET clause
+  const updateColumns = Object.keys(updates).map(col => parseSqlIdentifier(col, 'column name'));
+  const updateValues = Object.values(updates).map(transformToSqlValue);
+  const setClause = updateColumns.map(col => `${col} = ?`).join(', ');
+
+  const whereClause = prepareWhereClause(keys, schema);
+
+  return {
+    sql: `UPDATE ${parsedTableName} SET ${setClause}${whereClause.sql}`,
+    args: [...updateValues, ...whereClause.args],
+  };
+}
+
+export function transformToSqlValue(value: any): InValue {
+  if (typeof value === 'undefined' || value === null) {
+    return null;
+  }
+  if (value instanceof Date) {
+    return value.toISOString();
+  }
+  return typeof value === 'object' ? JSON.stringify(value) : value;
+}
+
+export function prepareDeleteStatement({ tableName, keys }: { tableName: TABLE_NAMES; keys: Record<string, any> }): {
+  sql: string;
+  args: InValue[];
+} {
+  const parsedTableName = parseSqlIdentifier(tableName, 'table name');
+  const whereClause = prepareWhereClause(keys, TABLE_SCHEMAS[tableName]);
+
+  return {
+    sql: `DELETE FROM ${parsedTableName}${whereClause.sql}`,
+    args: whereClause.args,
+  };
+}
+
+type WhereValue = InValue | { startAt?: InValue; endAt?: InValue };
+
+export function prepareWhereClause(
+  filters: Record<string, WhereValue>,
+  schema: Record<string, StorageColumn>,
+): {
+  sql: string;
+  args: InValue[];
+} {
+  const conditions: string[] = [];
+  const args: InValue[] = [];
+
+  for (const [columnName, filterValue] of Object.entries(filters)) {
+    const column = schema[columnName];
+    if (!column) {
+      throw new Error(`Unknown column: ${columnName}`);
+    }
+
+    const parsedColumn = parseSqlIdentifier(columnName, 'column name');
+    const result = buildCondition(parsedColumn, filterValue);
+
+    conditions.push(result.condition);
+    args.push(...result.args);
+  }
+
+  return {
+    sql: conditions.length > 0 ? ` WHERE ${conditions.join(' AND ')}` : '',
+    args,
+  };
+}
+
+function buildCondition(columnName: string, filterValue: WhereValue): { condition: string; args: InValue[] } {
+  // Handle null values - IS NULL
+  if (filterValue === null) {
+    return { condition: `${columnName} IS NULL`, args: [] };
+  }
+
+  // Handle date range objects
+  if (typeof filterValue === 'object' && filterValue !== null && ('startAt' in filterValue || 'endAt' in filterValue)) {
+    return buildDateRangeCondition(columnName, filterValue);
+  }
+
+  // Handle exact match
+  return {
+    condition: `${columnName} = ?`,
+    args: [transformToSqlValue(filterValue)],
+  };
+}
+
+function buildDateRangeCondition(
+  columnName: string,
+  range: { startAt?: InValue; endAt?: InValue },
+): { condition: string; args: InValue[] } {
+  const conditions: string[] = [];
+  const args: InValue[] = [];
+
+  if (range.startAt !== undefined) {
+    conditions.push(`${columnName} >= ?`);
+    args.push(transformToSqlValue(range.startAt));
+  }
+
+  if (range.endAt !== undefined) {
+    conditions.push(`${columnName} <= ?`);
+    args.push(transformToSqlValue(range.endAt));
+  }
+
+  if (conditions.length === 0) {
+    throw new Error('Date range must specify at least startAt or endAt');
+  }
+
+  return {
+    condition: conditions.join(' AND '),
+    args,
+  };
+}
+
+type DateRangeFilter = {
+  startAt?: string;
+  endAt?: string;
+};
+
+/**
+ * Converts pagination date range to where clause date range format
+ * @param dateRange - The date range from pagination
+ * @param columnName - The timestamp column to filter on (defaults to 'createdAt')
+ * @returns Object with the date range filter, or empty object if no date range
+ */
+export function buildDateRangeFilter(
+  dateRange?: PaginationArgs['dateRange'],
+  columnName: string = 'createdAt',
+): Record<string, DateRangeFilter> {
+  if (!dateRange?.start && !dateRange?.end) {
+    return {};
+  }
+
+  const filter: DateRangeFilter = {};
+
+  if (dateRange.start) {
+    filter.startAt = new Date(dateRange.start).toISOString();
+  }
+
+  if (dateRange.end) {
+    filter.endAt = new Date(dateRange.end).toISOString();
+  }
+
+  return { [columnName]: filter };
+}
+
+/**
+ * Transforms SQL row data back to a typed object format
+ * Reverses the transformations done in prepareStatement
+ */
+export function transformFromSqlRow<T>({
+  tableName,
+  sqlRow,
+}: {
+  tableName: TABLE_NAMES;
+  sqlRow: Record<string, any>;
+}): T {
+  const result: Record<string, any> = {};
+  const jsonColumns = new Set(
+    Object.keys(TABLE_SCHEMAS[tableName])
+      .filter(key => TABLE_SCHEMAS[tableName][key]!.type === 'jsonb')
+      .map(key => key),
+  );
+  const dateColumns = new Set(
+    Object.keys(TABLE_SCHEMAS[tableName])
+      .filter(key => TABLE_SCHEMAS[tableName][key]!.type === 'timestamp')
+      .map(key => key),
+  );
+
+  for (const [key, value] of Object.entries(sqlRow)) {
+    if (value === null || value === undefined) {
+      result[key] = value;
+      continue;
+    }
+
+    if (dateColumns.has(key) && typeof value === 'string') {
+      result[key] = new Date(value);
+      continue;
+    }
+
+    if (jsonColumns.has(key) && typeof value === 'string') {
+      result[key] = safelyParseJSON(value);
+      continue;
+    }
+
+    result[key] = value;
+  }
+
+  return result as T;
+}

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

@@ -1,7 +1,8 @@
 import type { Client, InValue } from '@libsql/client';
-import type { WorkflowRun, WorkflowRuns, WorkflowRunState } from '@mastra/core';
 import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import type { WorkflowRun, WorkflowRuns } from '@mastra/core/storage';
 import { TABLE_WORKFLOW_SNAPSHOT, WorkflowsStorage } from '@mastra/core/storage';
+import type { WorkflowRunState, StepResult } from '@mastra/core/workflows';
 import type { StoreOperationsLibSQL } from '../operations';
 
 function parseWorkflowRun(row: Record<string, any>): WorkflowRun {
@@ -27,10 +28,232 @@ function parseWorkflowRun(row: Record<string, any>): WorkflowRun {
 export class WorkflowsLibSQL extends WorkflowsStorage {
   operations: StoreOperationsLibSQL;
   client: Client;
-  constructor({ operations, client }: { operations: StoreOperationsLibSQL; client: Client }) {
+  private readonly maxRetries: number;
+  private readonly initialBackoffMs: number;
+
+  constructor({
+    operations,
+    client,
+    maxRetries = 5,
+    initialBackoffMs = 500,
+  }: {
+    operations: StoreOperationsLibSQL;
+    client: Client;
+    maxRetries?: number;
+    initialBackoffMs?: number;
+  }) {
     super();
     this.operations = operations;
     this.client = client;
+    this.maxRetries = maxRetries;
+    this.initialBackoffMs = initialBackoffMs;
+
+    // Set PRAGMA settings to help with database locks
+    // Note: This is async but we can't await in constructor, so we'll handle it as a fire-and-forget
+    this.setupPragmaSettings().catch(err =>
+      this.logger.warn('LibSQL Workflows: Failed to setup PRAGMA settings.', err),
+    );
+  }
+
+  private async setupPragmaSettings() {
+    try {
+      // Set busy timeout to wait longer before returning busy errors
+      await this.client.execute('PRAGMA busy_timeout = 10000;');
+      this.logger.debug('LibSQL Workflows: PRAGMA busy_timeout=10000 set.');
+
+      // Enable WAL mode for better concurrency (if supported)
+      try {
+        await this.client.execute('PRAGMA journal_mode = WAL;');
+        this.logger.debug('LibSQL Workflows: PRAGMA journal_mode=WAL set.');
+      } catch {
+        this.logger.debug('LibSQL Workflows: WAL mode not supported, using default journal mode.');
+      }
+
+      // Set synchronous mode for better durability vs performance trade-off
+      try {
+        await this.client.execute('PRAGMA synchronous = NORMAL;');
+        this.logger.debug('LibSQL Workflows: PRAGMA synchronous=NORMAL set.');
+      } catch {
+        this.logger.debug('LibSQL Workflows: Failed to set synchronous mode.');
+      }
+    } catch (err) {
+      this.logger.warn('LibSQL Workflows: Failed to set PRAGMA settings.', err);
+    }
+  }
+
+  private async executeWithRetry<T>(operation: () => Promise<T>): Promise<T> {
+    let attempts = 0;
+    let backoff = this.initialBackoffMs;
+
+    while (attempts < this.maxRetries) {
+      try {
+        return await operation();
+      } catch (error: any) {
+        // Log the error details for debugging
+        this.logger.debug('LibSQL Workflows: Error caught in retry loop', {
+          errorType: error.constructor.name,
+          errorCode: error.code,
+          errorMessage: error.message,
+          attempts,
+          maxRetries: this.maxRetries,
+        });
+
+        // Check for various database lock/busy conditions
+        const isLockError =
+          error.code === 'SQLITE_BUSY' ||
+          error.code === 'SQLITE_LOCKED' ||
+          error.message?.toLowerCase().includes('database is locked') ||
+          error.message?.toLowerCase().includes('database table is locked') ||
+          error.message?.toLowerCase().includes('table is locked') ||
+          (error.constructor.name === 'SqliteError' && error.message?.toLowerCase().includes('locked'));
+
+        if (isLockError) {
+          attempts++;
+          if (attempts >= this.maxRetries) {
+            this.logger.error(
+              `LibSQL Workflows: Operation failed after ${this.maxRetries} attempts due to database lock: ${error.message}`,
+              { error, attempts, maxRetries: this.maxRetries },
+            );
+            throw error;
+          }
+          this.logger.warn(
+            `LibSQL Workflows: Attempt ${attempts} failed due to database lock. Retrying in ${backoff}ms...`,
+            { errorMessage: error.message, attempts, backoff, maxRetries: this.maxRetries },
+          );
+          await new Promise(resolve => setTimeout(resolve, backoff));
+          backoff *= 2;
+        } else {
+          // Not a lock error, re-throw immediately
+          this.logger.error('LibSQL Workflows: Non-lock error occurred, not retrying', { error });
+          throw error;
+        }
+      }
+    }
+    throw new Error('LibSQL Workflows: Max retries reached, but no error was re-thrown from the loop.');
+  }
+
+  async updateWorkflowResults({
+    workflowName,
+    runId,
+    stepId,
+    result,
+    runtimeContext,
+  }: {
+    workflowName: string;
+    runId: string;
+    stepId: string;
+    result: StepResult<any, any, any, any>;
+    runtimeContext: Record<string, any>;
+  }): Promise<Record<string, StepResult<any, any, any, any>>> {
+    return this.executeWithRetry(async () => {
+      // Use a transaction to ensure atomicity
+      const tx = await this.client.transaction('write');
+      try {
+        // Load existing snapshot within transaction
+        const existingSnapshotResult = await tx.execute({
+          sql: `SELECT snapshot FROM ${TABLE_WORKFLOW_SNAPSHOT} WHERE workflow_name = ? AND run_id = ?`,
+          args: [workflowName, runId],
+        });
+
+        let snapshot: WorkflowRunState;
+        if (!existingSnapshotResult.rows?.[0]) {
+          // Create new snapshot if none exists
+          snapshot = {
+            context: {},
+            activePaths: [],
+            timestamp: Date.now(),
+            suspendedPaths: {},
+            serializedStepGraph: [],
+            value: {},
+            waitingPaths: {},
+            status: 'pending',
+            runId: runId,
+            runtimeContext: {},
+          } as WorkflowRunState;
+        } else {
+          // Parse existing snapshot
+          const existingSnapshot = existingSnapshotResult.rows[0].snapshot;
+          snapshot = typeof existingSnapshot === 'string' ? JSON.parse(existingSnapshot) : existingSnapshot;
+        }
+
+        // Merge the new step result and runtime context
+        snapshot.context[stepId] = result;
+        snapshot.runtimeContext = { ...snapshot.runtimeContext, ...runtimeContext };
+
+        // Update the snapshot within the same transaction
+        await tx.execute({
+          sql: `UPDATE ${TABLE_WORKFLOW_SNAPSHOT} SET snapshot = ? WHERE workflow_name = ? AND run_id = ?`,
+          args: [JSON.stringify(snapshot), workflowName, runId],
+        });
+
+        await tx.commit();
+        return snapshot.context;
+      } catch (error) {
+        if (!tx.closed) {
+          await tx.rollback();
+        }
+        throw error;
+      }
+    });
+  }
+
+  async updateWorkflowState({
+    workflowName,
+    runId,
+    opts,
+  }: {
+    workflowName: string;
+    runId: string;
+    opts: {
+      status: string;
+      result?: StepResult<any, any, any, any>;
+      error?: string;
+      suspendedPaths?: Record<string, number[]>;
+      waitingPaths?: Record<string, number[]>;
+    };
+  }): Promise<WorkflowRunState | undefined> {
+    return this.executeWithRetry(async () => {
+      // Use a transaction to ensure atomicity
+      const tx = await this.client.transaction('write');
+      try {
+        // Load existing snapshot within transaction
+        const existingSnapshotResult = await tx.execute({
+          sql: `SELECT snapshot FROM ${TABLE_WORKFLOW_SNAPSHOT} WHERE workflow_name = ? AND run_id = ?`,
+          args: [workflowName, runId],
+        });
+
+        if (!existingSnapshotResult.rows?.[0]) {
+          await tx.rollback();
+          return undefined;
+        }
+
+        // Parse existing snapshot
+        const existingSnapshot = existingSnapshotResult.rows[0].snapshot;
+        const snapshot = typeof existingSnapshot === 'string' ? JSON.parse(existingSnapshot) : existingSnapshot;
+
+        if (!snapshot || !snapshot?.context) {
+          await tx.rollback();
+          throw new Error(`Snapshot not found for runId ${runId}`);
+        }
+
+        // Merge the new options with the existing snapshot
+        const updatedSnapshot = { ...snapshot, ...opts };
+
+        // Update the snapshot within the same transaction
+        await tx.execute({
+          sql: `UPDATE ${TABLE_WORKFLOW_SNAPSHOT} SET snapshot = ? WHERE workflow_name = ? AND run_id = ?`,
+          args: [JSON.stringify(updatedSnapshot), workflowName, runId],
+        });
+
+        await tx.commit();
+        return updatedSnapshot;
+      } catch (error) {
+        if (!tx.closed) {
+          await tx.rollback();
+        }
+        throw error;
+      }
+    });
   }
 
   async persistWorkflowSnapshot({