@mastra/pg 1.27.1 → 1.27.2-alpha.0

package/dist/index.js

@@ -10349,6 +10349,12 @@ var MCPServersPG = class MCPServersPG extends MCPServersStorage {
 * versions that don't export TABLE_OBSERVATIONAL_MEMORY.
 */
 const OM_TABLE = "mastra_observational_memory";
+/**
+* Newest generation first. Databases written before generation creation was
+* serialized can hold several rows with the same generation; the earliest-created
+* one wins so the active record stays stable across reads.
+*/
+const OM_GENERATION_ORDER = `"generationCount" DESC, "createdAt" ASC, id ASC`;
 const POSTGRES_MAX_BIND_PARAMETERS = 65535;
 const MESSAGE_INSERT_BIND_PARAMETERS = 8;
 const MAX_MESSAGES_PER_INSERT = Math.floor(POSTGRES_MAX_BIND_PARAMETERS / MESSAGE_INSERT_BIND_PARAMETERS);
@@ -11950,6 +11956,23 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 	getOMKey(threadId, resourceId) {
 		return threadId ? `thread:${threadId}` : `resource:${resourceId}`;
 	}
+	/**
+	* Runs `fn` in a transaction that holds an advisory lock for one OM lookup key.
+	*
+	* The table has no unique constraint on ("lookupKey", "generationCount"), and
+	* adding one would fail on databases that already contain duplicates. The lock
+	* serializes generation creation across processes instead. It is
+	* transaction-scoped so it also works behind transaction-pooling proxies.
+	*/
+	async #withOMLookupKeyLock(tableName, lookupKey, fn) {
+		return this.#db.client.tx(async (t) => {
+			await t.query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))", [`${tableName}:${lookupKey}`]);
+			return fn(t);
+		});
+	}
+	async #getLatestOMRow(client, tableName, lookupKey) {
+		return client.oneOrNone(`SELECT * FROM ${tableName} WHERE "lookupKey" = $1 ORDER BY ${OM_GENERATION_ORDER} LIMIT 1`, [lookupKey]);
+	}
 	parseOMRow(row) {
 		return {
 			id: row.id,
@@ -11992,7 +12015,7 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 				indexName: OM_TABLE,
 				schemaName: getSchemaName$3(this.#schema)
 			});
-			const result = await this.#db.readClient.oneOrNone(`SELECT * FROM ${tableName} WHERE "lookupKey" = $1 ORDER BY "generationCount" DESC LIMIT 1`, [lookupKey]);
+			const result = await this.#getLatestOMRow(this.#db.readClient, tableName, lookupKey);
 			if (!result) return null;
 			return this.parseOMRow(result);
 		} catch (error) {
@@ -12028,7 +12051,7 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 				paramIndex++;
 			}
 			params.push(limit);
-			let sql = `SELECT * FROM ${tableName} WHERE ${conditions.join(" AND ")} ORDER BY "generationCount" DESC LIMIT $${paramIndex}`;
+			let sql = `SELECT * FROM ${tableName} WHERE ${conditions.join(" AND ")} ORDER BY ${OM_GENERATION_ORDER} LIMIT $${paramIndex}`;
 			paramIndex++;
 			if (options?.offset != null) {
 				params.push(options.offset);
@@ -12083,7 +12106,10 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 				schemaName: getSchemaName$3(this.#schema)
 			});
 			const nowStr = now.toISOString();
