@mastra/mssql 1.10.6-alpha.0 → 1.10.6

package/dist/index.js

@@ -2162,9 +2162,17 @@ function withRunFence(pool, claimsTable, fence, operation, write) {
 * duplicate key if they don't. On retry the row exists and the claim
 * serializes behind its row lock.
 */
-async function claimTransaction(pool, fn) {
+function claimTransaction(pool, fn) {
+	return retryWriteRace(() => inTransaction(pool, fn));
+}
+/**
+* Runs `write` again when it lost a race to insert the same row: a deadlock
+* victim or a duplicate key. `write` must be a whole statement or
+* transaction, since a deadlock victim's transaction has already ended.
+*/
+async function retryWriteRace(write) {
 	for (let attempt = 1;; attempt++) try {
-		return await inTransaction(pool, fn);
+		return await write();
 	} catch (error) {
 		const number = error.number;
 		if (attempt >= MAX_CLAIM_ATTEMPTS || !(number === ERR_DEADLOCK_VICTIM || number === ERR_PRIMARY_KEY || number === ERR_UNIQUE_INDEX)) throw error;
@@ -4636,6 +4644,9 @@ var WorkflowsMSSQL = class WorkflowsMSSQL extends WorkflowsStorage {
 		this.skipDefaultIndexes = skipDefaultIndexes;
 		this.indexes = indexes?.filter((idx) => WorkflowsMSSQL.MANAGED_TABLES.includes(idx.table));
 	}
+	supportsGuardedSnapshotWrites() {
+		return true;
+	}
 	supportsConcurrentUpdates() {
 		return true;
 	}
@@ -4945,14 +4956,16 @@ var WorkflowsMSSQL = class WorkflowsMSSQL extends WorkflowsStorage {
 			}, error);
 		}
 	}
-	async persistWorkflowSnapshot({ workflowName, runId, resourceId, snapshot, fence }) {
+	async persistWorkflowSnapshot({ workflowName, runId, resourceId, snapshot, fence, expectedStatus }) {
 		const table = getTableName({
 			indexName: TABLE_WORKFLOW_SNAPSHOT,
 			schemaName: getSchemaName(this.schema)
 		});
 		const now = (/* @__PURE__ */ new Date()).toISOString();
+		const expectedStatuses = expectedStatus === void 0 ? void 0 : Array.isArray(expectedStatus) ? expectedStatus : [expectedStatus];
+		const guard = expectedStatuses === void 0 ? "" : expectedStatuses.length === 0 ? " AND 1 = 0" : ` AND JSON_VALUE(target.[snapshot], '$.status') IN (${expectedStatuses.map((_, i) => `@expectedStatus${i}`).join(", ")})`;
 		try {
-			await withRunFence(this.pool, this.#runOwnersTable, resolveRunFence(this, fence, runId), "persistWorkflowSnapshot", async (q) => {
+			return await retryWriteRace(() => withRunFence(this.pool, this.#runOwnersTable, resolveRunFence(this, fence, runId), "persistWorkflowSnapshot", async (q) => {
 				const request = q.request();
 				request.input("workflow_name", workflowName);
 				request.input("run_id", runId);
@@ -4960,17 +4973,18 @@ var WorkflowsMSSQL = class WorkflowsMSSQL extends WorkflowsStorage {
 				request.input("snapshot", JSON.stringify(snapshot));
 				request.input("createdAt", sql.DateTime2, new Date(now));
 				request.input("updatedAt", sql.DateTime2, new Date(now));
+				expectedStatuses?.forEach((status, i) => request.input(`expectedStatus${i}`, status));
 				const mergeSql = `MERGE INTO ${table} AS target
             USING (SELECT @workflow_name AS workflow_name, @run_id AS run_id) AS src
             ON target.workflow_name = src.workflow_name AND target.run_id = src.run_id
-            WHEN MATCHED THEN UPDATE SET
+            WHEN MATCHED${guard} THEN UPDATE SET
               resourceId = @resourceId,
               snapshot = @snapshot,
               [updatedAt] = @updatedAt
             WHEN NOT MATCHED THEN INSERT (workflow_name, run_id, resourceId, snapshot, [createdAt], [updatedAt])
               VALUES (@workflow_name, @run_id, @resourceId, @snapshot, @createdAt, @updatedAt);`;
-				await request.query(mergeSql);
-			});
+				return ((await request.query(mergeSql)).rowsAffected[0] ?? 0) > 0;
+			}));
 		} catch (error) {
 			if (isRunFenceConflictError(error)) throw error;
 			throw new MastraError({
@@ -4984,6 +4998,27 @@ var WorkflowsMSSQL = class WorkflowsMSSQL extends WorkflowsStorage {
 			}, error);
 		}
 	}
+	async loadWorkflowRunStatus({ workflowName, runId }) {
+		try {
+			const request = this.pool.request();
+			request.input("workflow_name", workflowName);
+			request.input("run_id", runId);
+			return (await request.query(`SELECT TOP 1 JSON_VALUE([snapshot], '$.status') AS status FROM ${getTableName({
+				indexName: TABLE_WORKFLOW_SNAPSHOT,
+				schemaName: getSchemaName(this.schema)
+			})} WHERE workflow_name = @workflow_name AND run_id = @run_id`)).recordset?.[0]?.status ?? null;
+		} catch (error) {
+			throw new MastraError({
+				id: createStorageErrorId("MSSQL", "LOAD_WORKFLOW_RUN_STATUS", "FAILED"),
+				domain: ErrorDomain.STORAGE,
+				category: ErrorCategory.THIRD_PARTY,
+				details: {
+					workflowName,
+					runId
+				}
+			}, error);
+		}
+	}
 	async loadWorkflowSnapshot({ workflowName, runId }) {
 		try {
 			const result = await this.db.load({