@mastra/dynamodb 1.4.1-alpha.0 → 1.4.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,7 +1,7 @@
1
1
  import { DescribeTableCommand, DynamoDBClient } from "@aws-sdk/client-dynamodb";
2
2
  import { DynamoDBDocumentClient } from "@aws-sdk/lib-dynamodb";
3
3
  import { ErrorCategory, ErrorDomain, MastraError } from "@mastra/core/error";
4
- import { BackgroundTasksStorage, MastraCompositeStore, MemoryStorage, SCORERS_SCHEMA, ScoresStorage, TABLE_BACKGROUND_TASKS, TABLE_MESSAGES, TABLE_RESOURCES, TABLE_SCORERS, TABLE_THREADS, TABLE_WORKFLOW_SNAPSHOT, WorkflowsStorage, calculatePagination, createStorageErrorId, matchesExpectedWorkflowStatus, normalizePerPage, storageMessageMatchesMetadataFilter, validateStorageMetadataFilter } from "@mastra/core/storage";
4
+ import { BackgroundTasksStorage, MastraCompositeStore, MemoryStorage, RunFenceConflictError, SCORERS_SCHEMA, ScoresStorage, TABLE_BACKGROUND_TASKS, TABLE_MEMORY_RUN_FENCES, TABLE_MESSAGES, TABLE_RESOURCES, TABLE_SCORERS, TABLE_THREADS, TABLE_WORKFLOW_RUN_OWNERS, TABLE_WORKFLOW_SNAPSHOT, WorkflowsStorage, calculatePagination, createStorageErrorId, isRunFenceConflictError, matchesExpectedWorkflowStatus, matchesRunFence, normalizePerPage, resolveRunFence, storageMessageMatchesMetadataFilter, validateStorageMetadataFilter } from "@mastra/core/storage";
5
5
  import { Entity, Service } from "electrodb";
6
6
  import { MessageList } from "@mastra/core/agent";
7
7
  import { saveScorePayloadSchema } from "@mastra/core/evals";
@@ -498,6 +498,80 @@ const resourceEntity = new Entity({
498
498
  } }
499
499
  });
500
500
  //#endregion
501
+ //#region src/entities/run-fencing.ts
502
+ /**
503
+ * The current claim on a durable workflow run. `leaseExpiresAt` holds the
504
+ * lease's expiry in epoch milliseconds and is removed when the owner releases
505
+ * the run; `generation` and `ownerId` are kept. `updatedAt` is the epoch
506
+ * milliseconds of the last claim, renewal or release.
507
+ */
508
+ const workflowRunOwnerEntity = new Entity({
509
+ model: {
510
+ entity: "workflow_run_owner",
511
+ version: "1",
512
+ service: "mastra"
513
+ },
514
+ attributes: {
515
+ entity: {
516
+ type: "string",
517
+ required: true
518
+ },
519
+ run_id: {
520
+ type: "string",
521
+ required: true
522
+ },
523
+ generation: { type: "number" },
524
+ ownerId: { type: "string" },
525
+ leaseExpiresAt: { type: "number" },
526
+ updatedAt: { type: "number" }
527
+ },
528
+ indexes: { primary: {
529
+ pk: {
530
+ field: "pk",
531
+ composite: ["entity", "run_id"]
532
+ },
533
+ sk: {
534
+ field: "sk",
535
+ composite: ["entity"]
536
+ }
537
+ } }
538
+ });
539
+ /**
540
+ * The newest fence memory has seen for a run. Fenced memory writes check it.
541
+ * `retiredAt` is set, in epoch milliseconds, once the execution holding the
542
+ * fence settles, and removed when a fence is raised again.
543
+ */
544
+ const memoryRunFenceEntity = new Entity({
545
+ model: {
546
+ entity: "memory_run_fence",
547
+ version: "1",
548
+ service: "mastra"
549
+ },
550
+ attributes: {
551
+ entity: {
552
+ type: "string",
553
+ required: true
554
+ },
555
+ run_id: {
556
+ type: "string",
557
+ required: true
558
+ },
559
+ generation: { type: "number" },
560
+ ownerId: { type: "string" },
561
+ retiredAt: { type: "number" }
562
+ },
563
+ indexes: { primary: {
564
+ pk: {
565
+ field: "pk",
566
+ composite: ["entity", "run_id"]
567
+ },
568
+ sk: {
569
+ field: "sk",
570
+ composite: ["entity"]
571
+ }
572
+ } }
573
+ });
574
+ //#endregion
501
575
  //#region src/entities/score.ts
