@mastra/upstash 1.5.2-alpha.0 → 1.5.3

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 { RedisServerCache, RedisServerCache as RedisServerCache$1, upstashPreset, upstashPreset as upstashPreset$1 } from "@mastra/redis";
2
2
  import { Redis } from "@upstash/redis";
3
- import { BackgroundTasksStorage, MastraCompositeStore, MemoryStorage, ScoresStorage, TABLE_BACKGROUND_TASKS, TABLE_MESSAGES, TABLE_RESOURCES, TABLE_SCORERS, TABLE_THREADS, TABLE_WORKFLOW_SNAPSHOT, WorkflowsStorage, calculatePagination, createStorageErrorId, createVectorErrorId, ensureDate, filterByDateRange, normalizePerPage, serializeDate, storageMessageMatchesMetadataFilter, transformScoreRow, validateStorageMetadataFilter } from "@mastra/core/storage";
3
+ import { BackgroundTasksStorage, MastraCompositeStore, MemoryStorage, RunFenceConflictError, 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, createVectorErrorId, ensureDate, filterByDateRange, isRunFenceConflictError, normalizePerPage, resolveRunFence, serializeDate, storageMessageMatchesMetadataFilter, transformScoreRow, validateStorageMetadataFilter } from "@mastra/core/storage";
4
4
  import { ErrorCategory, ErrorDomain, MastraError } from "@mastra/core/error";
5
5
  import { MessageList } from "@mastra/core/agent";
6
6
  import { saveScorePayloadSchema } from "@mastra/core/evals";
@@ -350,6 +350,239 @@ var BackgroundTasksUpstash = class extends BackgroundTasksStorage {
350
350
  }
351
351
  };
352
352
  //#endregion
