@mastra/dsql 1.5.1 → 1.5.2-alpha.0

package/dist/index.js

@@ -1,6 +1,6 @@
 import { AuroraDSQLClient } from "@aws/aurora-dsql-node-postgres-connector";
 import { ErrorCategory, ErrorDomain, MastraError } from "@mastra/core/error";
-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";
+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";
 import { Pool } from "pg";
 import { MastraBase } from "@mastra/core/base";
 import { parseSqlIdentifier } from "@mastra/core/utils";
@@ -1824,6 +1824,49 @@ var AgentsDSQL = class AgentsDSQL extends AgentsStorage {
 	}
 };
 //#endregion
+//#region src/storage/db/run-fencing.ts
+/**
+* Epoch milliseconds on the database clock. Aurora DSQL never waits on a row
+* lock, so the transaction's start time is the time its statements run, and
+* every retry starts a new transaction with a fresh clock.
+*/
+const DB_NOW_MS = `(EXTRACT(EPOCH FROM now()) * 1000)::bigint`;
+/**
+* Retry predicate for writes to an ownership or fence row that may not exist
+* yet. Two first writers both insert the same key; the loser normally fails
+* at commit with an OCC conflict (40001), but a duplicate-key error (23505)
+* means the same thing here: the retry sees the row and takes the update path.
+*/
+function isRetriableRunFenceWrite(error) {
+	return isRetriableError(error) || getErrorCode(error) === "23505";
+}
+/**
+* The check half of a fenced write; call it first inside the write's
+* transaction. Aurora DSQL has no blocking row locks: `FOR UPDATE` puts the
+* claim row in this transaction's write set instead, so a claim that commits
+* against the row while the write is in flight makes the write fail at commit
+* with an OCC conflict. The retry then sees the new claim and is rejected.
+*/
+async function assertRunFence(t, claimsTable, fence, operation) {
+	const row = await t.oneOrNone(`SELECT generation, "ownerId" FROM ${claimsTable} WHERE "runId" = $1 FOR UPDATE`, [fence.runId]);
+	if (!matchesRunFence(row && {
+		generation: Number(row.generation),
+		ownerId: row.ownerId
+	}, fence)) throw new RunFenceConflictError(fence, operation);
+}
+/**
+* Runs `write` directly when it carries no fence, or in a transaction behind
+* the fence check when it does. Callers own retries: wrap the call in
+* `withRetry` so an OCC conflict re-runs the check as well as the write.
+*/
+function withRunFence(client, claimsTable, fence, operation, write) {
+	if (!fence) return write(client);
+	return client.tx(async (t) => {
+		await assertRunFence(t, claimsTable, fence, operation);
+		return write(t);
+	});
+}
+//#endregion
 //#region src/storage/domains/memory/index.ts
 /**
 * Generate SQL placeholder string for IN clauses.
@@ -1883,6 +1926,10 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 			schema: TABLE_SCHEMAS[TABLE_MESSAGES],
 			ifNotExists: ["resourceId"]
 		});
+		await this.#db.createTable({
+			tableName: TABLE_MEMORY_RUN_FENCES,
+			schema: RUN_FENCING_TABLE_SCHEMAS[TABLE_MEMORY_RUN_FENCES]
+		});
 		await this.createDefaultIndexes();
 		await this.createCustomIndexes();
 	}
@@ -1928,6 +1975,62 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 		await this.#db.clearTable({ tableName: TABLE_MESSAGES });
 		await this.#db.clearTable({ tableName: TABLE_THREADS });
 		await this.#db.clearTable({ tableName: TABLE_RESOURCES });
+		await this.#db.clearTable({ tableName: TABLE_MEMORY_RUN_FENCES });
+	}
+	supportsRunFencing() {
+		return true;
+	}
+	#runFencesTable() {
+		return getTableName({
+			indexName: TABLE_MEMORY_RUN_FENCES,
+			schemaName: getSchemaName(this.#schema)
+		});
+	}
+	/** Runs `write` behind the run's fence check when it carries a fence. */
+	#fenced(fence, operation, write) {
+		return withRunFence(this.#db.client, this.#runFencesTable(), resolveRunFence(this, fence), operation, write);
+	}
+	async raiseRunFence(fence) {
+		try {
+			const { result } = await withRetry(() => this.#db.client.oneOrNone(`INSERT INTO ${this.#runFencesTable()} AS f ("runId", generation, "ownerId")
+             VALUES ($1, $2, $3)
+             ON CONFLICT ("runId") DO UPDATE
+             SET generation = EXCLUDED.generation, "ownerId" = EXCLUDED."ownerId", "retiredAt" = NULL
+             WHERE f.generation < EXCLUDED.generation
+                OR (f.generation = EXCLUDED.generation AND f."ownerId" = EXCLUDED."ownerId")
+             RETURNING "runId"`, [
+				fence.runId,
+				fence.generation,
+				fence.ownerId
+			]), { isRetriable: isRetriableRunFenceWrite });
+			return result !== null;
+		} catch (error) {
+			throw new MastraError({
+				id: createStorageErrorId("DSQL", "RAISE_RUN_FENCE", "FAILED"),
+				domain: ErrorDomain.STORAGE,
+				category: ErrorCategory.THIRD_PARTY,
+				details: { runId: fence.runId }
+			}, error);
+		}
+	}
+	async retireRunFence(fence) {
+		try {
+			const { result } = await withRetry(() => this.#db.client.oneOrNone(`UPDATE ${this.#runFencesTable()} SET "retiredAt" = ${DB_NOW_MS}
+             WHERE "runId" = $1 AND generation = $2 AND "ownerId" = $3
+             RETURNING "runId"`, [
+				fence.runId,
+				fence.generation,
+				fence.ownerId
+			]), { isRetriable: isRetriableRunFenceWrite });
+			return result !== null;
+		} catch (error) {
+			throw new MastraError({
+				id: createStorageErrorId("DSQL", "RETIRE_RUN_FENCE", "FAILED"),
+				domain: ErrorDomain.STORAGE,
+				category: ErrorCategory.THIRD_PARTY,
+				details: { runId: fence.runId }
+			}, error);
+		}
 	}
 	/**
 	* Normalizes message row from database by applying createdAtZ fallback
@@ -2064,7 +2167,7 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 			throw mastraError;
 		}
 	}
-	async saveThread({ thread }) {
+	async saveThread({ thread, fence }) {
 		const tableName = getTableName({
 			indexName: TABLE_THREADS,
 			schemaName: getSchemaName(this.#schema)
@@ -2072,7 +2175,7 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 		const createdAt = toUtcISOString(thread.createdAt);
 		const updatedAt = toUtcISOString(thread.updatedAt);
 		await withRetry(async () => {
-			await this.#db.client.none(`INSERT INTO ${tableName} (
+			await this.#fenced(fence, "saveThread", (q) => q.none(`INSERT INTO ${tableName} (
             id,
             "resourceId",
             title,
@@ -2098,10 +2201,11 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 				createdAt,
 				updatedAt,
 				updatedAt
-			]);
+			]));
 		}, { onRetry: (error, attempt, delay) => {
 			this.logger?.warn?.(`saveThread retry ${attempt} for ${thread.id} after ${delay}ms: ${error.message}`);
 		} }).catch((error) => {
+			if (isRunFenceConflictError(error)) throw error;
 			throw new MastraError({
 				id: createStorageErrorId("DSQL", "SAVE_THREAD", "FAILED"),
 				domain: ErrorDomain.STORAGE,
@@ -2111,13 +2215,13 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 		});
 		return thread;
 	}
-	async updateThread({ id, title, metadata }) {
+	async updateThread({ id, title, metadata, fence }) {
 		const threadTableName = getTableName({
 			indexName: TABLE_THREADS,
 			schemaName: getSchemaName(this.#schema)
 		});
-		const { result } = await withRetry(async () => {
-			const existingThread = await this.getThreadById({ threadId: id });
+		const { result } = await withRetry(() => this.#fenced(fence, "updateThread", async (q) => {
+			const existingThread = await q.oneOrNone(`SELECT metadata FROM ${threadTableName} WHERE id = $1`, [id]);
 			if (!existingThread) throw new MastraError({
 				id: createStorageErrorId("DSQL", "UPDATE_THREAD", "NOT_FOUND"),
 				domain: ErrorDomain.STORAGE,
@@ -2129,11 +2233,11 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 				}
 			});
 			const mergedMetadata = {
-				...existingThread.metadata,
+				...typeof existingThread.metadata === "string" ? JSON.parse(existingThread.metadata) : existingThread.metadata,
 				...metadata
 			};
 			const now = (/* @__PURE__ */ new Date()).toISOString();
-			const thread = await this.#db.client.one(`UPDATE ${threadTableName}
+			return q.one(`UPDATE ${threadTableName}
                       SET 
                           title = COALESCE($1, title),
                           metadata = $2,
@@ -2148,15 +2252,14 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 				now,
 				id
 			]);
-			return {
-				id: thread.id,
-				resourceId: thread.resourceId,
-				title: thread.title,
-				metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
-				createdAt: thread.createdAtZ || thread.createdAt,
-				updatedAt: thread.updatedAtZ || thread.updatedAt
-			};
-		}, { onRetry: (error, attempt, delay) => {
+		}).then((thread) => ({
+			id: thread.id,
+			resourceId: thread.resourceId,
+			title: thread.title,
+			metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
+			createdAt: thread.createdAtZ || thread.createdAt,
+			updatedAt: thread.updatedAtZ || thread.updatedAt
+		})), { onRetry: (error, attempt, delay) => {
 			this.logger?.warn?.(`updateThread retry ${attempt} for ${id} after ${delay}ms: ${error.message}`);
 		} }).catch((error) => {
 			if (error instanceof MastraError) throw error;
@@ -2481,7 +2584,7 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 			throw mastraError;
 		}
 	}
-	async saveMessages({ messages }) {
+	async saveMessages({ messages, fence: explicitFence }) {
 		if (messages.length === 0) return { messages: [] };
 		const threadId = messages[0]?.threadId;
 		if (!threadId) throw new MastraError({
@@ -2505,8 +2608,10 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 			indexName: TABLE_THREADS,
 			schemaName: getSchemaName(this.#schema)
 		});
+		const fence = resolveRunFence(this, explicitFence);
 		await withRetry(async () => {
 			await this.#db.client.tx(async (t) => {
+				if (fence) await assertRunFence(t, this.#runFencesTable(), fence, "saveMessages");
 				const messageInserts = messages.map((message) => {
 					if (!message.threadId) throw new Error(`Expected to find a threadId for message, but couldn't find one. An unexpected error has occurred.`);
 					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 {
 		}, { onRetry: (error, attempt, delay) => {
 			this.logger?.warn?.(`saveMessages retry ${attempt} for thread ${threadId} after ${delay}ms: ${error.message}`);
 		} }).catch((error) => {
+			if (isRunFenceConflictError(error)) throw error;
 			throw new MastraError({
 				id: createStorageErrorId("DSQL", "SAVE_MESSAGES", "FAILED"),
 				domain: ErrorDomain.STORAGE,
@@ -2565,8 +2671,9 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 		});
 		return { messages: new MessageList().add(messagesWithParsedContent, "memory").get.all.db() };
 	}
-	async updateMessages({ messages }) {
+	async updateMessages({ messages, fence: explicitFence }) {
 		if (messages.length === 0) return [];
+		const fence = resolveRunFence(this, explicitFence);
 		const messageIds = messages.map((m) => m.id);
 		const selectQuery = `SELECT id, content, role, type, "createdAt", "createdAtZ", thread_id AS "threadId", "resourceId" FROM ${getTableName({
 			indexName: TABLE_MESSAGES,
@@ -2583,6 +2690,7 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 		const threadIdsToUpdate = /* @__PURE__ */ new Set();
 		await withRetry(async () => {
 			await this.#db.client.tx(async (t) => {
+				if (fence) await assertRunFence(t, this.#runFencesTable(), fence, "updateMessages");
 				const queries = [];
 				const columnMapping = { threadId: "thread_id" };
 				for (const existingMessage of existingMessages) {
@@ -2635,6 +2743,7 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 		}, { onRetry: (error, attempt, delay) => {
 			this.logger?.warn?.(`updateMessages retry ${attempt} for ${messageIds.length} messages after ${delay}ms: ${error.message}`);
 		} }).catch((error) => {
+			if (isRunFenceConflictError(error)) throw error;
 			throw new MastraError({
 				id: createStorageErrorId("DSQL", "UPDATE_MESSAGES", "FAILED"),
 				domain: ErrorDomain.STORAGE,
@@ -2653,8 +2762,9 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 			return message;
 		});
 	}
-	async deleteMessages(messageIds) {
+	async deleteMessages(messageIds, options) {
 		if (!messageIds || messageIds.length === 0) return;
+		const fence = resolveRunFence(this, options?.fence);
 		const messageTableName = getTableName({
 			indexName: TABLE_MESSAGES,
 			schemaName: getSchemaName(this.#schema)
@@ -2665,6 +2775,7 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 		});
 		await withRetry(async () => {
 			await this.#db.client.tx(async (t) => {
+				if (fence) await assertRunFence(t, this.#runFencesTable(), fence, "deleteMessages");
 				const placeholders = messageIds.map((_, idx) => `$${idx + 1}`).join(",");
 				const threadIds = (await t.manyOrNone(`SELECT DISTINCT thread_id FROM ${messageTableName} WHERE id IN (${placeholders})`, messageIds))?.map((msg) => msg.thread_id).filter(Boolean) || [];
 				await t.none(`DELETE FROM ${messageTableName} WHERE id IN (${placeholders})`, messageIds);
@@ -2676,6 +2787,7 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 		}, { onRetry: (error, attempt, delay) => {
 			this.logger?.warn?.(`deleteMessages retry ${attempt} for ${messageIds.length} messages after ${delay}ms: ${error.message}`);
 		} }).catch((error) => {
+			if (isRunFenceConflictError(error)) throw error;
 			throw new MastraError({
 				id: createStorageErrorId("DSQL", "DELETE_MESSAGES", "FAILED"),
 				domain: ErrorDomain.STORAGE,
@@ -2720,23 +2832,42 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 		});
 		return resource;
 	}
-	async updateResource({ resourceId, workingMemory, metadata }) {
+	async updateResource({ resourceId, workingMemory, metadata, fence }) {
 		const tableName = getTableName({
 			indexName: TABLE_RESOURCES,
 			schemaName: getSchemaName(this.#schema)
 		});
-		const { result } = await withRetry(async () => {
-			const existingResource = await this.getResourceById({ resourceId });
-			if (!existingResource) {
+		const { result } = await withRetry(() => this.#fenced(fence, "updateResource", async (q) => {
+			const row = await q.oneOrNone(`SELECT * FROM ${tableName} WHERE id = $1`, [resourceId]);
+			if (!row) {
+				const now = /* @__PURE__ */ new Date();
 				const newResource = {
 					id: resourceId,
 					workingMemory,
 					metadata: metadata || {},
-					createdAt: /* @__PURE__ */ new Date(),
-					updatedAt: /* @__PURE__ */ new Date()
+					createdAt: now,
+					updatedAt: now
 				};
-				return this.saveResource({ resource: newResource });
+				const nowIso = now.toISOString();
+				await q.none(`INSERT INTO ${tableName} (id, "workingMemory", metadata, "createdAt", "createdAtZ", "updatedAt", "updatedAtZ")
+               VALUES ($1, $2, $3, $4, $5, $6, $7)`, [
+					resourceId,
+					workingMemory ?? null,
+					JSON.stringify(newResource.metadata),
+					nowIso,
+					nowIso,
+					nowIso,
+					nowIso
+				]);
+				return newResource;
 			}
+			const existingResource = {
+				id: row.id,
+				createdAt: row.createdAtZ || row.createdAt,
+				updatedAt: row.updatedAtZ || row.updatedAt,
+				workingMemory: row.workingMemory,
+				metadata: typeof row.metadata === "string" ? JSON.parse(row.metadata) : row.metadata
+			};
 			const updatedResource = {
 				...existingResource,
 				workingMemory: workingMemory !== void 0 ? workingMemory : existingResource.workingMemory,
@@ -2766,9 +2897,9 @@ var MemoryDSQL = class MemoryDSQL extends MemoryStorage {
 			values.push(updatedResource.updatedAt.toISOString());
 			paramIndex++;
 			values.push(resourceId);
-			await this.#db.client.none(`UPDATE ${tableName} SET ${updates.join(", ")} WHERE id = $${paramIndex}`, values);
+			await q.none(`UPDATE ${tableName} SET ${updates.join(", ")} WHERE id = $${paramIndex}`, values);
 			return updatedResource;
-		}, { onRetry: (error, attempt, delay) => {
+		}), { onRetry: (error, attempt, delay) => {
 			this.logger?.warn?.(`updateResource retry ${attempt} for ${resourceId} after ${delay}ms: ${error.message}`);
 		} }).catch((error) => {
 			if (error instanceof MastraError) throw error;
@@ -3718,6 +3849,17 @@ var ScoresDSQL = class ScoresDSQL extends ScoresStorage {
 };
 //#endregion
 //#region src/storage/domains/workflows/index.ts
+const RUN_OWNER_COLUMNS = `generation, "ownerId", "leaseExpiresAt", ${DB_NOW_MS} AS "nowMs"`;
+function toRunOwnershipRecord(runId, row) {
+	const leaseExpiresAt = row.leaseExpiresAt === null ? null : Number(row.leaseExpiresAt);
+	return {
+		runId,
+		generation: Number(row.generation),
+		ownerId: row.ownerId,
+		leaseExpiresAt: leaseExpiresAt === null ? null : new Date(leaseExpiresAt),
+		live: leaseExpiresAt !== null && leaseExpiresAt > Number(row.nowMs)
+	};
+}
 function parseWorkflowRun(row) {
 	let parsedSnapshot = row.snapshot;
 	if (typeof parsedSnapshot === "string") try {
@@ -3779,6 +3921,10 @@ var WorkflowsDSQL = class WorkflowsDSQL extends WorkflowsStorage {
 			schema: TABLE_SCHEMAS[TABLE_WORKFLOW_SNAPSHOT],
 			ifNotExists: ["resourceId"]
 		});
+		await this.#db.createTable({
+			tableName: TABLE_WORKFLOW_RUN_OWNERS,
+			schema: RUN_FENCING_TABLE_SCHEMAS[TABLE_WORKFLOW_RUN_OWNERS]
+		});
 		await this.createDefaultIndexes();
 		await this.createCustomIndexes();
 	}
@@ -3795,11 +3941,130 @@ var WorkflowsDSQL = class WorkflowsDSQL extends WorkflowsStorage {
 	}
 	async dangerouslyClearAll() {
 		await this.#db.clearTable({ tableName: TABLE_WORKFLOW_SNAPSHOT });
+		await this.#db.clearTable({ tableName: TABLE_WORKFLOW_RUN_OWNERS });
 	}
-	async updateWorkflowResults({ workflowName, runId, stepId, result, requestContext, state }) {
+	supportsRunFencing() {
+		return true;
+	}
+	#runOwnersTable() {
+		return getTableName({
+			indexName: TABLE_WORKFLOW_RUN_OWNERS,
+			schemaName: getSchemaName(this.#schema)
+		});
+	}
+	async #readRunOwner(q, runId) {
+		const row = await q.oneOrNone(`SELECT ${RUN_OWNER_COLUMNS} FROM ${this.#runOwnersTable()} WHERE "runId" = $1`, [runId]);
+		return row ? toRunOwnershipRecord(runId, row) : null;
+	}
+	#ownershipError(operation, runId, error) {
+		return new MastraError({
+			id: createStorageErrorId("DSQL", operation, "FAILED"),
+			domain: ErrorDomain.STORAGE,
+			category: ErrorCategory.THIRD_PARTY,
+			details: { runId }
+		}, error);
+	}
+	async claimRunOwnership({ runId, ownerId, leaseMs, force, expectedGeneration }) {
+		const table = this.#runOwnersTable();
+		try {
+			const { result } = await withRetry(() => this.#db.client.tx(async (t) => {
+				if (!await t.oneOrNone(`SELECT 1 FROM ${table} WHERE "runId" = $1 FOR UPDATE`, [runId])) {
+					if (expectedGeneration !== void 0 && expectedGeneration !== 0) return {
+						acquired: false,
+						record: null
+					};
+					return {
+						acquired: true,
+						record: toRunOwnershipRecord(runId, await t.one(`INSERT INTO ${table} ("runId", generation, "ownerId", "leaseExpiresAt", "updatedAt")
+                 VALUES ($1, ${DB_NOW_MS}, $2, ${DB_NOW_MS} + $3, ${DB_NOW_MS})
+                 RETURNING ${RUN_OWNER_COLUMNS}`, [
+							runId,
+							ownerId,
+							leaseMs
+						]))
+					};
+				}
+				const claimed = await t.oneOrNone(`UPDATE ${table}
+               SET generation = generation + 1, "ownerId" = $2, "leaseExpiresAt" = ${DB_NOW_MS} + $3,
+                   "updatedAt" = ${DB_NOW_MS}
+               WHERE "runId" = $1
+                 AND ($4::bigint IS NULL OR generation = $4)
+                 AND ($5::boolean OR "leaseExpiresAt" IS NULL OR "leaseExpiresAt" <= ${DB_NOW_MS})
+               RETURNING ${RUN_OWNER_COLUMNS}`, [
+					runId,
+					ownerId,
+					leaseMs,
+					expectedGeneration ?? null,
+					force === true
+				]);
+				if (claimed) return {
+					acquired: true,
+					record: toRunOwnershipRecord(runId, claimed)
+				};
+				return {
+					acquired: false,
+					record: await this.#readRunOwner(t, runId)
+				};
+			}), { isRetriable: isRetriableRunFenceWrite });
+			return result;
+		} catch (error) {
+			throw this.#ownershipError("CLAIM_RUN_OWNERSHIP", runId, error);
+		}
+	}
+	async renewRunOwnership({ leaseMs, ...fence }) {
+		try {
+			const { result } = await withRetry(async () => {
+				const renewed = await this.#db.client.oneOrNone(`UPDATE ${this.#runOwnersTable()}
+           SET "leaseExpiresAt" = ${DB_NOW_MS} + $4, "updatedAt" = ${DB_NOW_MS}
+           WHERE "runId" = $1 AND generation = $2 AND "ownerId" = $3 AND "leaseExpiresAt" IS NOT NULL
+           RETURNING ${RUN_OWNER_COLUMNS}`, [
+					fence.runId,
+					fence.generation,
+					fence.ownerId,
+					leaseMs
+				]);
+				if (renewed) return {
+					renewed: true,
+					record: toRunOwnershipRecord(fence.runId, renewed)
+				};
+				return {
+					renewed: false,
+					record: await this.#readRunOwner(this.#db.client, fence.runId)
+				};
+			});
+			return result;
+		} catch (error) {
+			throw this.#ownershipError("RENEW_RUN_OWNERSHIP", fence.runId, error);
+		}
+	}
+	async releaseRunOwnership(fence) {
+		try {
+			const { result } = await withRetry(() => this.#db.client.oneOrNone(`UPDATE ${this.#runOwnersTable()}
+           SET "leaseExpiresAt" = NULL, "updatedAt" = ${DB_NOW_MS}
+           WHERE "runId" = $1 AND generation = $2 AND "ownerId" = $3
+           RETURNING "runId"`, [
+				fence.runId,
+				fence.generation,
+				fence.ownerId
+			]));
+			return result !== null;
+		} catch (error) {
+			throw this.#ownershipError("RELEASE_RUN_OWNERSHIP", fence.runId, error);
+		}
+	}
+	async getRunOwnership({ runId }) {
+		try {
+			return await this.#readRunOwner(this.#db.client, runId);
+		} catch (error) {
+			throw this.#ownershipError("GET_RUN_OWNERSHIP", runId, error);
+		}
+	}
+	async updateWorkflowResults({ workflowName, runId, stepId, result, requestContext, state, fence: explicitFence }) {
+		const fence = resolveRunFence(this, explicitFence, runId);
 		try {
 			const { result: context } = await withRetry(async () => {
 				return this.#db.client.tx(async (t) => {
+					if (fence) await assertRunFence(t, this.#runOwnersTable(), fence, "updateWorkflowResults");
 					const tableName = getTableName({
 						indexName: TABLE_WORKFLOW_SNAPSHOT,
 						schemaName: getSchemaName(this.#schema)
@@ -3848,6 +4113,7 @@ var WorkflowsDSQL = class WorkflowsDSQL extends WorkflowsStorage {
 			} });
 			return context;
 		} catch (error) {
+			if (isRunFenceConflictError(error)) throw error;
 			throw new MastraError({
 				id: createStorageErrorId("DSQL", "UPDATE_WORKFLOW_RESULTS", "FAILED"),
 				domain: ErrorDomain.STORAGE,
@@ -3860,10 +4126,12 @@ var WorkflowsDSQL = class WorkflowsDSQL extends WorkflowsStorage {
 			}, error);
 		}
 	}
-	async updateWorkflowState({ workflowName, runId, opts }) {
+	async updateWorkflowState({ workflowName, runId, opts, fence: explicitFence }) {
+		const fence = resolveRunFence(this, explicitFence, runId);
 		try {
 			const { result } = await withRetry(async () => {
 				return this.#db.client.tx(async (t) => {
+					if (fence) await assertRunFence(t, this.#runOwnersTable(), fence, "updateWorkflowState");
 					const tableName = getTableName({
 						indexName: TABLE_WORKFLOW_SNAPSHOT,
 						schemaName: getSchemaName(this.#schema)
@@ -3892,6 +4160,7 @@ var WorkflowsDSQL = class WorkflowsDSQL extends WorkflowsStorage {
 			} });
 			return result;
 		} catch (error) {
+			if (isRunFenceConflictError(error)) throw error;
 			throw new MastraError({
 				id: createStorageErrorId("DSQL", "UPDATE_WORKFLOW_STATE", "FAILED"),
 				domain: ErrorDomain.STORAGE,
@@ -3903,12 +4172,13 @@ var WorkflowsDSQL = class WorkflowsDSQL extends WorkflowsStorage {
 			}, error);
 		}
 	}
-	async persistWorkflowSnapshot({ workflowName, runId, resourceId, snapshot, createdAt, updatedAt }) {
+	async persistWorkflowSnapshot({ workflowName, runId, resourceId, snapshot, createdAt, updatedAt, fence }) {
 		try {
 			const now = /* @__PURE__ */ new Date();
 			const createdAtValue = createdAt ? createdAt : now;
 			const updatedAtValue = updatedAt ? updatedAt : now;
-			await this.#db.client.none(`INSERT INTO ${getTableName({
+			const resolvedFence = resolveRunFence(this, fence, runId);
+			await withRetry(() => withRunFence(this.#db.client, this.#runOwnersTable(), resolvedFence, "persistWorkflowSnapshot", (q) => q.none(`INSERT INTO ${getTableName({
 				indexName: TABLE_WORKFLOW_SNAPSHOT,
 				schemaName: getSchemaName(this.#schema)
 			})} (workflow_name, run_id, "resourceId", snapshot, "createdAt", "updatedAt")
@@ -3921,8 +4191,9 @@ var WorkflowsDSQL = class WorkflowsDSQL extends WorkflowsStorage {
 				JSON.stringify(snapshot),
 				createdAtValue.toISOString(),
 				updatedAtValue.toISOString()
-			]);
+			])));
 		} catch (error) {
+			if (isRunFenceConflictError(error)) throw error;
 			throw new MastraError({
 				id: createStorageErrorId("DSQL", "PERSIST_WORKFLOW_SNAPSHOT", "FAILED"),
 				domain: ErrorDomain.STORAGE,
@@ -3988,13 +4259,15 @@ var WorkflowsDSQL = class WorkflowsDSQL extends WorkflowsStorage {
 			}, error);
 		}
 	}
-	async deleteWorkflowRunById({ runId, workflowName }) {
+	async deleteWorkflowRunById({ runId, workflowName, fence }) {
 		try {
-			await this.#db.client.none(`DELETE FROM ${getTableName({
+			const resolvedFence = resolveRunFence(this, fence, runId);
+			await withRetry(() => withRunFence(this.#db.client, this.#runOwnersTable(), resolvedFence, "deleteWorkflowRunById", (q) => q.none(`DELETE FROM ${getTableName({
 				indexName: TABLE_WORKFLOW_SNAPSHOT,
 				schemaName: getSchemaName(this.#schema)
-			})} WHERE run_id = $1 AND workflow_name = $2`, [runId, workflowName]);
+			})} WHERE run_id = $1 AND workflow_name = $2`, [runId, workflowName])));
 		} catch (error) {
+			if (isRunFenceConflictError(error)) throw error;
 			throw new MastraError({
 				id: createStorageErrorId("DSQL", "DELETE_WORKFLOW_RUN_BY_ID", "FAILED"),
 				domain: ErrorDomain.STORAGE,