@mastra/mssql 1.10.5 → 1.10.6-alpha.1

package/dist/index.js

@@ -5,7 +5,7 @@ import { parseSqlIdentifier } from "@mastra/core/utils";
 import { MastraBase } from "@mastra/core/base";
 import { MessageList } from "@mastra/core/agent";
 import { saveScorePayloadSchema } from "@mastra/core/evals";
-//#region src/storage/db/utils.ts
+// src/storage/db/utils.ts
 function getSchemaName$1(schema) {
 	return schema ? `[${parseSqlIdentifier(schema, "schema name")}]` : void 0;
 }
@@ -14,8 +14,7 @@ function getTableName$1({ indexName, schemaName }) {
 	const quotedSchemaName = schemaName;
 	return quotedSchemaName ? `${quotedSchemaName}.${quotedIndexName}` : quotedIndexName;
 }
-//#endregion
-//#region src/storage/db/index.ts
+// src/storage/db/index.ts
 /**
 * Resolves MssqlDomainConfig to pool and schema.
 * Domain classes create their own MssqlDB instance from the returned pool.
@@ -1033,8 +1032,7 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
 		}
 	}
 };
-//#endregion
-//#region src/storage/domains/utils.ts
+// src/storage/domains/utils.ts
 function getSchemaName(schema) {
 	return schema ? `[${parseSqlIdentifier(schema, "schema name")}]` : void 0;
 }
@@ -1156,8 +1154,7 @@ function transformFromSqlRow({ tableName, sqlRow }) {
 	});
 	return result;
 }
-//#endregion
-//#region src/storage/domains/agents/index.ts
+// src/storage/domains/agents/index.ts
 var AgentsMSSQL = class AgentsMSSQL extends AgentsStorage {
 	pool;
 	schema;
@@ -1745,8 +1742,7 @@ var AgentsMSSQL = class AgentsMSSQL extends AgentsStorage {
 		}
 	}
 };
-//#endregion
-//#region src/storage/domains/background-tasks/index.ts
+// src/storage/domains/background-tasks/index.ts
 function serializeJson(v) {
 	if (typeof v === "object" && v != null) return JSON.stringify(v);
 	return v ?? null;
@@ -2107,8 +2103,7 @@ var BackgroundTasksMSSQL = class BackgroundTasksMSSQL extends BackgroundTasksSto
 		return Number(result.recordset[0]?.count ?? 0);
 	}
 };
-//#endregion
-//#region src/storage/db/run-fencing.ts
+// src/storage/db/run-fencing.ts
 /**
 * Epoch milliseconds on the database clock. SQL Server evaluates
 * `SYSUTCDATETIME()` once when a statement starts, so a statement that waits
@@ -2167,16 +2162,23 @@ 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;
 	}
 }
-//#endregion
-//#region src/storage/domains/memory/index.ts
+// src/storage/domains/memory/index.ts
 function bindMssqlMetadataParams(request, params) {
 	for (const [paramName, paramValue] of Object.entries(params)) request.input(paramName, paramValue);
 }
@@ -3336,8 +3338,7 @@ var MemoryMSSQL = class MemoryMSSQL extends MemoryStorage {
 		}
 	}
 };
-//#endregion
-//#region src/storage/domains/observability/index.ts
+// src/storage/domains/observability/index.ts
 var ObservabilityMSSQL = class ObservabilityMSSQL extends ObservabilityStorage {
 	static retentionTables = { spans: {
 		table: TABLE_SPANS,
@@ -3978,8 +3979,7 @@ var ObservabilityMSSQL = class ObservabilityMSSQL extends ObservabilityStorage {
 		}
 	}
 };
-//#endregion
-//#region src/storage/domains/scores/index.ts
+// src/storage/domains/scores/index.ts
 /**
 * MSSQL-specific score row transformation.
 * Converts timestamp strings to Date objects.
@@ -4395,8 +4395,7 @@ var ScoresMSSQL = class ScoresMSSQL extends ScoresStorage {
 		}
 	}
 };
-//#endregion
-//#region src/storage/domains/workflow-definitions/index.ts
+// src/storage/domains/workflow-definitions/index.ts
 function parseJson(value, column, rowId) {
 	if (value == null) return null;
 	if (typeof value === "string") try {
@@ -4599,8 +4598,7 @@ var WorkflowDefinitionsMSSQL = class WorkflowDefinitionsMSSQL extends WorkflowDe
 		await request.query(`DELETE FROM ${tableName} WHERE [id] = @id`);
 	}
 };
-//#endregion
-//#region src/storage/domains/workflows/index.ts
+// 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);
@@ -4646,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;
 	}
@@ -4955,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);
@@ -4970,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({
@@ -4994,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({
@@ -5137,8 +5162,7 @@ var WorkflowsMSSQL = class WorkflowsMSSQL extends WorkflowsStorage {
 		}
 	}
 };
-//#endregion
-//#region src/storage/index.ts
+// src/storage/index.ts
 /**
 * Type guard for pre-configured pool config
 */
@@ -5251,12 +5275,8 @@ var MSSQLStore = class extends MastraCompositeStore {
 		}
 	}
 	async _performInitializationAndStore() {
-		try {
-			await this.pool.connect();
-			return true;
-		} catch (err) {
-			throw err;
-		}
+		await this.pool.connect();
+		return true;
 	}
 	/**
 	* Closes the MSSQL connection pool.
@@ -5267,7 +5287,6 @@ var MSSQLStore = class extends MastraCompositeStore {
 		await this.pool.close();
 	}
 };
-//#endregion
 export { AgentsMSSQL, BackgroundTasksMSSQL, MSSQLStore, MemoryMSSQL, ObservabilityMSSQL, ScoresMSSQL, WorkflowDefinitionsMSSQL, WorkflowsMSSQL };
 
 //# sourceMappingURL=index.js.map