353
+ //#region src/storage/domains/run-fencing.ts
354
+ function runClaimKey(claims, runId) {
355
+ return `${claims}:run_id:${runId}`;
356
+ }
357
+ const RUN_FENCE_CONFLICT = "MASTRA_RUN_FENCE_CONFLICT";
358
+ const RUN_CLAIM_RETENTION_MS = 10080 * 60 * 1e3;
359
+ const FENCE_GUARD = `
360
+ local fenceClaim = redis.call('HMGET', KEYS[#KEYS], 'generation', 'ownerId')
361
+ if tonumber(fenceClaim[1]) ~= tonumber(ARGV[#ARGV - 1]) or fenceClaim[2] ~= ARGV[#ARGV] then
362
+ return '${RUN_FENCE_CONFLICT}'
363
+ end
364
+ `;
365
+ const OWNER = `
366
+ local function owner(flag)
367
+ local r = redis.call('HMGET', KEYS[1], 'generation', 'ownerId', 'leaseExpiresAt')
368
+ return cjson.encode({ flag, r[1], r[2], r[3] })
369
+ end
370
+ `;
371
+ const CLAIM = `${OWNER}
372
+ local generation = tonumber(redis.call('HGET', KEYS[1], 'generation'))
373
+ local expected = ARGV[3]
374
+ if not generation then
375
+ if expected ~= '' and expected ~= '0' then return owner(0) end
376
+ -- The first generation comes from the clock, so a claim after this hash is
377
+ -- removed still outranks a surviving fence.
378
+ redis.call('HSET', KEYS[1], 'generation', string.format('%d', math.max(1, tonumber(ARGV[4]))), 'ownerId', ARGV[1], 'leaseExpiresAt', ARGV[5], 'updatedAt', ARGV[4])
379
+ redis.call('PEXPIRE', KEYS[1], ARGV[6])
380
+ return owner(1)
381
+ end
382
+ if expected ~= '' and generation ~= tonumber(expected) then return owner(0) end
383
+ if ARGV[2] ~= '1' then
384
+ local expires = tonumber(redis.call('HGET', KEYS[1], 'leaseExpiresAt'))
385
+ if expires and expires > tonumber(ARGV[4]) then return owner(0) end
386
+ end
387
+ redis.call('HSET', KEYS[1], 'generation', string.format('%d', generation + 1), 'ownerId', ARGV[1], 'leaseExpiresAt', ARGV[5], 'updatedAt', ARGV[4])
388
+ redis.call('PEXPIRE', KEYS[1], ARGV[6])
389
+ return owner(1)
390
+ `;
391
+ const RENEW = `${OWNER}
392
+ local r = redis.call('HMGET', KEYS[1], 'generation', 'ownerId', 'leaseExpiresAt')
393
+ if tonumber(r[1]) ~= tonumber(ARGV[1]) or r[2] ~= ARGV[2] or not r[3] then return owner(0) end
394
+ redis.call('HSET', KEYS[1], 'leaseExpiresAt', ARGV[3], 'updatedAt', ARGV[4])
395
+ redis.call('PEXPIRE', KEYS[1], ARGV[5])
396
+ return owner(1)
397
+ `;
398
+ const RELEASE = `
399
+ local r = redis.call('HMGET', KEYS[1], 'generation', 'ownerId')
400
+ if tonumber(r[1]) ~= tonumber(ARGV[1]) or r[2] ~= ARGV[2] then return 0 end
401
+ redis.call('HDEL', KEYS[1], 'leaseExpiresAt')
402
+ redis.call('HSET', KEYS[1], 'updatedAt', ARGV[3])
403
+ redis.call('PEXPIRE', KEYS[1], ARGV[4])
404
+ return 1
405
+ `;
406
+ const READ = `${OWNER}
407
+ return owner(1)
408
+ `;
409
+ const RAISE = `
410
+ local r = redis.call('HMGET', KEYS[1], 'generation', 'ownerId')
411
+ local current = tonumber(r[1])
412
+ if not current or current < tonumber(ARGV[1]) then
413
+ redis.call('HSET', KEYS[1], 'generation', ARGV[1], 'ownerId', ARGV[2])
414
+ elseif current ~= tonumber(ARGV[1]) or r[2] ~= ARGV[2] then
415
+ return 0
416
+ end
417
+ -- Every write un-retires the fence.
418
+ redis.call('HDEL', KEYS[1], 'retiredAt')
419
+ redis.call('PERSIST', KEYS[1])
420
+ return 1
421
+ `;
422
+ const RETIRE = `
423
+ local r = redis.call('HMGET', KEYS[1], 'generation', 'ownerId')
424
+ if tonumber(r[1]) ~= tonumber(ARGV[1]) or r[2] ~= ARGV[2] then return 0 end
425
+ redis.call('HSET', KEYS[1], 'retiredAt', ARGV[3])
426
+ redis.call('PEXPIRE', KEYS[1], ARGV[4])
427
+ return 1
428
+ `;
429
+ const FENCED_WRITE = `
430
+ local i = 1
431
+ while i <= #ARGV - 2 do
432
+ local n = tonumber(ARGV[i])
433
+ redis.call(unpack(ARGV, i + 1, i + n))
434
+ i = i + n + 1
435
+ end
436
+ return 1
437
+ `;
438
+ /**
439
+ * Runs `script`, first checking the run's claim against `check`'s fence in
440
+ * the same script when `check` is set. A stale fence throws
441
+ * `RunFenceConflictError` before the script does anything.
442
+ */
443
+ async function evalFenced(client, script, keys, args, check) {
444
+ if (!check) return client.eval(script, keys, args);
445
+ const { claims, fence, operation } = check;
446
+ const result = await client.eval(FENCE_GUARD + script, [...keys, runClaimKey(claims, fence.runId)], [
447
+ ...args,
448
+ String(fence.generation),
449
+ fence.ownerId
450
+ ]);
451
+ if (result === RUN_FENCE_CONFLICT) throw new RunFenceConflictError(fence, operation);
452
+ return result;
453
+ }
454
+ function toRunOwnershipRecord(runId, reply, now) {
455
+ const [, generation, ownerId, leaseExpiresAt] = reply;
456
+ if (!generation) return null;
457
+ const expiresAt = leaseExpiresAt ? Number(leaseExpiresAt) : null;
458
+ return {
459
+ runId,
460
+ generation: Number(generation),
461
+ ownerId: String(ownerId),
462
+ leaseExpiresAt: expiresAt === null ? null : new Date(expiresAt),
463
+ live: expiresAt !== null && expiresAt > now
464
+ };
465
+ }
466
+ async function evalOwner(client, script, runId, args, now) {
467
+ const raw = await client.eval(script, [runClaimKey(TABLE_WORKFLOW_RUN_OWNERS, runId)], args);
468
+ const reply = typeof raw === "string" ? JSON.parse(raw) : raw;
469
+ return {
470
+ applied: Number(reply[0]) === 1,
471
+ record: toRunOwnershipRecord(runId, reply, now)
472
+ };
473
+ }
474
+ async function claimRunOwnership(client, { runId, ownerId, leaseMs, force, expectedGeneration }) {
475
+ const now = Date.now();
476
+ const { applied, record } = await evalOwner(client, CLAIM, runId, [
477
+ ownerId,
478
+ force ? "1" : "0",
479
+ expectedGeneration === void 0 ? "" : String(expectedGeneration),
480
+ String(now),
481
+ String(now + leaseMs),
482
+ String(RUN_CLAIM_RETENTION_MS)
483
+ ], now);
484
+ return applied ? {
485
+ acquired: true,
486
+ record
487
+ } : {
488
+ acquired: false,
489
+ record
490
+ };
491
+ }
492
+ async function renewRunOwnership(client, { leaseMs, ...fence }) {
493
+ const now = Date.now();
494
+ const { applied, record } = await evalOwner(client, RENEW, fence.runId, [
495
+ String(fence.generation),
496
+ fence.ownerId,
497
+ String(now + leaseMs),
498
+ String(now),
499
+ String(RUN_CLAIM_RETENTION_MS)
500
+ ], now);
501
+ return applied ? {
502
+ renewed: true,
503
+ record
504
+ } : {
505
+ renewed: false,
506
+ record
507
+ };
508
+ }
509
+ async function releaseRunOwnership(client, fence) {
510
+ const released = await client.eval(RELEASE, [runClaimKey(TABLE_WORKFLOW_RUN_OWNERS, fence.runId)], [
511
+ String(fence.generation),
512
+ fence.ownerId,
513
+ String(Date.now()),
514
+ String(RUN_CLAIM_RETENTION_MS)
515
+ ]);
516
+ return Number(released) === 1;
517
+ }
518
+ async function getRunOwnership(client, runId) {
519
+ return (await evalOwner(client, READ, runId, [], Date.now())).record;
520
+ }
521
+ async function raiseRunFence(client, fence) {
522
+ const raised = await client.eval(RAISE, [runClaimKey(TABLE_MEMORY_RUN_FENCES, fence.runId)], [String(fence.generation), fence.ownerId]);
523
+ return Number(raised) === 1;
524
+ }
525
+ async function retireRunFence(client, fence) {
526
+ const retired = await client.eval(RETIRE, [runClaimKey(TABLE_MEMORY_RUN_FENCES, fence.runId)], [
527
+ String(fence.generation),
528
+ fence.ownerId,
529
+ String(Date.now()),
530
+ String(RUN_CLAIM_RETENTION_MS)
531
+ ]);
532
+ return Number(retired) === 1;
533
+ }
534
+ function serialize(value) {
535
+ return typeof value === "string" ? value : JSON.stringify(value);
536
+ }
537
+ /**
538
+ * Collects writes and applies them in one script that first checks the run's
539
+ * claim against the fence. A fence that is no longer current makes `exec()`
540
+ * throw `RunFenceConflictError` without writing anything. Scripts run
541
+ * atomically, so a takeover lands either before the batch (failing its check)
542
+ * or after it.
543
+ */
544
+ var FencedWriteBatch = class {
545
+ client;
546
+ check;
547
+ #keys = /* @__PURE__ */ new Set();
548
+ #args = [];
549
+ constructor(client, check) {
550
+ this.client = client;
551
+ this.check = check;
552
+ }
553
+ #add(key, ...command) {
554
+ this.#keys.add(key);
555
+ this.#args.push(String(command.length), ...command);
556
+ return this;
557
+ }
558
+ set(key, value) {
559
+ return this.#add(key, "SET", key, serialize(value));
560
+ }
561
+ del(key) {
562
+ return this.#add(key, "DEL", key);
563
+ }
564
+ zadd(key, { score, member }) {
565
+ return this.#add(key, "ZADD", key, String(score), member);
566
+ }
567
+ zrem(key, member) {
568
+ return this.#add(key, "ZREM", key, member);
569
+ }
570
+ async exec() {
571
+ await evalFenced(this.client, FENCED_WRITE, [...this.#keys], this.#args, this.check);
572
+ }
573
+ };
574
+ /** A pipeline, or a fenced batch when `check` is set. */
575
+ function writeBatch(client, check) {
576
+ return check ? new FencedWriteBatch(client, check) : client.pipeline();
577
+ }
578
+ /**
579
+ * Rejects a fenced write that turns out to have nothing to write when its
580
+ * fence is no longer current, as the write itself would have been rejected.
581
+ */
582
+ async function assertRunFence(client, check) {
583
+ if (check) await new FencedWriteBatch(client, check).exec();
584
+ }
585
+ //#endregion
353
586
  //#region src/storage/domains/memory/index.ts
354
587
  function getThreadMessagesKey(threadId) {
355
588
  return `thread:${threadId}:messages`;
@@ -373,10 +606,46 @@ var StoreMemoryUpstash = class extends MemoryStorage {
373
606
  this.client = client;
374
607
  this.#db = new UpstashDB({ client });
375
608
  }
609
+ supportsRunFencing() {
610
+ return true;
611
+ }
612
+ async raiseRunFence(fence) {
613
+ try {
614
+ return await raiseRunFence(this.client, fence);
615
+ } catch (error) {
616
+ throw new MastraError({
617
+ id: createStorageErrorId("UPSTASH", "RAISE_RUN_FENCE", "FAILED"),
618
+ domain: ErrorDomain.STORAGE,
619
+ category: ErrorCategory.THIRD_PARTY,
620
+ details: { runId: fence.runId }
621
+ }, error);
622
+ }
623
+ }
624
+ async retireRunFence(fence) {
625
+ try {
626
+ return await retireRunFence(this.client, fence);
627
+ } catch (error) {
628
+ throw new MastraError({
629
+ id: createStorageErrorId("UPSTASH", "RETIRE_RUN_FENCE", "FAILED"),
630
+ domain: ErrorDomain.STORAGE,
631
+ category: ErrorCategory.THIRD_PARTY,
632
+ details: { runId: fence.runId }
633
+ }, error);
634
+ }
635
+ }
636
+ #runFenceCheck(fence, operation) {
637
+ const resolved = resolveRunFence(this, fence);
638
+ return resolved && {
639
+ claims: TABLE_MEMORY_RUN_FENCES,
640
+ fence: resolved,
641
+ operation
642
+ };
643
+ }
376
644
  async dangerouslyClearAll() {
377
645
  await this.#db.deleteData({ tableName: TABLE_THREADS });
378
646
  await this.#db.deleteData({ tableName: TABLE_MESSAGES });
379
647
  await this.#db.deleteData({ tableName: TABLE_RESOURCES });
648
+ await this.#db.scanAndDelete(`${runClaimKey(TABLE_MEMORY_RUN_FENCES, "")}*`);
380
649
  }
381
650
  async getThreadById({ threadId, resourceId }) {
382
651
  try {
@@ -485,14 +754,16 @@ var StoreMemoryUpstash = class extends MemoryStorage {
485
754
  throw mastraError;
486
755
  }
487
756
  }
488
- async saveThread({ thread }) {
757
+ async saveThread({ thread, fence }) {
758
+ const check = this.#runFenceCheck(fence, "saveThread");
489
759
  try {
490
- await this.#db.insert({
491
- tableName: TABLE_THREADS,
492
- record: thread
493
- });
760
+ const { key, processedRecord } = processRecord(TABLE_THREADS, thread);
761
+ const batch = writeBatch(this.client, check);
762
+ batch.set(key, processedRecord);
763
+ await batch.exec();
494
764
  return thread;
495
765
  } catch (error) {
766
+ if (isRunFenceConflictError(error)) throw error;
496
767
  const mastraError = new MastraError({
497
768
  id: createStorageErrorId("UPSTASH", "SAVE_THREAD", "FAILED"),
498
769
  domain: ErrorDomain.STORAGE,
@@ -504,7 +775,7 @@ var StoreMemoryUpstash = class extends MemoryStorage {
504
775
  throw mastraError;
505
776
  }
506
777
  }
507
- async updateThread({ id, title, metadata }) {
778
+ async updateThread({ id, title, metadata, fence }) {
508
779
  const thread = await this.getThreadById({ threadId: id });
509
780
  if (!thread) throw new MastraError({
510
781
  id: createStorageErrorId("UPSTASH", "UPDATE_THREAD", "FAILED"),
@@ -524,9 +795,13 @@ var StoreMemoryUpstash = class extends MemoryStorage {
524
795
  updatedAt: now
525
796
  };
526
797
  try {
527
- await this.saveThread({ thread: updatedThread });
798
+ await this.saveThread({
799
+ thread: updatedThread,
800
+ fence
801
+ });
528
802
  return updatedThread;
529
803
  } catch (error) {
804
+ if (isRunFenceConflictError(error)) throw error;
530
805
  throw new MastraError({
531
806
  id: createStorageErrorId("UPSTASH", "UPDATE_THREAD", "FAILED"),
532
807
  domain: ErrorDomain.STORAGE,
@@ -561,6 +836,7 @@ var StoreMemoryUpstash = class extends MemoryStorage {
561
836
  }
562
837
  async saveMessages(args) {
563
838
  const { messages } = args;
839
+ const check = this.#runFenceCheck(args.fence, "saveMessages");
564
840
  if (messages.length === 0) return { messages: [] };
565
841
  try {
566
842
  for (const message of messages) {
@@ -615,7 +891,7 @@ var StoreMemoryUpstash = class extends MemoryStorage {
615
891
  }, /* @__PURE__ */ new Error(`Thread ${targetThreadId} not found`));
616
892
  for (let i = 0; i < messagesWithIndex.length; i += batchSize) {
617
893
  const batch = messagesWithIndex.slice(i, i + batchSize);
618
- const pipeline = this.client.pipeline();
894
+ const pipeline = writeBatch(this.client, check);
619
895
  const batchTouchedThreadIds = /* @__PURE__ */ new Set();
620
896
  const batchExistingThreadIds = existingThreadIds.slice(i, i + batch.length);
621
897
  for (const [batchIndex, message] of batch.entries()) {
@@ -930,6 +1206,9 @@ var StoreMemoryUpstash = class extends MemoryStorage {
930
1206
  }
931
1207
  }
932
1208
  async saveResource({ resource }) {
1209
+ return this.#saveResource(resource, void 0);
1210
+ }
1211
+ async #saveResource(resource, check) {
933
1212
  try {
934
1213
  const key = `${TABLE_RESOURCES}:${resource.id}`;
935
1214
  const serializedResource = {
@@ -938,14 +1217,17 @@ var StoreMemoryUpstash = class extends MemoryStorage {
938
1217
  createdAt: resource.createdAt.toISOString(),
939
1218
  updatedAt: resource.updatedAt.toISOString()
940
1219
  };
941
- await this.client.set(key, serializedResource);
1220
+ const batch = writeBatch(this.client, check);
1221
+ batch.set(key, serializedResource);
1222
+ await batch.exec();
942
1223
  return resource;
943
1224
  } catch (error) {
944
- this.logger.error("Error saving resource:", error);
1225
+ if (!isRunFenceConflictError(error)) this.logger.error("Error saving resource:", error);
945
1226
  throw error;
946
1227
  }
947
1228
  }
948
- async updateResource({ resourceId, workingMemory, metadata }) {
1229
+ async updateResource({ resourceId, workingMemory, metadata, fence }) {
1230
+ const check = this.#runFenceCheck(fence, "updateResource");
949
1231
  try {
950
1232
  const existingResource = await this.getResourceById({ resourceId });
951
1233
  if (!existingResource) {
@@ -956,7 +1238,7 @@ var StoreMemoryUpstash = class extends MemoryStorage {
956
1238
  createdAt: /* @__PURE__ */ new Date(),
957
1239
  updatedAt: /* @__PURE__ */ new Date()
958
1240
  };
959
- return this.saveResource({ resource: newResource });
1241
+ return await this.#saveResource(newResource, check);
960
1242
  }
961
1243
  const updatedResource = {
962
1244
  ...existingResource,
@@ -967,16 +1249,17 @@ var StoreMemoryUpstash = class extends MemoryStorage {
967
1249
  },
968
1250
  updatedAt: /* @__PURE__ */ new Date()
969
1251
  };
970
- await this.saveResource({ resource: updatedResource });
1252
+ await this.#saveResource(updatedResource, check);
971
1253
  return updatedResource;
972
1254
  } catch (error) {
973
- this.logger.error("Error updating resource:", error);
1255
+ if (!isRunFenceConflictError(error)) this.logger.error("Error updating resource:", error);
974
1256
  throw error;
975
1257
  }
976
1258
  }
977
1259
  async updateMessages(args) {
978
1260
  const { messages } = args;
979
1261
  if (messages.length === 0) return [];
1262
+ const check = this.#runFenceCheck(args.fence, "updateMessages");
980
1263
  try {
981
1264
  const messageIds = messages.map((m) => m.id);
982
1265
  const updatesById = new Map(messages.map((message) => [message.id, message]));
@@ -1029,7 +1312,10 @@ var StoreMemoryUpstash = class extends MemoryStorage {
1029
1312
  for (const [messageId, threadId] of Object.entries(backfillIndexValues)) backfillPipeline.set(getMessageIndexKey(messageId), threadId);
1030
1313
  await backfillPipeline.exec();
1031
1314
  }
1032
- if (existingMessages.length === 0) return [];
1315
+ if (existingMessages.length === 0) {
1316
+ await assertRunFence(this.client, check);
1317
+ return [];
1318
+ }
1033
1319
  const threadIdsToUpdate = /* @__PURE__ */ new Set();
1034
1320
  const destinationThreadIds = /* @__PURE__ */ new Set();
1035
1321
  for (const existingMessage of existingMessages) {
@@ -1061,7 +1347,7 @@ var StoreMemoryUpstash = class extends MemoryStorage {
1061
1347
  domain: ErrorDomain.STORAGE,
1062
1348
  category: ErrorCategory.USER
1063
1349
  }, /* @__PURE__ */ new Error(`Thread ${destinationThreadId} not found`));
1064
- const pipeline = this.client.pipeline();
1350
+ const pipeline = writeBatch(this.client, check);
1065
1351
  for (const existingMessage of existingMessages) {
1066
1352
  const updatePayload = updatesById.get(existingMessage.id);
1067
1353
  if (!updatePayload) continue;
@@ -1131,8 +1417,9 @@ var StoreMemoryUpstash = class extends MemoryStorage {
1131
1417
  }, error);
1132
1418
  }
1133
1419
  }
1134
- async deleteMessages(messageIds) {
1420
+ async deleteMessages(messageIds, options) {
1135
1421
  if (!messageIds || messageIds.length === 0) return;
1422
+ const check = this.#runFenceCheck(options?.fence, "deleteMessages");
1136
1423
  try {
1137
1424
  const threadIds = /* @__PURE__ */ new Set();
1138
1425
  const messageKeys = [];
@@ -1168,8 +1455,11 @@ var StoreMemoryUpstash = class extends MemoryStorage {
1168
1455
  }
1169
1456
  }
1170
1457
  }
1171
- if (messageKeys.length === 0) return;
1172
- const pipeline = this.client.pipeline();
1458
+ if (messageKeys.length === 0) {
1459
+ await assertRunFence(this.client, check);
1460
+ return;
1461
+ }
1462
+ const pipeline = writeBatch(this.client, check);
1173
1463
  for (const key of messageKeys) pipeline.del(key);
1174
1464
  for (const messageId of foundMessageIds) pipeline.del(getMessageIndexKey(messageId));
1175
1465
  if (threadIds.size > 0) for (const threadId of threadIds) {
@@ -1185,6 +1475,7 @@ var StoreMemoryUpstash = class extends MemoryStorage {
1185
1475
  }
1186
1476
  await pipeline.exec();
1187
1477
  } catch (error) {
1478
+ if (isRunFenceConflictError(error)) throw error;
1188
1479
  throw new MastraError({
1189
1480
  id: createStorageErrorId("UPSTASH", "DELETE_MESSAGES", "FAILED"),
1190
1481
  domain: ErrorDomain.STORAGE,
@@ -1612,8 +1903,56 @@ var WorkflowsUpstash = class extends WorkflowsStorage {
1612
1903
  resourceId: row.resourceId
1613
1904
  };
1614
1905
  }
1906
+ supportsRunFencing() {
1907
+ return true;
1908
+ }
1615
1909
  async dangerouslyClearAll() {
1616
1910
  await this.#db.deleteData({ tableName: TABLE_WORKFLOW_SNAPSHOT });
1911
+ await this.#db.scanAndDelete(`${runClaimKey(TABLE_WORKFLOW_RUN_OWNERS, "")}*`);
1912
+ }
1913
+ #runFenceCheck(fence, runId, operation) {
1914
+ const resolved = resolveRunFence(this, fence, runId);
1915
+ return resolved && {
1916
+ claims: TABLE_WORKFLOW_RUN_OWNERS,
1917
+ fence: resolved,
1918
+ operation
1919
+ };
1920
+ }
1921
+ #ownershipError(operation, runId, error) {
1922
+ return new MastraError({
1923
+ id: createStorageErrorId("UPSTASH", operation, "FAILED"),
1924
+ domain: ErrorDomain.STORAGE,
1925
+ category: ErrorCategory.THIRD_PARTY,
1926
+ details: { runId }
1927
+ }, error);
1928
+ }
1929
+ async claimRunOwnership(args) {
1930
+ try {
1931
+ return await claimRunOwnership(this.client, args);
1932
+ } catch (error) {
1933
+ throw this.#ownershipError("CLAIM_RUN_OWNERSHIP", args.runId, error);
1934
+ }
1935
+ }
1936
+ async renewRunOwnership(args) {
1937
+ try {
1938
+ return await renewRunOwnership(this.client, args);
1939
+ } catch (error) {
1940
+ throw this.#ownershipError("RENEW_RUN_OWNERSHIP", args.runId, error);
1941
+ }
1942
+ }
1943
+ async releaseRunOwnership(fence) {
1944
+ try {
1945
+ return await releaseRunOwnership(this.client, fence);
1946
+ } catch (error) {
1947
+ throw this.#ownershipError("RELEASE_RUN_OWNERSHIP", fence.runId, error);
1948
+ }
1949
+ }
1950
+ async getRunOwnership({ runId }) {
1951
+ try {
1952
+ return await getRunOwnership(this.client, runId);
1953
+ } catch (error) {
1954
+ throw this.#ownershipError("GET_RUN_OWNERSHIP", runId, error);
1955
+ }
1617
1956
  }
1618
1957
  async getWorkflowRunRecord({ namespace, runId, resourceId, workflowName }) {
1619
1958
  if (workflowName) {
@@ -1641,7 +1980,8 @@ var WorkflowsUpstash = class extends WorkflowsStorage {
1641
1980
  return this.client.get(key);
1642
1981
  }))).find((w) => w?.run_id === runId && (!workflowName || w.workflow_name === workflowName) && (!resourceId || w.resourceId === resourceId)) ?? null;
1643
1982
  }
1644
- async updateWorkflowResults({ workflowName, runId, stepId, result, requestContext, state }) {
1983
+ async updateWorkflowResults({ workflowName, runId, stepId, result, requestContext, state, fence }) {
1984
+ const check = this.#runFenceCheck(fence, runId, "updateWorkflowResults");
1645
1985
  try {
1646
1986
  const key = getKey(TABLE_WORKFLOW_SNAPSHOT, {
1647
1987
  namespace: "workflows",
@@ -1649,7 +1989,7 @@ var WorkflowsUpstash = class extends WorkflowsStorage {
1649
1989
  run_id: runId
1650
1990
  });
1651
1991
  const now = (/* @__PURE__ */ new Date()).toISOString();
1652
- const resultJson = await this.client.eval(`
1992
+ const resultJson = await evalFenced(this.client, `
1653
1993
  local key = KEYS[1]
1654
1994
  local stepId = ARGV[1]
1655
1995
  local resultJson = ARGV[2]
@@ -1739,7 +2079,7 @@ var WorkflowsUpstash = class extends WorkflowsStorage {
1739
2079
  runId,
1740
2080
  String(Date.now()),
1741
2081
  state === void 0 ? "" : JSON.stringify(state)
1742
- ]);
2082
+ ], check);
1743
2083
  let data;
1744
2084
  if (typeof resultJson === "string") data = JSON.parse(resultJson);
1745
2085
  else data = resultJson;
@@ -1758,7 +2098,8 @@ var WorkflowsUpstash = class extends WorkflowsStorage {
1758
2098
  }, error);
1759
2099
  }
1760
2100
  }
1761
- async updateWorkflowState({ workflowName, runId, opts }) {
2101
+ async updateWorkflowState({ workflowName, runId, opts, fence }) {
2102
+ const check = this.#runFenceCheck(fence, runId, "updateWorkflowState");
1762
2103
  try {
1763
2104
  const key = getKey(TABLE_WORKFLOW_SNAPSHOT, {
1764
2105
  namespace: "workflows",
@@ -1823,11 +2164,11 @@ var WorkflowsUpstash = class extends WorkflowsStorage {
1823
2164
  `;
1824
2165
  const { expectedStatus, ...state } = opts;
1825
2166
  const expectedStatusJson = expectedStatus === void 0 ? "" : JSON.stringify(Array.isArray(expectedStatus) ? expectedStatus : [expectedStatus]);
1826
- const resultJson = await this.client.eval(luaScript, [key], [
2167
+ const resultJson = await evalFenced(this.client, luaScript, [key], [
1827
2168
  JSON.stringify(state),
1828
2169
  now,
1829
2170
  expectedStatusJson
1830
- ]);
2171
+ ], check);
1831
2172
  if (!resultJson) return;
1832
2173
  let data;
1833
2174
  if (typeof resultJson === "string") data = JSON.parse(resultJson);
@@ -1848,6 +2189,7 @@ var WorkflowsUpstash = class extends WorkflowsStorage {
1848
2189
  }
1849
2190
  async persistWorkflowSnapshot(params) {
1850
2191
  const { namespace = "workflows", workflowName, runId, resourceId, snapshot, createdAt, updatedAt } = params;
2192
+ const check = this.#runFenceCheck(params.fence, runId, "persistWorkflowSnapshot");
1851
2193
  try {
1852
2194
  const now = /* @__PURE__ */ new Date();
1853
2195
  const key = getKey(TABLE_WORKFLOW_SNAPSHOT, {
@@ -1865,7 +2207,7 @@ var WorkflowsUpstash = class extends WorkflowsStorage {
1865
2207
  createdAt: (createdAt ?? now).toISOString(),
1866
2208
  updatedAt: (updatedAt ?? now).toISOString()
1867
2209
  };
1868
- await this.client.eval(`
2210
+ await evalFenced(this.client, `
1869
2211
  local existing = redis.call("GET", KEYS[1])
1870
2212
  local next = cjson.decode(ARGV[1])
1871
2213
  if existing then
@@ -1876,8 +2218,9 @@ var WorkflowsUpstash = class extends WorkflowsStorage {
1876
2218
  end
1877
2219
  redis.call("SET", KEYS[1], cjson.encode(next))
1878
2220
  return 1
1879
- `, [key], [JSON.stringify(record)]);
2221
+ `, [key], [JSON.stringify(record)], check);
1880
2222
  } catch (error) {
2223
+ if (isRunFenceConflictError(error)) throw error;
1881
2224
  throw new MastraError({
1882
2225
  id: createStorageErrorId("UPSTASH", "PERSIST_WORKFLOW_SNAPSHOT", "FAILED"),
1883
2226
  domain: ErrorDomain.STORAGE,
@@ -1935,7 +2278,8 @@ var WorkflowsUpstash = class extends WorkflowsStorage {
1935
2278
  }, error);
1936
2279
  }
1937
2280
  }
1938
- async deleteWorkflowRunById({ runId, workflowName }) {
2281
+ async deleteWorkflowRunById({ runId, workflowName, fence }) {
2282
+ const check = this.#runFenceCheck(fence, runId, "deleteWorkflowRunById");
1939
2283
  try {
1940
2284
  const record = await this.getWorkflowRunRecord({
1941
2285
  namespace: "workflows",
@@ -1948,8 +2292,11 @@ var WorkflowsUpstash = class extends WorkflowsStorage {
1948
2292
  run_id: runId,
1949
2293
  ...record?.resourceId ? { resourceId: record.resourceId } : {}
1950
2294
  });
1951
- await this.client.del(key);
2295
+ const batch = writeBatch(this.client, check);
2296
+ batch.del(key);
2297
+ await batch.exec();
1952
2298
  } catch (error) {
2299
+ if (isRunFenceConflictError(error)) throw error;
1953
2300
  throw new MastraError({
1954
2301
  id: createStorageErrorId("UPSTASH", "DELETE_WORKFLOW_RUN_BY_ID", "FAILED"),
1955
2302
  domain: ErrorDomain.STORAGE,