@mastra/pg 1.7.0 → 1.7.1

package/dist/index.js

@@ -1439,11 +1439,9 @@ var PgVector = class extends MastraVector {
             WHERE c.relname = $1
             AND n.nspname = $2;
             `;
-      const [dimResult, countResult, indexResult] = await Promise.all([
-        client.query(dimensionQuery, [tableName]),
-        client.query(countQuery),
-        client.query(indexQuery, [`${indexName}_vector_idx`, this.schema || "public"])
-      ]);
+      const dimResult = await client.query(dimensionQuery, [tableName]);
+      const countResult = await client.query(countQuery);
+      const indexResult = await client.query(indexQuery, [`${indexName}_vector_idx`, this.schema || "public"]);
       const { index_method, index_def, operator_class } = indexResult.rows[0] || {
         index_method: "flat",
         index_def: "",
@@ -6723,7 +6721,8 @@ var MemoryPG = class _MemoryPG extends MemoryStorage {
           "isBufferingObservation",
           "isBufferingReflection",
           "lastBufferedAtTokens",
-          "lastBufferedAtTime"
+          "lastBufferedAtTime",
+          "metadata"
         ]
       });
     }
@@ -7897,9 +7896,11 @@ var MemoryPG = class _MemoryPG extends MemoryStorage {
           ]
         );
         const clonedMessages = [];
+        const messageIdMap = {};
         const targetResourceId = resourceId || sourceThread.resourceId;
         for (const sourceMsg of sourceMessages) {
           const newMessageId = crypto.randomUUID();
+          messageIdMap[sourceMsg.id] = newMessageId;
           const normalizedMsg = this.normalizeMessageRow(sourceMsg);
           let parsedContent = normalizedMsg.content;
           try {
@@ -7932,7 +7933,8 @@ var MemoryPG = class _MemoryPG extends MemoryStorage {
         }
         return {
           thread: newThread,
-          clonedMessages
+          clonedMessages,
+          messageIdMap
         };
       });
     } catch (error) {
@@ -8142,6 +8144,80 @@ var MemoryPG = class _MemoryPG extends MemoryStorage {
       );
     }
   }
+  async insertObservationalMemoryRecord(record) {
+    try {
+      const lookupKey = this.getOMKey(record.threadId, record.resourceId);
+      const tableName = getTableName3({
+        indexName: OM_TABLE,
+        schemaName: getSchemaName3(this.#schema)
+      });
+      const lastObservedAtStr = record.lastObservedAt ? record.lastObservedAt.toISOString() : null;
+      const lastBufferedAtTimeStr = record.lastBufferedAtTime ? record.lastBufferedAtTime.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",
+          "observedMessageIds", "bufferedObservationChunks",
+          "bufferedReflection", "bufferedReflectionTokens", "bufferedReflectionInputTokens",
+          "reflectedObservationLineCount",
+          "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, $30, $31, $32, $33, $34, $35)`,
+        [
+          record.id,
+          lookupKey,
+          record.scope,
+          record.resourceId,
+          record.threadId || null,
+          record.activeObservations || "",
+          null,
+          record.originType || "initial",
+          record.config ? JSON.stringify(record.config) : null,
+          record.generationCount || 0,
+          lastObservedAtStr,
+          lastObservedAtStr,
+          null,
+          // lastReflectionAt
+          null,
+          // lastReflectionAtZ
+          record.pendingMessageTokens || 0,
+          record.totalTokensObserved || 0,
+          record.observationTokenCount || 0,
+          record.observedMessageIds ? JSON.stringify(record.observedMessageIds) : null,
+          record.bufferedObservationChunks ? JSON.stringify(record.bufferedObservationChunks) : null,
+          record.bufferedReflection || null,
+          record.bufferedReflectionTokens ?? null,
+          record.bufferedReflectionInputTokens ?? null,
+          record.reflectedObservationLineCount ?? null,
+          record.isObserving || false,
+          record.isReflecting || false,
+          record.isBufferingObservation || false,
+          record.isBufferingReflection || false,
+          record.lastBufferedAtTokens || 0,
+          lastBufferedAtTimeStr,
+          record.observedTimezone || null,
+          record.metadata ? JSON.stringify(record.metadata) : null,
+          record.createdAt.toISOString(),
+          record.createdAt.toISOString(),
+          record.updatedAt.toISOString(),
+          record.updatedAt.toISOString()
+        ]
+      );
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: createStorageErrorId("PG", "INSERT_OBSERVATIONAL_MEMORY_RECORD", "FAILED"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { id: record.id, threadId: record.threadId, resourceId: record.resourceId }
+        },
+        error
+      );
+    }
+  }
   async updateActiveObservations(input) {
     try {
       const now = /* @__PURE__ */ new Date();
@@ -8242,8 +8318,8 @@ var MemoryPG = class _MemoryPG extends MemoryStorage {
           "originType", config, "generationCount", "lastObservedAt", "lastObservedAtZ", "lastReflectionAt", "lastReflectionAtZ",
           "pendingMessageTokens", "totalTokensObserved", "observationTokenCount",
           "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)`,
+          "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,
@@ -8279,6 +8355,7 @@ var MemoryPG = class _MemoryPG extends MemoryStorage {
           null,
           // lastBufferedAtTime
           record.observedTimezone || null,
+          record.metadata ? JSON.stringify(record.metadata) : null,
           nowStr,
           // createdAt
           nowStr,
@@ -11889,6 +11966,9 @@ var WorkflowsPG = class _WorkflowsPG extends WorkflowsStorage {
     this.#skipDefaultIndexes = skipDefaultIndexes;
     this.#indexes = indexes?.filter((idx) => _WorkflowsPG.MANAGED_TABLES.includes(idx.table));
   }
+  supportsConcurrentUpdates() {
+    return true;
+  }
   parseWorkflowRun(row) {
     let parsedSnapshot = row.snapshot;
     if (typeof parsedSnapshot === "string") {
@@ -11967,21 +12047,111 @@ var WorkflowsPG = class _WorkflowsPG extends WorkflowsStorage {
   async dangerouslyClearAll() {
     await this.#db.clearTable({ tableName: TABLE_WORKFLOW_SNAPSHOT });
   }
-  updateWorkflowResults({
-    // workflowName,
-    // runId,
-    // stepId,
-    // result,
-    // requestContext,
+  async updateWorkflowResults({
+    workflowName,
+    runId,
+    stepId,
+    result,
+    requestContext
   }) {
-    throw new Error("Method not implemented.");
+    try {
+      return await this.#db.client.tx(async (t) => {
+        const tableName = getTableName5({ indexName: TABLE_WORKFLOW_SNAPSHOT, schemaName: getSchemaName5(this.#schema) });
+        const existingSnapshotResult = await t.oneOrNone(
+          `SELECT snapshot FROM ${tableName} WHERE workflow_name = $1 AND run_id = $2 FOR UPDATE`,
+          [workflowName, runId]
+        );
+        let snapshot;
+        if (!existingSnapshotResult) {
+          snapshot = {
+            context: {},
+            activePaths: [],
+            timestamp: Date.now(),
+            suspendedPaths: {},
+            activeStepsPath: {},
+            resumeLabels: {},
+            serializedStepGraph: [],
+            status: "pending",
+            value: {},
+            waitingPaths: {},
+            runId,
+            requestContext: {}
+          };
+        } else {
+          const existingSnapshot = existingSnapshotResult.snapshot;
+          snapshot = typeof existingSnapshot === "string" ? JSON.parse(existingSnapshot) : existingSnapshot;
+        }
+        snapshot.context[stepId] = result;
+        snapshot.requestContext = { ...snapshot.requestContext, ...requestContext };
+        const now = /* @__PURE__ */ new Date();
+        const sanitizedSnapshot = sanitizeJsonForPg(JSON.stringify(snapshot));
+        await t.none(
+          `INSERT INTO ${tableName} (workflow_name, run_id, snapshot, "createdAt", "updatedAt")
+           VALUES ($1, $2, $3, $4, $5)
+           ON CONFLICT (workflow_name, run_id) DO UPDATE
+           SET snapshot = $3, "updatedAt" = $5`,
+          [workflowName, runId, sanitizedSnapshot, now, now]
+        );
+        return snapshot.context;
+      });
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: createStorageErrorId("PG", "UPDATE_WORKFLOW_RESULTS", "FAILED"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            workflowName,
+            runId,
+            stepId
+          }
+        },
+        error
+      );
+    }
   }
-  updateWorkflowState({
-    // workflowName,
-    // runId,
-    // opts,
+  async updateWorkflowState({
+    workflowName,
+    runId,
+    opts
   }) {
-    throw new Error("Method not implemented.");
+    try {
+      return await this.#db.client.tx(async (t) => {
+        const tableName = getTableName5({ indexName: TABLE_WORKFLOW_SNAPSHOT, schemaName: getSchemaName5(this.#schema) });
+        const existingSnapshotResult = await t.oneOrNone(
+          `SELECT snapshot FROM ${tableName} WHERE workflow_name = $1 AND run_id = $2 FOR UPDATE`,
+          [workflowName, runId]
+        );
+        if (!existingSnapshotResult) {
+          return void 0;
+        }
+        const existingSnapshot = existingSnapshotResult.snapshot;
+        const snapshot = typeof existingSnapshot === "string" ? JSON.parse(existingSnapshot) : existingSnapshot;
+        if (!snapshot || !snapshot?.context) {
+          throw new Error(`Snapshot not found for runId ${runId}`);
+        }
+        const updatedSnapshot = { ...snapshot, ...opts };
+        const sanitizedSnapshot = sanitizeJsonForPg(JSON.stringify(updatedSnapshot));
+        await t.none(
+          `UPDATE ${tableName} SET snapshot = $1, "updatedAt" = $2 WHERE workflow_name = $3 AND run_id = $4`,
+          [sanitizedSnapshot, /* @__PURE__ */ new Date(), workflowName, runId]
+        );
+        return updatedSnapshot;
+      });
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: createStorageErrorId("PG", "UPDATE_WORKFLOW_STATE", "FAILED"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            workflowName,
+            runId
+          }
+        },
+        error
+      );
+    }
   }
   async persistWorkflowSnapshot({
     workflowName,