@mastra/dsql 1.5.1-alpha.1 → 1.5.2-alpha.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/dist/index.js CHANGED
@@ -1,6 +1,6 @@
1
1
  import { AuroraDSQLClient } from "@aws/aurora-dsql-node-postgres-connector";
2
2
  import { ErrorCategory, ErrorDomain, MastraError } from "@mastra/core/error";
3
- import { AgentsStorage, MastraStorage, MemoryStorage, ObservabilityStorage, ScoresStorage, TABLE_AGENTS, TABLE_AGENT_VERSIONS, TABLE_MESSAGES, TABLE_RESOURCES, TABLE_SCHEMAS, TABLE_SCORERS, TABLE_SPANS, TABLE_THREADS, TABLE_WORKFLOW_SNAPSHOT, TraceStatus, WorkflowsStorage, calculatePagination, createStorageErrorId, executeRetentionPrune, getDefaultValue, getSqlType, listTracesArgsSchema, matchesExpectedWorkflowStatus, normalizePerPage, resolveRetentionTargets, retentionCutoffMs, storageMessageMatchesMetadataFilter, toTraceSpans, transformScoreRow, validateStorageMetadataFilter } from "@mastra/core/storage";
3
+ import { AgentsStorage, MastraStorage, MemoryStorage, ObservabilityStorage, RUN_FENCING_TABLE_SCHEMAS, RunFenceConflictError, ScoresStorage, TABLE_AGENTS, TABLE_AGENT_VERSIONS, TABLE_MEMORY_RUN_FENCES, TABLE_MESSAGES, TABLE_RESOURCES, TABLE_SCHEMAS, TABLE_SCORERS, TABLE_SPANS, TABLE_THREADS, TABLE_WORKFLOW_RUN_OWNERS, TABLE_WORKFLOW_SNAPSHOT, TraceStatus, WorkflowsStorage, calculatePagination, createStorageErrorId, executeRetentionPrune, getDefaultValue, getSqlType, isRunFenceConflictError, listTracesArgsSchema, matchesExpectedWorkflowStatus, matchesRunFence, normalizePerPage, resolveRetentionTargets, resolveRunFence, retentionCutoffMs, storageMessageMatchesMetadataFilter, toTraceSpans, transformScoreRow, validateStorageMetadataFilter } from "@mastra/core/storage";
4
4
  import { Pool } from "pg";
5
5
  import { MastraBase } from "@mastra/core/base";
6
6
  import { parseSqlIdentifier } from "@mastra/core/utils";
@@ -1824,6 +1824,49 @@ var AgentsDSQL = class AgentsDSQL extends AgentsStorage {
1824
1824
  }
1825
1825
  };
1826
1826
  //#endregion
1827
+ //#region src/storage/db/run-fencing.ts
1828
+ /**
1829
+ * Epoch milliseconds on the database clock. Aurora DSQL never waits on a row
1830
+ * lock, so the transaction's start time is the time its statements run, and
1831
+ * every retry starts a new transaction with a fresh clock.
1832
+ */
1833
+ const DB_NOW_MS = `(EXTRACT(EPOCH FROM now()) * 1000)::bigint`;
1834
+ /**
1835
+ * Retry predicate for writes to an ownership or fence row that may not exist
1836
+ * yet. Two first writers both insert the same key; the loser normally fails
1837
+ * at commit with an OCC conflict (40001), but a duplicate-key error (23505)
1838
+ * means the same thing here: the retry sees the row and takes the update path.
1839
+ */
1840
+ function isRetriableRunFenceWrite(error) {
1841
+ return isRetriableError(error) || getErrorCode(error) === "23505";
1842
+ }
1843
+ /**
1844
+ * The check half of a fenced write; call it first inside the write's
1845
+ * transaction. Aurora DSQL has no blocking row locks: `FOR UPDATE` puts the
1846
+ * claim row in this transaction's write set instead, so a claim that commits
1847
+ * against the row while the write is in flight makes the write fail at commit
1848
+ * with an OCC conflict. The retry then sees the new claim and is rejected.
1849
+ */
1850
+ async function assertRunFence(t, claimsTable, fence, operation) {
1851
+ const row = await t.oneOrNone(`SELECT generation, "ownerId" FROM ${claimsTable} WHERE "runId" = $1 FOR UPDATE`, [fence.runId]);
1852
+ if (!matchesRunFence(row && {
1853
+ generation: Number(row.generation),
1854
+ ownerId: row.ownerId
1855
+ }, fence)) throw new RunFenceConflictError(fence, operation);
1856
+ }
1857
+ /**
1858
+ * Runs `write` directly when it carries no fence, or in a transaction behind
1859
+ * the fence check when it does. Callers own retries: wrap the call in
1860
+ * `withRetry` so an OCC conflict re-runs the check as well as the write.
1861
+ */
1862
+ function withRunFence(client, claimsTable, fence, operation, write) {
1863
+ if (!fence) return write(client);
1864
+ return client.tx(async (t) => {
1865
+ await assertRunFence(t, claimsTable, fence, operation);
1866
+ return write(t);
1867
+ });
1868
+ }
1869
+ //#endregion
1827
1870
  //#region src/storage/domains/memory/index.ts