-			await this.#db.client.none(`INSERT INTO ${tableName} (
+			return await this.#withOMLookupKeyLock(tableName, lookupKey, async (t) => {
+				const existing = await this.#getLatestOMRow(t, tableName, lookupKey);
+				if (existing) return this.parseOMRow(existing);
+				await t.none(`INSERT INTO ${tableName} (
           id, "lookupKey", scope, "resourceId", "threadId",
           "activeObservations", "activeObservationsPendingUpdate",
           "originType", config, "generationCount", "lastObservedAt", "lastObservedAtZ", "lastReflectionAt", "lastReflectionAtZ",
@@ -12091,36 +12117,37 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
           "isObserving", "isReflecting", "isBufferingObservation", "isBufferingReflection", "lastBufferedAtTokens", "lastBufferedAtTime",
           "observedTimezone", "createdAt", "createdAtZ", "updatedAt", "updatedAtZ"
         ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26, $27, $28)`, [
-				id,
-				lookupKey,
-				input.scope,
-				input.resourceId,
-				input.threadId || null,
-				"",
-				null,
-				"initial",
-				toPgJson(input.config),
-				0,
-				null,
-				null,
-				null,
-				null,
-				0,
-				0,
-				0,
-				false,
-				false,
-				false,
-				false,
-				0,
-				null,
-				input.observedTimezone || null,
-				nowStr,
-				nowStr,
-				nowStr,
-				nowStr
-			]);
-			return record;
+					id,
+					lookupKey,
+					input.scope,
+					input.resourceId,
+					input.threadId || null,
+					"",
+					null,
+					"initial",
+					toPgJson(input.config),
+					0,
+					null,
+					null,
+					null,
+					null,
+					0,
+					0,
+					0,
+					false,
+					false,
+					false,
+					false,
+					0,
+					null,
+					input.observedTimezone || null,
+					nowStr,
+					nowStr,
+					nowStr,
+					nowStr
+				]);
+				return record;
+			});
 		} catch (error) {
 			throw new MastraError({
 				id: createStorageErrorId("PG", "INITIALIZE_OBSERVATIONAL_MEMORY", "FAILED"),
@@ -12252,78 +12279,16 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 	}
 	async createReflectionGeneration(input) {
 		try {
-			const id = crypto.randomUUID();
-			const now = /* @__PURE__ */ new Date();
 			const lookupKey = this.getOMKey(input.currentRecord.threadId, input.currentRecord.resourceId);
-			const record = {
-				id,
-				scope: input.currentRecord.scope,
-				threadId: input.currentRecord.threadId,
-				resourceId: input.currentRecord.resourceId,
-				createdAt: now,
-				updatedAt: now,
-				lastObservedAt: input.currentRecord.lastObservedAt,
-				originType: "reflection",
-				generationCount: input.currentRecord.generationCount + 1,
-				activeObservations: input.reflection,
-				totalTokensObserved: input.currentRecord.totalTokensObserved,
-				observationTokenCount: input.tokenCount,
-				pendingMessageTokens: 0,
-				isReflecting: false,
-				isObserving: false,
-				isBufferingObservation: false,
-				isBufferingReflection: false,
-				lastBufferedAtTokens: 0,
-				lastBufferedAtTime: null,
-				config: input.currentRecord.config,
-				metadata: input.currentRecord.metadata,
-				observedTimezone: input.currentRecord.observedTimezone
-			};
 			const tableName = getTableName$3({
 				indexName: OM_TABLE,
 				schemaName: getSchemaName$3(this.#schema)
 			});
-			const nowStr = now.toISOString();
-			const lastObservedAtStr = record.lastObservedAt?.toISOString() || null;
-			await this.#db.client.none(`INSERT INTO ${tableName} (
-          id, "lookupKey", scope, "resourceId", "threadId",
-          "activeObservations", "activeObservationsPendingUpdate",
-          "originType", config, "generationCount", "lastObservedAt", "lastObservedAtZ", "lastReflectionAt", "lastReflectionAtZ",
-          "pendingMessageTokens", "totalTokensObserved", "observationTokenCount",
-          "isObserving", "isReflecting", "isBufferingObservation", "isBufferingReflection", "lastBufferedAtTokens", "lastBufferedAtTime",
-          "observedTimezone", metadata, "createdAt", "createdAtZ", "updatedAt", "updatedAtZ"
-        ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26, $27, $28, $29)`, [
-				id,
-				lookupKey,
-				record.scope,
-				record.resourceId,
-				record.threadId || null,
-				input.reflection,
-				null,
-				"reflection",
-				toPgJson(record.config),
-				input.currentRecord.generationCount + 1,
-				lastObservedAtStr,
-				lastObservedAtStr,
-				nowStr,
-				nowStr,
-				record.pendingMessageTokens,
-				Math.round(record.totalTokensObserved),
-				Math.round(record.observationTokenCount),
-				false,
-				false,
-				false,
-				false,
-				0,
-				null,
-				record.observedTimezone || null,
-				record.metadata ? toPgJson(record.metadata) : null,
-				nowStr,
-				nowStr,
-				nowStr,
-				nowStr
-			]);
-			return record;
+			return await this.#withOMLookupKeyLock(tableName, lookupKey, async (t) => {
+				const newer = await this.#getNewerOMGeneration(t, tableName, lookupKey, input.currentRecord.generationCount);
+				if (newer) return newer;
+				return this.#insertReflectionGeneration(t, tableName, lookupKey, input);
+			});
 		} catch (error) {
 			throw new MastraError({
 				id: createStorageErrorId("PG", "CREATE_REFLECTION_GENERATION", "FAILED"),
@@ -12333,6 +12298,84 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 			}, error);
 		}
 	}
+	/**
+	* Returns the active record when another writer already created a generation
+	* after `generationCount`. Reflecting from the older generation again would
+	* add a second record with the same generation.
+	*/
+	async #getNewerOMGeneration(t, tableName, lookupKey, generationCount) {
+		const latest = await this.#getLatestOMRow(t, tableName, lookupKey);
+		return latest && Number(latest.generationCount) > generationCount ? this.parseOMRow(latest) : null;
+	}
+	async #insertReflectionGeneration(t, tableName, lookupKey, input) {
+		const id = crypto.randomUUID();
+		const now = /* @__PURE__ */ new Date();
+		const record = {
+			id,
+			scope: input.currentRecord.scope,
+			threadId: input.currentRecord.threadId,
+			resourceId: input.currentRecord.resourceId,
+			createdAt: now,
+			updatedAt: now,
+			lastObservedAt: input.currentRecord.lastObservedAt,
+			originType: "reflection",
+			generationCount: input.currentRecord.generationCount + 1,
+			activeObservations: input.reflection,
+			totalTokensObserved: input.currentRecord.totalTokensObserved,
+			observationTokenCount: input.tokenCount,
+			pendingMessageTokens: 0,
+			isReflecting: false,
+			isObserving: false,
+			isBufferingObservation: false,
+			isBufferingReflection: false,
+			lastBufferedAtTokens: 0,
+			lastBufferedAtTime: null,
+			config: input.currentRecord.config,
+			metadata: input.currentRecord.metadata,
+			observedTimezone: input.currentRecord.observedTimezone
+		};
+		const nowStr = now.toISOString();
+		const lastObservedAtStr = record.lastObservedAt?.toISOString() || null;
+		await t.none(`INSERT INTO ${tableName} (
+        id, "lookupKey", scope, "resourceId", "threadId",
+        "activeObservations", "activeObservationsPendingUpdate",
+        "originType", config, "generationCount", "lastObservedAt", "lastObservedAtZ", "lastReflectionAt", "lastReflectionAtZ",
+        "pendingMessageTokens", "totalTokensObserved", "observationTokenCount",
+        "isObserving", "isReflecting", "isBufferingObservation", "isBufferingReflection", "lastBufferedAtTokens", "lastBufferedAtTime",
+        "observedTimezone", metadata, "createdAt", "createdAtZ", "updatedAt", "updatedAtZ"
+      ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26, $27, $28, $29)`, [
+			id,
+			lookupKey,
+			record.scope,
+			record.resourceId,
+			record.threadId || null,
+			input.reflection,
+			null,
+			"reflection",
+			toPgJson(record.config),
+			input.currentRecord.generationCount + 1,
+			lastObservedAtStr,
+			lastObservedAtStr,
+			nowStr,
+			nowStr,
+			record.pendingMessageTokens,
+			Math.round(record.totalTokensObserved),
+			Math.round(record.observationTokenCount),
+			false,
+			false,
+			false,
+			false,
+			0,
+			null,
+			record.observedTimezone || null,
+			record.metadata ? toPgJson(record.metadata) : null,
+			nowStr,
+			nowStr,
+			nowStr,
+			nowStr
+		]);
+		return record;
+	}
 	async setReflectingFlag(id, isReflecting) {
 		try {
 			const tableName = getTableName$3({
@@ -12827,32 +12870,36 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 				indexName: OM_TABLE,
 				schemaName: getSchemaName$3(this.#schema)
 			});
-			const record = await this.#db.client.oneOrNone(`SELECT * FROM ${tableName} WHERE id = $1`, [input.currentRecord.id]);
-			if (!record) throw new MastraError({
-				id: createStorageErrorId("PG", "SWAP_BUFFERED_REFLECTION_TO_ACTIVE", "NOT_FOUND"),
-				text: `Observational memory record not found: ${input.currentRecord.id}`,
-				domain: ErrorDomain.STORAGE,
-				category: ErrorCategory.THIRD_PARTY,
-				details: { id: input.currentRecord.id }
-			});
-			const bufferedReflection = record.bufferedReflection || "";
-			const reflectedLineCount = Number(record.reflectedObservationLineCount || 0);
-			if (!bufferedReflection) throw new MastraError({
-				id: createStorageErrorId("PG", "SWAP_BUFFERED_REFLECTION_TO_ACTIVE", "NO_CONTENT"),
-				text: "No buffered reflection to swap",
-				domain: ErrorDomain.STORAGE,
-				category: ErrorCategory.USER,
-				details: { id: input.currentRecord.id }
-			});
-			const unreflectedContent = (record.activeObservations || "").split("\n").slice(reflectedLineCount).join("\n").trim();
-			const newObservations = unreflectedContent ? `${bufferedReflection}\n\n${unreflectedContent}` : bufferedReflection;
-			const newRecord = await this.createReflectionGeneration({
-				currentRecord: input.currentRecord,
-				reflection: newObservations,
-				tokenCount: input.tokenCount
-			});
-			const nowStr = (/* @__PURE__ */ new Date()).toISOString();
-			await this.#db.client.query(`UPDATE ${tableName} SET
+			const lookupKey = this.getOMKey(input.currentRecord.threadId, input.currentRecord.resourceId);
+			return await this.#withOMLookupKeyLock(tableName, lookupKey, async (t) => {
+				const newer = await this.#getNewerOMGeneration(t, tableName, lookupKey, input.currentRecord.generationCount);
+				if (newer) return newer;
+				const record = await t.oneOrNone(`SELECT * FROM ${tableName} WHERE id = $1`, [input.currentRecord.id]);
+				if (!record) throw new MastraError({
+					id: createStorageErrorId("PG", "SWAP_BUFFERED_REFLECTION_TO_ACTIVE", "NOT_FOUND"),
+					text: `Observational memory record not found: ${input.currentRecord.id}`,
+					domain: ErrorDomain.STORAGE,
+					category: ErrorCategory.THIRD_PARTY,
+					details: { id: input.currentRecord.id }
+				});
+				const bufferedReflection = record.bufferedReflection || "";
+				const reflectedLineCount = Number(record.reflectedObservationLineCount || 0);
+				if (!bufferedReflection) throw new MastraError({
+					id: createStorageErrorId("PG", "SWAP_BUFFERED_REFLECTION_TO_ACTIVE", "NO_CONTENT"),
+					text: "No buffered reflection to swap",
+					domain: ErrorDomain.STORAGE,
+					category: ErrorCategory.USER,
+					details: { id: input.currentRecord.id }
+				});
+				const unreflectedContent = (record.activeObservations || "").split("\n").slice(reflectedLineCount).join("\n").trim();
+				const newObservations = unreflectedContent ? `${bufferedReflection}\n\n${unreflectedContent}` : bufferedReflection;
+				const newRecord = await this.#insertReflectionGeneration(t, tableName, lookupKey, {
+					currentRecord: input.currentRecord,
+					reflection: newObservations,
+					tokenCount: input.tokenCount
+				});
+				const nowStr = (/* @__PURE__ */ new Date()).toISOString();
+				await t.query(`UPDATE ${tableName} SET
           "bufferedReflection" = NULL,
           "bufferedReflectionTokens" = NULL,
           "bufferedReflectionInputTokens" = NULL,
@@ -12860,11 +12907,12 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
           "updatedAt" = $1,
           "updatedAtZ" = $2
         WHERE id = $3`, [
-				nowStr,
-				nowStr,
-				input.currentRecord.id
-			]);
-			return newRecord;
+					nowStr,
+					nowStr,
+					input.currentRecord.id
+				]);
+				return newRecord;
+			});
 		} catch (error) {
 			if (error instanceof MastraError) throw error;
 			throw new MastraError({