502
576
  const scoreEntity = new Entity({
503
577
  model: {
@@ -1128,7 +1202,9 @@ function getElectroDbService(client, tableName) {
1128
1202
  workflow_snapshot: workflowSnapshotEntity,
1129
1203
  resource: resourceEntity,
1130
1204
  score: scoreEntity,
1131
- background_task: backgroundTaskEntity
1205
+ background_task: backgroundTaskEntity,
1206
+ workflow_run_owner: workflowRunOwnerEntity,
1207
+ memory_run_fence: memoryRunFenceEntity
1132
1208
  }, {
1133
1209
  client,
1134
1210
  table: tableName
@@ -1166,7 +1242,9 @@ const ENTITY_MAP = {
1166
1242
  [TABLE_RESOURCES]: "resource",
1167
1243
  [TABLE_WORKFLOW_SNAPSHOT]: "workflow_snapshot",
1168
1244
  [TABLE_SCORERS]: "score",
1169
- [TABLE_BACKGROUND_TASKS]: "background_task"
1245
+ [TABLE_BACKGROUND_TASKS]: "background_task",
1246
+ [TABLE_WORKFLOW_RUN_OWNERS]: "workflow_run_owner",
1247
+ [TABLE_MEMORY_RUN_FENCES]: "memory_run_fence"
1170
1248
  };
1171
1249
  /**
1172
1250
  * Gets the primary key fields for an entity based on the scanned item
@@ -1185,6 +1263,10 @@ function getDeleteKey(entityName, item) {
1185
1263
  key.workflow_name = item.workflow_name;
1186
1264
  key.run_id = item.run_id;
1187
1265
  break;
1266
+ case "workflow_run_owner":
1267
+ case "memory_run_fence":
1268
+ key.run_id = item.run_id;
1269
+ break;
1188
1270
  default: key.id = item.id;
1189
1271
  }
1190
1272
  return key;
@@ -1520,6 +1602,91 @@ function getTtlProps(entityName, ttlConfig, customTtlSeconds) {
1520
1602
  return { [getTtlAttributeName(entityName, ttlConfig)]: ttlValue };
1521
1603
  }
1522
1604
  //#endregion
1605
+ //#region src/storage/domains/run-fencing.ts
1606
+ const MAX_CONFLICT_RETRIES = 8;
1607
+ const BASE_DELAY_MS$1 = 20;
1608
+ async function backoff(attempt) {
1609
+ const delay = BASE_DELAY_MS$1 * 2 ** attempt;
1610
+ await new Promise((resolve) => setTimeout(resolve, delay + Math.random() * delay * .5));
1611
+ }
1612
+ function hasErrorName(error, errorName) {
1613
+ let current = error;
1614
+ for (let depth = 0; depth < 4 && current && typeof current === "object"; depth++) {
1615
+ const { name, code } = current;
1616
+ if (name === errorName || code === errorName) return true;
1617
+ current = current.cause;
1618
+ }
1619
+ return false;
1620
+ }
1621
+ function isTransactionConflict(error) {
1622
+ return hasErrorName(error, "TransactionConflictException");
1623
+ }
1624
+ /** Whether a write failed its condition, either directly or wrapped by ElectroDB. */
1625
+ function isConditionalCheckFailed(error) {
1626
+ return hasErrorName(error, "ConditionalCheckFailedException");
1627
+ }
1628
+ /**
1629
+ * Runs a single-item write, retrying while a transaction is writing the same
1630
+ * item. DynamoDB rejects such a write with TransactionConflictException
1631
+ * instead of waiting for the transaction, and the SDK does not retry it.
1632
+ */
1633
+ async function retryOnTransactionConflict(write) {
1634
+ for (let attempt = 0;; attempt++) try {
1635
+ return await write();
1636
+ } catch (error) {
1637
+ if (attempt >= MAX_CONFLICT_RETRIES || !isTransactionConflict(error)) throw error;
1638
+ await backoff(attempt);
1639
+ }
1640
+ }
1641
+ /**
1642
+ * Rejects a fenced write that turns out to have nothing to write when its
1643
+ * fence is no longer current, as the write itself would have been rejected.
1644
+ */
1645
+ async function assertRunFence(service, check) {
1646
+ if (!check) return;
1647
+ const { claims, fence, operation } = check;
1648
+ const { data } = await service.entities[claims].get({
1649
+ entity: claims,
1650
+ run_id: fence.runId
1651
+ }).go({ consistent: true });
1652
+ if (!matchesRunFence(data, fence)) throw new RunFenceConflictError(fence, operation);
1653
+ }
1654
+ /**
1655
+ * Runs a write, fenced when `check` is set. A fenced write is a transaction
1656
+ * whose first item is a condition check on the run's claim, matching it on
1657
+ * the fence. A fence that is no longer current fails the check, and the write
1658
+ * throws `RunFenceConflictError` without writing anything. Transactions are
1659
+ * serializable with the single-item writes that claim a run, so a takeover
1660
+ * lands either before the transaction (failing its check) or after it.
1661
+ *
1662
+ * A fenced write whose own condition fails throws an error named
1663
+ * `ConditionalCheckFailedException`, as the unfenced write would.
1664
+ */
1665
+ async function goWithRunFence(service, check, write) {
1666
+ if (!check) {
1667
+ await write.go();
1668
+ return;
1669
+ }
1670
+ const { claims, fence, operation } = check;
1671
+ for (let attempt = 0;; attempt++) {
1672
+ const result = await service.transaction.write((entities) => [entities[claims].check({
1673
+ entity: claims,
1674
+ run_id: fence.runId
1675
+ }).where((attr, op) => `${op.eq(attr.generation, fence.generation)} AND ${op.eq(attr.ownerId, fence.ownerId)}`).commit(), write.commit()]).go();
1676
+ if (!result.canceled) return;
1677
+ const [claim, target] = result.data;
1678
+ if (claim?.code === "ConditionalCheckFailed") throw new RunFenceConflictError(fence, operation);
1679
+ if (target?.code === "ConditionalCheckFailed") {
1680
+ const error = new Error(target.message ?? "The conditional request failed");
1681
+ error.name = "ConditionalCheckFailedException";
1682
+ throw error;
1683
+ }
1684
+ const reasons = result.data.map((item) => item?.code ?? "None");
1685
+ if (attempt >= MAX_CONFLICT_RETRIES || !reasons.includes("TransactionConflict")) throw new Error(`DynamoDB canceled the fenced ${operation} transaction: ${reasons.join(", ")}`);
1686
+ await backoff(attempt);
1687
+ }
1688
+ }
1689
+ //#endregion
1523
1690
  //#region src/storage/domains/memory/index.ts
1524
1691
  var MemoryStorageDynamoDB = class extends MemoryStorage {
1525
1692
  supportsPartialThreadUpdate = true;
@@ -1531,15 +1698,81 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
1531
1698
  this.service = resolved.service;
1532
1699
  this.ttlConfig = resolved.ttl;
1533
1700
  }
1701
+ supportsRunFencing() {
1702
+ return true;
1703
+ }
1534
1704
  async dangerouslyClearAll() {
1535
1705
  await deleteTableData(this.service, TABLE_THREADS);
1536
1706
  await deleteTableData(this.service, TABLE_MESSAGES);
1537
1707
  await deleteTableData(this.service, TABLE_RESOURCES);
1708
+ await deleteTableData(this.service, TABLE_MEMORY_RUN_FENCES);
1538
1709
  }
1539
- async deleteMessages(messageIds) {
1710
+ async raiseRunFence(fence) {
1711
+ const fences = this.service.entities.memory_run_fence;
1712
+ const key = {
1713
+ entity: "memory_run_fence",
1714
+ run_id: fence.runId
1715
+ };
1716
+ try {
1717
+ await retryOnTransactionConflict(() => fences.update(key).set({
1718
+ generation: fence.generation,
1719
+ ownerId: fence.ownerId
1720
+ }).remove(["retiredAt"]).where((attr, op) => `${op.notExists(attr.generation)} OR ${op.lt(attr.generation, fence.generation)} OR (${op.eq(attr.generation, fence.generation)} AND ${op.eq(attr.ownerId, fence.ownerId)})`).go());
1721
+ return true;
1722
+ } catch (error) {
1723
+ try {
1724
+ if (!isConditionalCheckFailed(error)) throw error;
1725
+ const { data } = await fences.get(key).go({ consistent: true });
1726
+ return matchesRunFence(data, fence);
1727
+ } catch (cause) {
1728
+ throw new MastraError({
1729
+ id: createStorageErrorId("DYNAMODB", "RAISE_RUN_FENCE", "FAILED"),
1730
+ domain: ErrorDomain.STORAGE,
1731
+ category: ErrorCategory.THIRD_PARTY,
1732
+ details: { runId: fence.runId }
1733
+ }, cause);
1734
+ }
1735
+ }
1736
+ }
1737
+ async retireRunFence(fence) {
1738
+ try {
1739
+ await retryOnTransactionConflict(() => this.service.entities.memory_run_fence.patch({
1740
+ entity: "memory_run_fence",
1741
+ run_id: fence.runId
1742
+ }).set({ retiredAt: Date.now() }).where((attr, op) => `${op.eq(attr.generation, fence.generation)} AND ${op.eq(attr.ownerId, fence.ownerId)}`).go());
1743
+ return true;
1744
+ } catch (error) {
1745
+ if (isConditionalCheckFailed(error)) return false;
1746
+ throw new MastraError({
1747
+ id: createStorageErrorId("DYNAMODB", "RETIRE_RUN_FENCE", "FAILED"),
1748
+ domain: ErrorDomain.STORAGE,
1749
+ category: ErrorCategory.THIRD_PARTY,
1750
+ details: { runId: fence.runId }
1751
+ }, error);
1752
+ }
1753
+ }
1754
+ /** The check a write must pass: its own fence, otherwise the one in scope. */
1755
+ #runFenceCheck(fence, operation) {
1756
+ const resolved = resolveRunFence(this, fence);
1757
+ return resolved && {
1758
+ claims: "memory_run_fence",
1759
+ fence: resolved,
1760
+ operation
1761
+ };
1762
+ }
1763
+ /** Bumps a thread's updatedAt after a write to its messages. */
1764
+ #touchThread(threadId, check) {
1765
+ return goWithRunFence(this.service, check, this.service.entities.thread.update({
1766
+ entity: "thread",
1767
+ id: threadId
1768
+ }).set({ updatedAt: (/* @__PURE__ */ new Date()).toISOString() }));
1769
+ }
1770
+ async deleteMessages(messageIds, options) {
1540
1771
  if (!messageIds || messageIds.length === 0) return;
1541
1772
  this.logger.debug("Deleting messages", { count: messageIds.length });
1773
+ const check = this.#runFenceCheck(options?.fence, "deleteMessages");
1542
1774
  try {
1775
+ let deleted = 0;
1543
1776
  const threadIds = /* @__PURE__ */ new Set();
1544
1777
  const batchSize = 25;
1545
1778
  for (let i = 0; i < messageIds.length; i += batchSize) {
@@ -1552,18 +1785,17 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
1552
1785
  }));
1553
1786
  for (const message of messagesToDelete) if (message) {
1554
1787
  if (message.threadId) threadIds.add(message.threadId);
1555
- await this.service.entities.message.delete({
1788
+ await goWithRunFence(this.service, check, this.service.entities.message.delete({
1556
1789
  entity: "message",
1557
1790
  id: message.id
1558
- }).go();
1791
+ }));
1792
+ deleted++;
1559
1793
  }
1560
1794
  }
1561
- const now = (/* @__PURE__ */ new Date()).toISOString();
1562
- for (const threadId of threadIds) await this.service.entities.thread.update({
1563
- entity: "thread",
1564
- id: threadId
1565
- }).set({ updatedAt: now }).go();
1795
+ if (deleted === 0) await assertRunFence(this.service, check);
1796
+ for (const threadId of threadIds) await this.#touchThread(threadId, check);
1566
1797
  } catch (error) {
1798
+ if (isRunFenceConflictError(error)) throw error;
1567
1799
  throw new MastraError({
1568
1800
  id: createStorageErrorId("DYNAMODB", "DELETE_MESSAGES", "FAILED"),
1569
1801
  domain: ErrorDomain.STORAGE,
@@ -1617,7 +1849,7 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
1617
1849
  }, error);
1618
1850
  }
1619
1851
  }