1828
1871
  /**
1829
1872
  * Generate SQL placeholder string for IN clauses.
@@ -1883,6 +1926,10 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
1883
1926
  schema: TABLE_SCHEMAS[TABLE_MESSAGES],
1884
1927
  ifNotExists: ["resourceId"]
1885
1928
  });
1929
+ await this.#db.createTable({
1930
+ tableName: TABLE_MEMORY_RUN_FENCES,
1931
+ schema: RUN_FENCING_TABLE_SCHEMAS[TABLE_MEMORY_RUN_FENCES]
1932
+ });
1886
1933
  await this.createDefaultIndexes();
1887
1934
  await this.createCustomIndexes();
1888
1935
  }
@@ -1928,6 +1975,62 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
1928
1975
  await this.#db.clearTable({ tableName: TABLE_MESSAGES });
1929
1976
  await this.#db.clearTable({ tableName: TABLE_THREADS });
1930
1977
  await this.#db.clearTable({ tableName: TABLE_RESOURCES });
1978
+ await this.#db.clearTable({ tableName: TABLE_MEMORY_RUN_FENCES });
1979
+ }
1980
+ supportsRunFencing() {
1981
+ return true;
1982
+ }
1983
+ #runFencesTable() {
1984
+ return getTableName({
1985
+ indexName: TABLE_MEMORY_RUN_FENCES,
1986
+ schemaName: getSchemaName(this.#schema)
1987
+ });
1988
+ }
1989
+ /** Runs `write` behind the run's fence check when it carries a fence. */
1990
+ #fenced(fence, operation, write) {
1991
+ return withRunFence(this.#db.client, this.#runFencesTable(), resolveRunFence(this, fence), operation, write);
1992
+ }
1993
+ async raiseRunFence(fence) {
1994
+ try {
1995
+ const { result } = await withRetry(() => this.#db.client.oneOrNone(`INSERT INTO ${this.#runFencesTable()} AS f ("runId", generation, "ownerId")
1996
+ VALUES ($1, $2, $3)
1997
+ ON CONFLICT ("runId") DO UPDATE
1998
+ SET generation = EXCLUDED.generation, "ownerId" = EXCLUDED."ownerId", "retiredAt" = NULL
1999
+ WHERE f.generation < EXCLUDED.generation
2000
+ OR (f.generation = EXCLUDED.generation AND f."ownerId" = EXCLUDED."ownerId")
2001
+ RETURNING "runId"`, [
2002
+ fence.runId,
2003
+ fence.generation,
2004
+ fence.ownerId
2005
+ ]), { isRetriable: isRetriableRunFenceWrite });
2006
+ return result !== null;
2007
+ } catch (error) {
2008
+ throw new MastraError({
2009
+ id: createStorageErrorId("DSQL", "RAISE_RUN_FENCE", "FAILED"),
2010
+ domain: ErrorDomain.STORAGE,
2011
+ category: ErrorCategory.THIRD_PARTY,
2012
+ details: { runId: fence.runId }
2013
+ }, error);
2014
+ }
2015
+ }
2016
+ async retireRunFence(fence) {
2017
+ try {
2018
+ const { result } = await withRetry(() => this.#db.client.oneOrNone(`UPDATE ${this.#runFencesTable()} SET "retiredAt" = ${DB_NOW_MS}
2019
+ WHERE "runId" = $1 AND generation = $2 AND "ownerId" = $3
2020
+ RETURNING "runId"`, [
2021
+ fence.runId,
2022
+ fence.generation,
2023
+ fence.ownerId
2024
+ ]), { isRetriable: isRetriableRunFenceWrite });
2025
+ return result !== null;
2026
+ } catch (error) {
2027
+ throw new MastraError({
2028
+ id: createStorageErrorId("DSQL", "RETIRE_RUN_FENCE", "FAILED"),
2029
+ domain: ErrorDomain.STORAGE,
2030
+ category: ErrorCategory.THIRD_PARTY,
2031
+ details: { runId: fence.runId }
2032
+ }, error);
2033
+ }
1931
2034
  }
1932
2035
  /**
1933
2036
  * Normalizes message row from database by applying createdAtZ fallback
@@ -2064,7 +2167,7 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2064
2167
  throw mastraError;
2065
2168
  }
2066
2169
  }
2067
- async saveThread({ thread }) {
2170
+ async saveThread({ thread, fence }) {
2068
2171
  const tableName = getTableName({
2069
2172
  indexName: TABLE_THREADS,
2070
2173
  schemaName: getSchemaName(this.#schema)
@@ -2072,7 +2175,7 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2072
2175
  const createdAt = toUtcISOString(thread.createdAt);
2073
2176
  const updatedAt = toUtcISOString(thread.updatedAt);
2074
2177
  await withRetry(async () => {
2075
- await this.#db.client.none(`INSERT INTO ${tableName} (
2178
+ await this.#fenced(fence, "saveThread", (q) => q.none(`INSERT INTO ${tableName} (
2076
2179
  id,
2077
2180
  "resourceId",
2078
2181
  title,
@@ -2098,10 +2201,11 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2098
2201
  createdAt,
2099
2202
  updatedAt,
2100
2203
  updatedAt
2101
- ]);
2204
+ ]));
2102
2205
  }, { onRetry: (error, attempt, delay) => {
2103
2206
  this.logger?.warn?.(`saveThread retry ${attempt} for ${thread.id} after ${delay}ms: ${error.message}`);
2104
2207
  } }).catch((error) => {
2208
+ if (isRunFenceConflictError(error)) throw error;
2105
2209
  throw new MastraError({
2106
2210
  id: createStorageErrorId("DSQL", "SAVE_THREAD", "FAILED"),
2107
2211
  domain: ErrorDomain.STORAGE,
@@ -2111,13 +2215,13 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2111
2215
  });
2112
2216
  return thread;
2113
2217
  }
2114
- async updateThread({ id, title, metadata }) {
2218
+ async updateThread({ id, title, metadata, fence }) {
2115
2219
  const threadTableName = getTableName({
2116
2220
  indexName: TABLE_THREADS,
2117
2221
  schemaName: getSchemaName(this.#schema)
2118
2222
  });
2119
- const { result } = await withRetry(async () => {
2120
- const existingThread = await this.getThreadById({ threadId: id });
2223
+ const { result } = await withRetry(() => this.#fenced(fence, "updateThread", async (q) => {
2224
+ const existingThread = await q.oneOrNone(`SELECT metadata FROM ${threadTableName} WHERE id = $1`, [id]);
2121
2225
  if (!existingThread) throw new MastraError({
2122
2226
  id: createStorageErrorId("DSQL", "UPDATE_THREAD", "NOT_FOUND"),
2123
2227
  domain: ErrorDomain.STORAGE,
@@ -2129,11 +2233,11 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2129
2233
  }
2130
2234
  });
2131
2235
  const mergedMetadata = {
2132
- ...existingThread.metadata,
2236
+ ...typeof existingThread.metadata === "string" ? JSON.parse(existingThread.metadata) : existingThread.metadata,
2133
2237
  ...metadata
2134
2238
  };
2135
2239
  const now = (/* @__PURE__ */ new Date()).toISOString();
2136
- const thread = await this.#db.client.one(`UPDATE ${threadTableName}
2240
+ return q.one(`UPDATE ${threadTableName}
2137
2241
  SET
2138
2242
  title = COALESCE($1, title),
2139
2243
  metadata = $2,
@@ -2148,15 +2252,14 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2148
2252
  now,
2149
2253
  id
2150
2254
  ]);
2151
- return {
2152
- id: thread.id,
2153
- resourceId: thread.resourceId,
2154
- title: thread.title,
2155
- metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
2156
- createdAt: thread.createdAtZ || thread.createdAt,
2157
- updatedAt: thread.updatedAtZ || thread.updatedAt
2158
- };
2159
- }, { onRetry: (error, attempt, delay) => {
2255
+ }).then((thread) => ({
2256
+ id: thread.id,
2257
+ resourceId: thread.resourceId,
2258
+ title: thread.title,
2259
+ metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
2260
+ createdAt: thread.createdAtZ || thread.createdAt,
2261
+ updatedAt: thread.updatedAtZ || thread.updatedAt
2262
+ })), { onRetry: (error, attempt, delay) => {
2160
2263
  this.logger?.warn?.(`updateThread retry ${attempt} for ${id} after ${delay}ms: ${error.message}`);
2161
2264
  } }).catch((error) => {
2162
2265
  if (error instanceof MastraError) throw error;
@@ -2481,7 +2584,7 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2481
2584
  throw mastraError;
2482
2585
  }
2483
2586
  }
2484
- async saveMessages({ messages }) {
2587
+ async saveMessages({ messages, fence: explicitFence }) {
2485
2588
  if (messages.length === 0) return { messages: [] };
2486
2589
  const threadId = messages[0]?.threadId;
2487
2590
  if (!threadId) throw new MastraError({
@@ -2505,8 +2608,10 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2505
2608
  indexName: TABLE_THREADS,
2506
2609
  schemaName: getSchemaName(this.#schema)
2507
2610
  });
2611
+ const fence = resolveRunFence(this, explicitFence);
2508
2612
  await withRetry(async () => {
2509
2613
  await this.#db.client.tx(async (t) => {
2614
+ if (fence) await assertRunFence(t, this.#runFencesTable(), fence, "saveMessages");
2510
2615
  const messageInserts = messages.map((message) => {
2511
2616
  if (!message.threadId) throw new Error(`Expected to find a threadId for message, but couldn't find one. An unexpected error has occurred.`);
2512
2617
  if (!message.resourceId) throw new Error(`Expected to find a resourceId for message, but couldn't find one. An unexpected error has occurred.`);
@@ -2545,6 +2650,7 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2545
2650
  }, { onRetry: (error, attempt, delay) => {
2546
2651
  this.logger?.warn?.(`saveMessages retry ${attempt} for thread ${threadId} after ${delay}ms: ${error.message}`);
2547
2652
  } }).catch((error) => {
2653
+ if (isRunFenceConflictError(error)) throw error;
2548
2654
  throw new MastraError({
2549
2655
  id: createStorageErrorId("DSQL", "SAVE_MESSAGES", "FAILED"),
2550
2656
  domain: ErrorDomain.STORAGE,
@@ -2565,8 +2671,9 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2565
2671
  });
2566
2672
  return { messages: new MessageList().add(messagesWithParsedContent, "memory").get.all.db() };
2567
2673
  }
2568
- async updateMessages({ messages }) {
2674
+ async updateMessages({ messages, fence: explicitFence }) {
2569
2675
  if (messages.length === 0) return [];
2676
+ const fence = resolveRunFence(this, explicitFence);
2570
2677
  const messageIds = messages.map((m) => m.id);
2571
2678
  const selectQuery = `SELECT id, content, role, type, "createdAt", "createdAtZ", thread_id AS "threadId", "resourceId" FROM ${getTableName({
2572
2679
  indexName: TABLE_MESSAGES,
@@ -2583,6 +2690,7 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2583
2690
  const threadIdsToUpdate = /* @__PURE__ */ new Set();
2584
2691
  await withRetry(async () => {
2585
2692
  await this.#db.client.tx(async (t) => {
2693
+ if (fence) await assertRunFence(t, this.#runFencesTable(), fence, "updateMessages");
2586
2694
  const queries = [];
2587
2695
  const columnMapping = { threadId: "thread_id" };
2588
2696
  for (const existingMessage of existingMessages) {
@@ -2635,6 +2743,7 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2635
2743
  }, { onRetry: (error, attempt, delay) => {
2636
2744
  this.logger?.warn?.(`updateMessages retry ${attempt} for ${messageIds.length} messages after ${delay}ms: ${error.message}`);
2637
2745
  } }).catch((error) => {
2746
+ if (isRunFenceConflictError(error)) throw error;
2638
2747
  throw new MastraError({
2639
2748
  id: createStorageErrorId("DSQL", "UPDATE_MESSAGES", "FAILED"),
2640
2749
  domain: ErrorDomain.STORAGE,
@@ -2653,8 +2762,9 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2653
2762
  return message;
2654
2763
  });
2655
2764
  }
2656
- async deleteMessages(messageIds) {
2765
+ async deleteMessages(messageIds, options) {
2657
2766
  if (!messageIds || messageIds.length === 0) return;
2767
+ const fence = resolveRunFence(this, options?.fence);
2658
2768
  const messageTableName = getTableName({
2659
2769
  indexName: TABLE_MESSAGES,
2660
2770
  schemaName: getSchemaName(this.#schema)
@@ -2665,6 +2775,7 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2665
2775
  });
2666
2776
  await withRetry(async () => {
2667
2777
  await this.#db.client.tx(async (t) => {
2778
+ if (fence) await assertRunFence(t, this.#runFencesTable(), fence, "deleteMessages");
2668
2779
  const placeholders = messageIds.map((_, idx) => `$${idx + 1}`).join(",");
2669
2780
  const threadIds = (await t.manyOrNone(`SELECT DISTINCT thread_id FROM ${messageTableName} WHERE id IN (${placeholders})`, messageIds))?.map((msg) => msg.thread_id).filter(Boolean) || [];
2670
2781
  await t.none(`DELETE FROM ${messageTableName} WHERE id IN (${placeholders})`, messageIds);
@@ -2676,6 +2787,7 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2676
2787
  }, { onRetry: (error, attempt, delay) => {
2677
2788
  this.logger?.warn?.(`deleteMessages retry ${attempt} for ${messageIds.length} messages after ${delay}ms: ${error.message}`);
2678
2789
  } }).catch((error) => {
2790
+ if (isRunFenceConflictError(error)) throw error;
2679
2791
  throw new MastraError({
2680
2792
  id: createStorageErrorId("DSQL", "DELETE_MESSAGES", "FAILED"),
2681
2793
  domain: ErrorDomain.STORAGE,
@@ -2720,23 +2832,42 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2720
2832
  });
2721
2833
  return resource;
2722
2834
  }
2723
- async updateResource({ resourceId, workingMemory, metadata }) {
2835
+ async updateResource({ resourceId, workingMemory, metadata, fence }) {
2724
2836
  const tableName = getTableName({
2725
2837
  indexName: TABLE_RESOURCES,
2726
2838
  schemaName: getSchemaName(this.#schema)
2727
2839
  });
2728
- const { result } = await withRetry(async () => {
2729
- const existingResource = await this.getResourceById({ resourceId });
2730
- if (!existingResource) {
2840
+ const { result } = await withRetry(() => this.#fenced(fence, "updateResource", async (q) => {
2841
+ const row = await q.oneOrNone(`SELECT * FROM ${tableName} WHERE id = $1`, [resourceId]);
2842
+ if (!row) {
2843
+ const now = /* @__PURE__ */ new Date();
2731
2844
  const newResource = {
2732
2845
  id: resourceId,
2733
2846
  workingMemory,
2734
2847
  metadata: metadata || {},
2735
- createdAt: /* @__PURE__ */ new Date(),
2736
- updatedAt: /* @__PURE__ */ new Date()
2848
+ createdAt: now,
2849
+ updatedAt: now
2737
2850
  };
2738
- return this.saveResource({ resource: newResource });
2851
+ const nowIso = now.toISOString();
2852
+ await q.none(`INSERT INTO ${tableName} (id, "workingMemory", metadata, "createdAt", "createdAtZ", "updatedAt", "updatedAtZ")
2853
+ VALUES ($1, $2, $3, $4, $5, $6, $7)`, [
2854
+ resourceId,
2855
+ workingMemory ?? null,
2856
+ JSON.stringify(newResource.metadata),
2857
+ nowIso,
2858
+ nowIso,
2859
+ nowIso,
2860
+ nowIso
2861
+ ]);
2862
+ return newResource;
2739
2863
  }
2864
+ const existingResource = {
2865
+ id: row.id,
2866
+ createdAt: row.createdAtZ || row.createdAt,
2867
+ updatedAt: row.updatedAtZ || row.updatedAt,
2868
+ workingMemory: row.workingMemory,
2869
+ metadata: typeof row.metadata === "string" ? JSON.parse(row.metadata) : row.metadata
2870
+ };
2740
2871
  const updatedResource = {
2741
2872
  ...existingResource,
2742
2873
  workingMemory: workingMemory !== void 0 ? workingMemory : existingResource.workingMemory,
@@ -2766,9 +2897,9 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
2766
2897
  values.push(updatedResource.updatedAt.toISOString());
2767
2898
  paramIndex++;
2768
2899
  values.push(resourceId);
2769
- await this.#db.client.none(`UPDATE ${tableName} SET ${updates.join(", ")} WHERE id = $${paramIndex}`, values);
2900
+ await q.none(`UPDATE ${tableName} SET ${updates.join(", ")} WHERE id = $${paramIndex}`, values);
2770
2901
  return updatedResource;
2771
- }, { onRetry: (error, attempt, delay) => {
2902
+ }), { onRetry: (error, attempt, delay) => {
2772
2903
  this.logger?.warn?.(`updateResource retry ${attempt} for ${resourceId} after ${delay}ms: ${error.message}`);
2773
2904
  } }).catch((error) => {
2774
2905
  if (error instanceof MastraError) throw error;
@@ -3718,6 +3849,17 @@ var ScoresDSQL = class ScoresDSQL extends ScoresStorage {
3718
3849
  };
3719
3850
  //#endregion
3720
3851
  //#region src/storage/domains/workflows/index.ts
3852
+ const RUN_OWNER_COLUMNS = `generation, "ownerId", "leaseExpiresAt", ${DB_NOW_MS} AS "nowMs"`;
3853
+ function toRunOwnershipRecord(runId, row) {
3854
+ const leaseExpiresAt = row.leaseExpiresAt === null ? null : Number(row.leaseExpiresAt);
3855
+ return {
3856
+ runId,
3857
+ generation: Number(row.generation),
3858
+ ownerId: row.ownerId,
3859
+ leaseExpiresAt: leaseExpiresAt === null ? null : new Date(leaseExpiresAt),
3860
+ live: leaseExpiresAt !== null && leaseExpiresAt > Number(row.nowMs)
3861
+ };
3862
+ }
3721
3863
  function parseWorkflowRun(row) {
3722
3864
  let parsedSnapshot = row.snapshot;
3723
3865
  if (typeof parsedSnapshot === "string") try {
@@ -3779,6 +3921,10 @@ var WorkflowsDSQL = class WorkflowsDSQL extends WorkflowsStorage {
3779
3921
  schema: TABLE_SCHEMAS[TABLE_WORKFLOW_SNAPSHOT],
3780
3922
  ifNotExists: ["resourceId"]
3781
3923
  });
3924
+ await this.#db.createTable({
3925
+ tableName: TABLE_WORKFLOW_RUN_OWNERS,
3926
+ schema: RUN_FENCING_TABLE_SCHEMAS[TABLE_WORKFLOW_RUN_OWNERS]
3927
+ });
3782
3928
  await this.createDefaultIndexes();
3783
3929
  await this.createCustomIndexes();
3784
3930
  }
@@ -3795,11 +3941,130 @@ var WorkflowsDSQL = class WorkflowsDSQL extends WorkflowsStorage {
3795
3941
  }
3796
3942
  async dangerouslyClearAll() {
3797
3943
  await this.#db.clearTable({ tableName: TABLE_WORKFLOW_SNAPSHOT });
3944
+ await this.#db.clearTable({ tableName: TABLE_WORKFLOW_RUN_OWNERS });
3798
3945
  }
3799
- async updateWorkflowResults({ workflowName, runId, stepId, result, requestContext, state }) {
3946
+ supportsRunFencing() {
3947
+ return true;
3948
+ }
3949
+ #runOwnersTable() {
3950
+ return getTableName({
3951
+ indexName: TABLE_WORKFLOW_RUN_OWNERS,
3952
+ schemaName: getSchemaName(this.#schema)
3953
+ });
3954
+ }
3955
+ async #readRunOwner(q, runId) {
3956
+ const row = await q.oneOrNone(`SELECT ${RUN_OWNER_COLUMNS} FROM ${this.#runOwnersTable()} WHERE "runId" = $1`, [runId]);
3957
+ return row ? toRunOwnershipRecord(runId, row) : null;
3958
+ }
3959
+ #ownershipError(operation, runId, error) {
3960
+ return new MastraError({
3961
+ id: createStorageErrorId("DSQL", operation, "FAILED"),
3962
+ domain: ErrorDomain.STORAGE,
3963
+ category: ErrorCategory.THIRD_PARTY,
3964
+ details: { runId }
3965
+ }, error);
3966
+ }
3967
+ async claimRunOwnership({ runId, ownerId, leaseMs, force, expectedGeneration }) {
3968
+ const table = this.#runOwnersTable();
3969
+ try {
3970
+ const { result } = await withRetry(() => this.#db.client.tx(async (t) => {
3971
+ if (!await t.oneOrNone(`SELECT 1 FROM ${table} WHERE "runId" = $1 FOR UPDATE`, [runId])) {
3972
+ if (expectedGeneration !== void 0 && expectedGeneration !== 0) return {
3973
+ acquired: false,
3974
+ record: null
3975
+ };
3976
+ return {
3977
+ acquired: true,
3978
+ record: toRunOwnershipRecord(runId, await t.one(`INSERT INTO ${table} ("runId", generation, "ownerId", "leaseExpiresAt", "updatedAt")
3979
+ VALUES ($1, ${DB_NOW_MS}, $2, ${DB_NOW_MS} + $3, ${DB_NOW_MS})
3980
+ RETURNING ${RUN_OWNER_COLUMNS}`, [
3981
+ runId,
3982
+ ownerId,
3983
+ leaseMs
3984
+ ]))
3985
+ };
3986
+ }
3987
+ const claimed = await t.oneOrNone(`UPDATE ${table}
3988
+ SET generation = generation + 1, "ownerId" = $2, "leaseExpiresAt" = ${DB_NOW_MS} + $3,
3989
+ "updatedAt" = ${DB_NOW_MS}
3990
+ WHERE "runId" = $1
3991
+ AND ($4::bigint IS NULL OR generation = $4)
3992
+ AND ($5::boolean OR "leaseExpiresAt" IS NULL OR "leaseExpiresAt" <= ${DB_NOW_MS})
3993
+ RETURNING ${RUN_OWNER_COLUMNS}`, [
3994
+ runId,
3995
+ ownerId,
3996
+ leaseMs,
3997
+ expectedGeneration ?? null,
3998
+ force === true
3999
+ ]);
4000
+ if (claimed) return {
4001
+ acquired: true,
4002
+ record: toRunOwnershipRecord(runId, claimed)
4003
+ };
4004
+ return {
4005
+ acquired: false,
4006
+ record: await this.#readRunOwner(t, runId)
4007
+ };
4008
+ }), { isRetriable: isRetriableRunFenceWrite });
4009
+ return result;
4010
+ } catch (error) {
4011
+ throw this.#ownershipError("CLAIM_RUN_OWNERSHIP", runId, error);
4012
+ }
4013
+ }
4014
+ async renewRunOwnership({ leaseMs, ...fence }) {
4015
+ try {
4016
+ const { result } = await withRetry(async () => {
4017
+ const renewed = await this.#db.client.oneOrNone(`UPDATE ${this.#runOwnersTable()}
4018
+ SET "leaseExpiresAt" = ${DB_NOW_MS} + $4, "updatedAt" = ${DB_NOW_MS}
4019
+ WHERE "runId" = $1 AND generation = $2 AND "ownerId" = $3 AND "leaseExpiresAt" IS NOT NULL
4020
+ RETURNING ${RUN_OWNER_COLUMNS}`, [
4021
+ fence.runId,
4022
+ fence.generation,
4023
+ fence.ownerId,
4024
+ leaseMs
4025
+ ]);
4026
+ if (renewed) return {
4027
+ renewed: true,
4028
+ record: toRunOwnershipRecord(fence.runId, renewed)
4029
+ };
4030
+ return {
4031
+ renewed: false,
4032
+ record: await this.#readRunOwner(this.#db.client, fence.runId)
4033
+ };
4034
+ });
4035
+ return result;
4036
+ } catch (error) {
4037
+ throw this.#ownershipError("RENEW_RUN_OWNERSHIP", fence.runId, error);
4038
+ }
4039
+ }
4040
+ async releaseRunOwnership(fence) {
4041
+ try {
4042
+ const { result } = await withRetry(() => this.#db.client.oneOrNone(`UPDATE ${this.#runOwnersTable()}
4043
+ SET "leaseExpiresAt" = NULL, "updatedAt" = ${DB_NOW_MS}
4044
+ WHERE "runId" = $1 AND generation = $2 AND "ownerId" = $3
4045
+ RETURNING "runId"`, [
4046
+ fence.runId,
4047
+ fence.generation,
4048
+ fence.ownerId
4049
+ ]));
4050
+ return result !== null;
4051
+ } catch (error) {
4052
+ throw this.#ownershipError("RELEASE_RUN_OWNERSHIP", fence.runId, error);
4053
+ }
4054
+ }
4055
+ async getRunOwnership({ runId }) {
4056
+ try {
4057
+ return await this.#readRunOwner(this.#db.client, runId);
4058
+ } catch (error) {
4059
+ throw this.#ownershipError("GET_RUN_OWNERSHIP", runId, error);
4060
+ }
4061
+ }
4062
+ async updateWorkflowResults({ workflowName, runId, stepId, result, requestContext, state, fence: explicitFence }) {
4063
+ const fence = resolveRunFence(this, explicitFence, runId);
3800
4064
  try {
3801
4065
  const { result: context } = await withRetry(async () => {
3802
4066
  return this.#db.client.tx(async (t) => {
4067
+ if (fence) await assertRunFence(t, this.#runOwnersTable(), fence, "updateWorkflowResults");
3803
4068
  const tableName = getTableName({
3804
4069
  indexName: TABLE_WORKFLOW_SNAPSHOT,
3805
4070
  schemaName: getSchemaName(this.#schema)
@@ -3848,6 +4113,7 @@ var WorkflowsDSQL = class WorkflowsDSQL extends WorkflowsStorage {
3848
4113
  } });
3849
4114
  return context;
3850
4115
  } catch (error) {
4116
+ if (isRunFenceConflictError(error)) throw error;
3851
4117
  throw new MastraError({
3852
4118
  id: createStorageErrorId("DSQL", "UPDATE_WORKFLOW_RESULTS", "FAILED"),
3853
4119
  domain: ErrorDomain.STORAGE,
@@ -3860,10 +4126,12 @@ var WorkflowsDSQL = class WorkflowsDSQL extends WorkflowsStorage {
3860
4126
  }, error);
3861
4127
  }
3862
4128
  }
3863
- async updateWorkflowState({ workflowName, runId, opts }) {
4129
+ async updateWorkflowState({ workflowName, runId, opts, fence: explicitFence }) {
4130
+ const fence = resolveRunFence(this, explicitFence, runId);
3864
4131
  try {
3865
4132
  const { result } = await withRetry(async () => {
3866
4133
  return this.#db.client.tx(async (t) => {
4134
+ if (fence) await assertRunFence(t, this.#runOwnersTable(), fence, "updateWorkflowState");
3867
4135
  const tableName = getTableName({
3868
4136
  indexName: TABLE_WORKFLOW_SNAPSHOT,
3869
4137
  schemaName: getSchemaName(this.#schema)
@@ -3892,6 +4160,7 @@ var WorkflowsDSQL = class WorkflowsDSQL extends WorkflowsStorage {
3892
4160
  } });
3893
4161
  return result;
3894
4162
  } catch (error) {
4163
+ if (isRunFenceConflictError(error)) throw error;
3895
4164
  throw new MastraError({
3896
4165
  id: createStorageErrorId("DSQL", "UPDATE_WORKFLOW_STATE", "FAILED"),
3897
4166
  domain: ErrorDomain.STORAGE,
@@ -3903,12 +4172,13 @@ var WorkflowsDSQL = class WorkflowsDSQL extends WorkflowsStorage {
3903
4172
  }, error);
3904
4173
  }
3905
4174
  }
3906
- async persistWorkflowSnapshot({ workflowName, runId, resourceId, snapshot, createdAt, updatedAt }) {
4175
+ async persistWorkflowSnapshot({ workflowName, runId, resourceId, snapshot, createdAt, updatedAt, fence }) {
3907
4176
  try {
3908
4177
  const now = /* @__PURE__ */ new Date();
3909
4178
  const createdAtValue = createdAt ? createdAt : now;
3910
4179
  const updatedAtValue = updatedAt ? updatedAt : now;
3911
- await this.#db.client.none(`INSERT INTO ${getTableName({
4180
+ const resolvedFence = resolveRunFence(this, fence, runId);
4181
+ await withRetry(() => withRunFence(this.#db.client, this.#runOwnersTable(), resolvedFence, "persistWorkflowSnapshot", (q) => q.none(`INSERT INTO ${getTableName({
3912
4182
  indexName: TABLE_WORKFLOW_SNAPSHOT,
3913
4183
  schemaName: getSchemaName(this.#schema)
3914
4184
  })} (workflow_name, run_id, "resourceId", snapshot, "createdAt", "updatedAt")
@@ -3921,8 +4191,9 @@ var WorkflowsDSQL = class WorkflowsDSQL extends WorkflowsStorage {
3921
4191
  JSON.stringify(snapshot),
3922
4192
  createdAtValue.toISOString(),
3923
4193
  updatedAtValue.toISOString()
3924
- ]);
4194
+ ])));
3925
4195
  } catch (error) {
4196
+ if (isRunFenceConflictError(error)) throw error;
3926
4197
  throw new MastraError({
3927
4198
  id: createStorageErrorId("DSQL", "PERSIST_WORKFLOW_SNAPSHOT", "FAILED"),
3928
4199
  domain: ErrorDomain.STORAGE,
@@ -3988,13 +4259,15 @@ var WorkflowsDSQL = class WorkflowsDSQL extends WorkflowsStorage {
3988
4259
  }, error);
3989
4260
  }
3990
4261
  }
3991
- async deleteWorkflowRunById({ runId, workflowName }) {
4262
+ async deleteWorkflowRunById({ runId, workflowName, fence }) {
3992
4263
  try {
3993
- await this.#db.client.none(`DELETE FROM ${getTableName({
4264
+ const resolvedFence = resolveRunFence(this, fence, runId);
4265
+ await withRetry(() => withRunFence(this.#db.client, this.#runOwnersTable(), resolvedFence, "deleteWorkflowRunById", (q) => q.none(`DELETE FROM ${getTableName({
3994
4266
  indexName: TABLE_WORKFLOW_SNAPSHOT,
3995
4267
  schemaName: getSchemaName(this.#schema)
3996
- })} WHERE run_id = $1 AND workflow_name = $2`, [runId, workflowName]);
4268
+ })} WHERE run_id = $1 AND workflow_name = $2`, [runId, workflowName])));
3997
4269
  } catch (error) {
4270
+ if (isRunFenceConflictError(error)) throw error;
3998
4271
  throw new MastraError({
3999
4272
  id: createStorageErrorId("DSQL", "DELETE_WORKFLOW_RUN_BY_ID", "FAILED"),
4000
4273
  domain: ErrorDomain.STORAGE,