1620
- async saveThread({ thread }) {
1852
+ async saveThread({ thread, fence }) {
1621
1853
  this.logger.debug("Saving thread", { threadId: thread.id });
1622
1854
  const now = /* @__PURE__ */ new Date();
1623
1855
  const threadData = {
@@ -1631,7 +1863,7 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
1631
1863
  ...getTtlProps("thread", this.ttlConfig)
1632
1864
  };
1633
1865
  try {
1634
- await this.service.entities.thread.upsert(threadData).go();
1866
+ await goWithRunFence(this.service, this.#runFenceCheck(fence, "saveThread"), this.service.entities.thread.upsert(threadData));
1635
1867
  return {
1636
1868
  id: thread.id,
1637
1869
  resourceId: thread.resourceId,
@@ -1641,6 +1873,7 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
1641
1873
  metadata: thread.metadata
1642
1874
  };
1643
1875
  } catch (error) {
1876
+ if (isRunFenceConflictError(error)) throw error;
1644
1877
  throw new MastraError({
1645
1878
  id: createStorageErrorId("DYNAMODB", "SAVE_THREAD", "FAILED"),
1646
1879
  domain: ErrorDomain.STORAGE,
@@ -1649,8 +1882,9 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
1649
1882
  }, error);
1650
1883
  }
1651
1884
  }
1652
- async updateThread({ id, title, metadata }) {
1885
+ async updateThread({ id, title, metadata, fence }) {
1653
1886
  this.logger.debug("Updating thread", { threadId: id });
1887
+ const check = this.#runFenceCheck(fence, "updateThread");
1654
1888
  try {
1655
1889
  const existingThread = await this.getThreadById({ threadId: id });
1656
1890
  if (!existingThread) throw new Error(`Thread not found: ${id}`);
@@ -1664,10 +1898,10 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
1664
1898
  };
1665
1899
  updateData.metadata = JSON.stringify(mergedMetadata);
1666
1900
  }
1667
- await this.service.entities.thread.update({
1901
+ await goWithRunFence(this.service, check, this.service.entities.thread.update({
1668
1902
  entity: "thread",
1669
1903
  id
1670
- }).set(updateData).go();
1904
+ }).set(updateData));
1671
1905
  return {
1672
1906
  ...existingThread,
1673
1907
  title: title || existingThread.title,
@@ -1678,6 +1912,7 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
1678
1912
  updatedAt: now
1679
1913
  };
1680
1914
  } catch (error) {
1915
+ if (isRunFenceConflictError(error)) throw error;
1681
1916
  throw new MastraError({
1682
1917
  id: createStorageErrorId("DYNAMODB", "UPDATE_THREAD", "FAILED"),
1683
1918
  domain: ErrorDomain.STORAGE,
@@ -1925,6 +2160,7 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
1925
2160
  ...getTtlProps("message", this.ttlConfig)
1926
2161
  };
1927
2162
  });
2163
+ const check = this.#runFenceCheck(args.fence, "saveMessages");
1928
2164
  try {
1929
2165
  const savedMessageIds = [];
1930
2166
  for (const messageData of messagesToSave) {
@@ -1933,14 +2169,14 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
1933
2169
  throw new Error("Internal error: Missing entity property during saveMessages");
1934
2170
  }
1935
2171
  try {
1936
- await this.service.entities.message.put(messageData).go();
2172
+ await goWithRunFence(this.service, check, this.service.entities.message.put(messageData));
1937
2173
  savedMessageIds.push(messageData.id);
1938
2174
  } catch (error) {
1939
- for (const savedId of savedMessageIds) try {
1940
- await this.service.entities.message.delete({
2175
+ for (const savedId of isRunFenceConflictError(error) ? [] : savedMessageIds) try {
2176
+ await goWithRunFence(this.service, check, this.service.entities.message.delete({
1941
2177
  entity: "message",
1942
2178
  id: savedId
1943
- }).go();
2179
+ }));
1944
2180
  } catch (rollbackError) {
1945
2181
  this.logger.error("Failed to rollback message during save error", {
1946
2182
  messageId: savedId,
@@ -1950,12 +2186,10 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
1950
2186
  throw error;
1951
2187
  }
1952
2188
  }
1953
- await this.service.entities.thread.update({
1954
- entity: "thread",
1955
- id: threadId
1956
- }).set({ updatedAt: (/* @__PURE__ */ new Date()).toISOString() }).go();
2189
+ await this.#touchThread(threadId, check);
1957
2190
  return { messages: new MessageList().add(messages, "memory").get.all.db() };
1958
2191
  } catch (error) {
2192
+ if (isRunFenceConflictError(error)) throw error;
1959
2193
  throw new MastraError({
1960
2194
  id: createStorageErrorId("DYNAMODB", "SAVE_MESSAGES", "FAILED"),
1961
2195
  domain: ErrorDomain.STORAGE,
@@ -2109,6 +2343,7 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
2109
2343
  const { messages } = args;
2110
2344
  this.logger.debug("Updating messages", { count: messages.length });
2111
2345
  if (!messages.length) return [];
2346
+ const check = this.#runFenceCheck(args.fence, "updateMessages");
2112
2347
  const updatedMessages = [];
2113
2348
  const affectedThreadIds = /* @__PURE__ */ new Set();
2114
2349
  try {
@@ -2144,22 +2379,21 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
2144
2379
  if ("parts" in updates.content && updates.content.parts !== void 0) newContent.parts = updates.content.parts;
2145
2380
  updatePayload.content = JSON.stringify(newContent);
2146
2381
  }
2147
- await this.service.entities.message.update({
2382
+ await goWithRunFence(this.service, check, this.service.entities.message.update({
2148
2383
  entity: "message",
2149
2384
  id
2150
- }).set(updatePayload).go();
2385
+ }).set(updatePayload));
2151
2386
  const updatedMessage = await this.service.entities.message.get({
2152
2387
  entity: "message",
2153
2388
  id
2154
2389
  }).go();
2155
2390
  if (updatedMessage.data) updatedMessages.push(this.parseMessageData(updatedMessage.data));
2156
2391
  }
2157
- for (const threadId of affectedThreadIds) await this.service.entities.thread.update({
2158
- entity: "thread",
2159
- id: threadId
2160
- }).set({ updatedAt: (/* @__PURE__ */ new Date()).toISOString() }).go();
2392
+ if (affectedThreadIds.size === 0) await assertRunFence(this.service, check);
2393
+ for (const threadId of affectedThreadIds) await this.#touchThread(threadId, check);
2161
2394
  return updatedMessages;
2162
2395
  } catch (error) {
2396
+ if (isRunFenceConflictError(error)) throw error;
2163
2397
  throw new MastraError({
2164
2398
  id: createStorageErrorId("DYNAMODB", "UPDATE_MESSAGES", "FAILED"),
2165
2399
  domain: ErrorDomain.STORAGE,
@@ -2193,6 +2427,9 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
2193
2427
  }
2194
2428
  }
2195
2429
  async saveResource({ resource }) {
2430
+ return this.#saveResource(resource, void 0);
2431
+ }
2432
+ async #saveResource(resource, check) {
2196
2433
  this.logger.debug("Saving resource", { resourceId: resource.id });
2197
2434
  const now = /* @__PURE__ */ new Date();
2198
2435
  const resourceData = {
@@ -2205,7 +2442,7 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
2205
2442
  ...getTtlProps("resource", this.ttlConfig)
2206
2443
  };
2207
2444
  try {
2208
- await this.service.entities.resource.upsert(resourceData).go();
2445
+ await goWithRunFence(this.service, check, this.service.entities.resource.upsert(resourceData));
2209
2446
  return {
2210
2447
  id: resource.id,
2211
2448
  workingMemory: resource.workingMemory,
@@ -2214,6 +2451,7 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
2214
2451
  updatedAt: now
2215
2452
  };
2216
2453
  } catch (error) {
2454
+ if (isRunFenceConflictError(error)) throw error;
2217
2455
  throw new MastraError({
2218
2456
  id: createStorageErrorId("DYNAMODB", "SAVE_RESOURCE", "FAILED"),
2219
2457
  domain: ErrorDomain.STORAGE,
@@ -2222,8 +2460,9 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
2222
2460
  }, error);
2223
2461
  }
2224
2462
  }
2225
- async updateResource({ resourceId, workingMemory, metadata }) {
2463
+ async updateResource({ resourceId, workingMemory, metadata, fence }) {
2226
2464
  this.logger.debug("Updating resource", { resourceId });
2465
+ const check = this.#runFenceCheck(fence, "updateResource");
2227
2466
  try {
2228
2467
  const existingResource = await this.getResourceById({ resourceId });
2229
2468
  if (!existingResource) {
@@ -2234,7 +2473,7 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
2234
2473
  createdAt: /* @__PURE__ */ new Date(),
2235
2474
  updatedAt: /* @__PURE__ */ new Date()
2236
2475
  };
2237
- return this.saveResource({ resource: newResource });
2476
+ return this.#saveResource(newResource, check);
2238
2477
  }
2239
2478
  const now = /* @__PURE__ */ new Date();
2240
2479
  const updateData = { updatedAt: now.toISOString() };
@@ -2246,10 +2485,10 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
2246
2485
  };
2247
2486
  updateData.metadata = JSON.stringify(mergedMetadata);
2248
2487
  }
2249
- await this.service.entities.resource.update({
2488
+ await goWithRunFence(this.service, check, this.service.entities.resource.update({
2250
2489
  entity: "resource",
2251
2490
  id: resourceId
2252
- }).set(updateData).go();
2491
+ }).set(updateData));
2253
2492
  return {
2254
2493
  ...existingResource,
2255
2494
  workingMemory: workingMemory !== void 0 ? workingMemory : existingResource.workingMemory,
@@ -2260,6 +2499,7 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
2260
2499
  updatedAt: now
2261
2500
  };
2262
2501
  } catch (error) {
2502
+ if (isRunFenceConflictError(error)) throw error;
2263
2503
  throw new MastraError({
2264
2504
  id: createStorageErrorId("DYNAMODB", "UPDATE_RESOURCE", "FAILED"),
2265
2505
  domain: ErrorDomain.STORAGE,
@@ -2581,6 +2821,15 @@ function formatWorkflowRun(snapshotData) {
2581
2821
  resourceId: snapshotData.resourceId
2582
2822
  };
2583
2823
  }
2824
+ function toRunOwnershipRecord(item) {
2825
+ return {
2826
+ runId: item.run_id,
2827
+ generation: Number(item.generation),
2828
+ ownerId: item.ownerId,
2829
+ leaseExpiresAt: item.leaseExpiresAt === void 0 ? null : new Date(item.leaseExpiresAt),
2830
+ live: item.leaseExpiresAt !== void 0 && item.leaseExpiresAt > Date.now()
2831
+ };
2832
+ }
2584
2833
  const MAX_RETRIES = 5;
2585
2834
  const BASE_DELAY_MS = 50;
2586
2835
  var WorkflowStorageDynamoDB = class extends WorkflowsStorage {
@@ -2595,8 +2844,126 @@ var WorkflowStorageDynamoDB = class extends WorkflowsStorage {
2595
2844
  supportsConcurrentUpdates() {
2596
2845
  return true;
2597
2846
  }
2847
+ supportsRunFencing() {
2848
+ return true;
2849
+ }
2598
2850
  async dangerouslyClearAll() {
2599
2851
  await deleteTableData(this.service, TABLE_WORKFLOW_SNAPSHOT);
2852
+ await deleteTableData(this.service, TABLE_WORKFLOW_RUN_OWNERS);
2853
+ }
2854
+ #runFenceCheck(fence, runId, operation) {
2855
+ const resolved = resolveRunFence(this, fence, runId);
2856
+ return resolved && {
2857
+ claims: "workflow_run_owner",
2858
+ fence: resolved,
2859
+ operation
2860
+ };
2861
+ }
2862
+ #runOwnerKey(runId) {
2863
+ return {
2864
+ entity: "workflow_run_owner",
2865
+ run_id: runId
2866
+ };
2867
+ }
2868
+ async #readRunOwner(runId) {
2869
+ const { data } = await this.service.entities.workflow_run_owner.get(this.#runOwnerKey(runId)).go({ consistent: true });
2870
+ return data ? toRunOwnershipRecord(data) : null;
2871
+ }
2872
+ #ownershipError(operation, runId, error) {
2873
+ return new MastraError({
2874
+ id: createStorageErrorId("DYNAMODB", operation, "FAILED"),
2875
+ domain: ErrorDomain.STORAGE,
2876
+ category: ErrorCategory.THIRD_PARTY,
2877
+ details: { runId }
2878
+ }, error);
2879
+ }
2880
+ async claimRunOwnership({ runId, ownerId, leaseMs, force, expectedGeneration }) {
2881
+ try {
2882
+ const owners = this.service.entities.workflow_run_owner;
2883
+ const now = Date.now();
2884
+ if (!expectedGeneration) {
2885
+ try {
2886
+ const item = {
2887
+ ...this.#runOwnerKey(runId),
2888
+ generation: Math.max(1, now),
2889
+ ownerId,
2890
+ leaseExpiresAt: now + leaseMs,
2891
+ updatedAt: now
2892
+ };
2893
+ await retryOnTransactionConflict(() => owners.create(item).go());
2894
+ return {
2895
+ acquired: true,
2896
+ record: toRunOwnershipRecord(item)
2897
+ };
2898
+ } catch (error) {
2899
+ if (!this.isConditionalCheckFailed(error)) throw error;
2900
+ }
2901
+ if (expectedGeneration === 0) return {
2902
+ acquired: false,
2903
+ record: await this.#readRunOwner(runId)
2904
+ };
2905
+ }
2906
+ try {
2907
+ const { data } = await retryOnTransactionConflict(() => owners.patch(this.#runOwnerKey(runId)).set({
2908
+ ownerId,
2909
+ leaseExpiresAt: now + leaseMs,
2910
+ updatedAt: now
2911
+ }).add({ generation: 1 }).where((attr, op) => [expectedGeneration !== void 0 ? op.eq(attr.generation, expectedGeneration) : "", force ? "" : `(${op.notExists(attr.leaseExpiresAt)} OR ${op.lte(attr.leaseExpiresAt, now)})`].filter(Boolean).join(" AND ")).go({ response: "all_new" }));
2912
+ return {
2913
+ acquired: true,
2914
+ record: toRunOwnershipRecord(data)
2915
+ };
2916
+ } catch (error) {
2917
+ if (!this.isConditionalCheckFailed(error)) throw error;
2918
+ return {
2919
+ acquired: false,
2920
+ record: await this.#readRunOwner(runId)
2921
+ };
2922
+ }
2923
+ } catch (error) {
2924
+ throw this.#ownershipError("CLAIM_RUN_OWNERSHIP", runId, error);
2925
+ }
2926
+ }
2927
+ async renewRunOwnership({ leaseMs, ...fence }) {
2928
+ try {
2929
+ const { data } = await retryOnTransactionConflict(() => {
2930
+ const now = Date.now();
2931
+ return this.service.entities.workflow_run_owner.patch(this.#runOwnerKey(fence.runId)).set({
2932
+ leaseExpiresAt: now + leaseMs,
2933
+ updatedAt: now
2934
+ }).where((attr, op) => `${op.eq(attr.generation, fence.generation)} AND ${op.eq(attr.ownerId, fence.ownerId)} AND ${op.exists(attr.leaseExpiresAt)}`).go({ response: "all_new" });
2935
+ });
2936
+ return {
2937
+ renewed: true,
2938
+ record: toRunOwnershipRecord(data)
2939
+ };
2940
+ } catch (error) {
2941
+ if (!this.isConditionalCheckFailed(error)) throw this.#ownershipError("RENEW_RUN_OWNERSHIP", fence.runId, error);
2942
+ try {
2943
+ return {
2944
+ renewed: false,
2945
+ record: await this.#readRunOwner(fence.runId)
2946
+ };
2947
+ } catch (readError) {
2948
+ throw this.#ownershipError("RENEW_RUN_OWNERSHIP", fence.runId, readError);
2949
+ }
2950
+ }
2951
+ }
2952
+ async releaseRunOwnership(fence) {
2953
+ try {
2954
+ await retryOnTransactionConflict(() => this.service.entities.workflow_run_owner.patch(this.#runOwnerKey(fence.runId)).set({ updatedAt: Date.now() }).remove(["leaseExpiresAt"]).where((attr, op) => `${op.eq(attr.generation, fence.generation)} AND ${op.eq(attr.ownerId, fence.ownerId)}`).go());
2955
+ return true;
2956
+ } catch (error) {
2957
+ if (this.isConditionalCheckFailed(error)) return false;
2958
+ throw this.#ownershipError("RELEASE_RUN_OWNERSHIP", fence.runId, error);
2959
+ }
2960
+ }
2961
+ async getRunOwnership({ runId }) {
2962
+ try {
2963
+ return await this.#readRunOwner(runId);
2964
+ } catch (error) {
2965
+ throw this.#ownershipError("GET_RUN_OWNERSHIP", runId, error);
2966
+ }
2600
2967
  }
2601
2968
  /**
2602
2969
  * Helper to check if an error is a DynamoDB conditional check failure
@@ -2621,7 +2988,8 @@ var WorkflowStorageDynamoDB = class extends WorkflowsStorage {
2621
2988
  const jitter = Math.random() * backoff * .5;
2622
2989
  await new Promise((resolve) => setTimeout(resolve, backoff + jitter));
2623
2990
  }
2624
- async updateWorkflowResults({ workflowName, runId, stepId, result, requestContext, state }) {
2991
+ async updateWorkflowResults({ workflowName, runId, stepId, result, requestContext, state, fence }) {
2992
+ const check = this.#runFenceCheck(fence, runId, "updateWorkflowResults");
2625
2993
  for (let attempt = 0; attempt < MAX_RETRIES; attempt++) try {
2626
2994
  const existingRecord = await this.service.entities.workflow_snapshot.get({
2627
2995
  entity: "workflow_snapshot",
@@ -2664,8 +3032,8 @@ var WorkflowStorageDynamoDB = class extends WorkflowsStorage {
2664
3032
  updatedAt: now.toISOString(),
2665
3033
  ...getTtlProps("workflow_snapshot", this.ttlConfig)
2666
3034
  };
2667
- if (previousUpdatedAt) await this.service.entities.workflow_snapshot.upsert(data).where((attr, op) => op.eq(attr.updatedAt, previousUpdatedAt)).go();
2668
- else await this.service.entities.workflow_snapshot.create(data).where((attr, op) => op.notExists(attr.run_id)).go();
3035
+ const write = previousUpdatedAt ? this.service.entities.workflow_snapshot.upsert(data).where((attr, op) => op.eq(attr.updatedAt, previousUpdatedAt)) : this.service.entities.workflow_snapshot.create(data).where((attr, op) => op.notExists(attr.run_id));
3036
+ await goWithRunFence(this.service, check, write);
2669
3037
  return snapshot.context;
2670
3038
  } catch (error) {
2671
3039
  if (this.isConditionalCheckFailed(error)) {
@@ -2698,19 +3066,21 @@ var WorkflowStorageDynamoDB = class extends WorkflowsStorage {
2698
3066
  }
2699
3067
  }, /* @__PURE__ */ new Error("Max retries exceeded for optimistic locking"));
2700
3068
  }
2701
- async updateWorkflowState({ workflowName, runId, opts }) {
3069
+ async updateWorkflowState({ workflowName, runId, opts, fence }) {
3070
+ const check = this.#runFenceCheck(fence, runId, "updateWorkflowState");
2702
3071
  for (let attempt = 0; attempt < MAX_RETRIES; attempt++) try {
2703
3072
  const existingRecord = await this.service.entities.workflow_snapshot.get({
2704
3073
  entity: "workflow_snapshot",
2705
3074
  workflow_name: workflowName,
2706
3075
  run_id: runId
2707
3076
  }).go();
2708
- if (!existingRecord.data) return;
2709
- const existingSnapshot = existingRecord.data.snapshot;
2710
- if (!existingSnapshot || !existingSnapshot.context) return;
2711
- const previousUpdatedAt = existingRecord.data.updatedAt;
3077
+ const existingSnapshot = existingRecord.data?.snapshot;
2712
3078
  const { expectedStatus, ...state } = opts;
2713
- if (!matchesExpectedWorkflowStatus(existingSnapshot.status, expectedStatus)) return;
3079
+ if (!existingSnapshot?.context || !matchesExpectedWorkflowStatus(existingSnapshot.status, expectedStatus)) {
3080
+ await assertRunFence(this.service, check);
3081
+ return;
3082
+ }
3083
+ const previousUpdatedAt = existingRecord.data.updatedAt;
2714
3084
  const updatedSnapshot = {
2715
3085
  ...existingSnapshot,
2716
3086
  ...state
@@ -2726,7 +3096,7 @@ var WorkflowStorageDynamoDB = class extends WorkflowsStorage {
2726
3096
  resourceId: existingRecord.data.resourceId,
2727
3097
  ...getTtlProps("workflow_snapshot", this.ttlConfig)
2728
3098
  };
2729
- await this.service.entities.workflow_snapshot.upsert(data).where((attr, op) => op.eq(attr.updatedAt, previousUpdatedAt)).go();
3099
+ await goWithRunFence(this.service, check, this.service.entities.workflow_snapshot.upsert(data).where((attr, op) => op.eq(attr.updatedAt, previousUpdatedAt)));
2730
3100
  return updatedSnapshot;
2731
3101
  } catch (error) {
2732
3102
  if (this.isConditionalCheckFailed(error)) {
@@ -2757,11 +3127,12 @@ var WorkflowStorageDynamoDB = class extends WorkflowsStorage {
2757
3127
  }
2758
3128
  }, /* @__PURE__ */ new Error("Max retries exceeded for optimistic locking"));
2759
3129
  }
2760
- async persistWorkflowSnapshot({ workflowName, runId, resourceId, snapshot, createdAt, updatedAt }) {
3130
+ async persistWorkflowSnapshot({ workflowName, runId, resourceId, snapshot, createdAt, updatedAt, fence }) {
2761
3131
  this.logger.debug("Persisting workflow snapshot", {
2762
3132
  workflowName,
2763
3133
  runId
2764
3134
  });
3135
+ const check = this.#runFenceCheck(fence, runId, "persistWorkflowSnapshot");
2765
3136
  try {
2766
3137
  const now = /* @__PURE__ */ new Date();
2767
3138
  const createdAtValue = (createdAt ?? now).toISOString();
@@ -2777,10 +3148,10 @@ var WorkflowStorageDynamoDB = class extends WorkflowsStorage {
2777
3148
  ...getTtlProps("workflow_snapshot", this.ttlConfig)
2778
3149
  };
2779
3150
  try {
2780
- await this.service.entities.workflow_snapshot.create(data).where((attr, op) => op.notExists(attr.run_id)).go();
3151
+ await goWithRunFence(this.service, check, this.service.entities.workflow_snapshot.create(data).where((attr, op) => op.notExists(attr.run_id)));
2781
3152
  } catch (error) {
2782
3153
  if (!this.isConditionalCheckFailed(error)) throw error;
2783
- await this.service.entities.workflow_snapshot.update({
3154
+ await goWithRunFence(this.service, check, this.service.entities.workflow_snapshot.update({
2784
3155
  entity: "workflow_snapshot",
2785
3156
  workflow_name: workflowName,
2786
3157
  run_id: runId
@@ -2789,9 +3160,10 @@ var WorkflowStorageDynamoDB = class extends WorkflowsStorage {
2789
3160
  updatedAt: updatedAtValue,
2790
3161
  resourceId,
2791
3162
  ...getTtlProps("workflow_snapshot", this.ttlConfig)
2792
- }).go();
3163
+ }));
2793
3164
  }
2794
3165
  } catch (error) {
3166
+ if (isRunFenceConflictError(error)) throw error;
2795
3167
  throw new MastraError({
2796
3168
  id: createStorageErrorId("DYNAMODB", "PERSIST_WORKFLOW_SNAPSHOT", "FAILED"),
2797
3169
  domain: ErrorDomain.STORAGE,
@@ -2950,18 +3322,19 @@ var WorkflowStorageDynamoDB = class extends WorkflowsStorage {
2950
3322
  }, error);
2951
3323
  }
2952
3324
  }
2953
- async deleteWorkflowRunById({ runId, workflowName }) {
3325
+ async deleteWorkflowRunById({ runId, workflowName, fence }) {
2954
3326
  this.logger.debug("Deleting workflow run by ID", {
2955
3327
  runId,
2956
3328
  workflowName
2957
3329
  });
2958
3330
  try {
2959
- await this.service.entities.workflow_snapshot.delete({
3331
+ await goWithRunFence(this.service, this.#runFenceCheck(fence, runId, "deleteWorkflowRunById"), this.service.entities.workflow_snapshot.delete({
2960
3332
  entity: "workflow_snapshot",
2961
3333
  workflow_name: workflowName,
2962
3334
  run_id: runId
2963
- }).go();
3335
+ }));
2964
3336
  } catch (error) {
3337
+ if (isRunFenceConflictError(error)) throw error;
2965
3338
  throw new MastraError({
2966
3339
  id: createStorageErrorId("DYNAMODB", "DELETE_WORKFLOW_RUN_BY_ID", "FAILED"),
2967
3340
  domain: ErrorDomain.STORAGE,