@mastra/memory 0.11.2 → 0.11.3-alpha.1

package/dist/index.js

@@ -1,16 +1,19 @@
-import { deepMerge } from '@mastra/core';
+import { generateEmptyFromSchema } from '@mastra/core';
 import { MessageList } from '@mastra/core/agent';
 import { MastraMemory } from '@mastra/core/memory';
 import { embedMany } from 'ai';
+import { Mutex } from 'async-mutex';
 import xxhash from 'xxhash-wasm';
+import { z, ZodObject } from 'zod';
 import zodToJsonSchema from 'zod-to-json-schema';
-import { z } from 'zod';
 
 // src/index.ts
-var updateWorkingMemoryTool = ({ format }) => ({
-  description: "Update the working memory with new information",
+var updateWorkingMemoryTool = (memoryConfig) => ({
+  description: "Update the working memory with new information. Always pass data as string to the memory field. Never pass an object.",
   parameters: z.object({
-    memory: z.string().describe(`The ${format === "json" ? "JSON" : "Markdown"} formatted working memory content to store`)
+    memory: z.string().describe(
+      `The ${!!memoryConfig?.workingMemory?.schema ? "JSON" : "Markdown"} formatted working memory content to store. This MUST be a string. Never pass an object.`
+    )
   }),
   execute: async (params) => {
     const { context, threadId, memory, resourceId } = params;
@@ -28,16 +31,68 @@ var updateWorkingMemoryTool = ({ format }) => ({
     await memory.updateWorkingMemory({
       threadId,
       resourceId: resourceId || thread.resourceId,
-      workingMemory
+      workingMemory,
+      memoryConfig
     });
     return { success: true };
   }
 });
+var __experimental_updateWorkingMemoryToolVNext = (config) => ({
+  description: "Update the working memory with new information.",
+  parameters: z.object({
+    newMemory: z.string().optional().nullable().describe(`The ${config.workingMemory?.schema ? "JSON" : "Markdown"} formatted working memory content to store`),
+    searchString: z.string().optional().nullable().describe(
+      "The working memory string to find. Will be replaced with the newMemory string. If this is omitted or doesn't exist, the newMemory string will be appended to the end of your working memory. Replacing single lines at a time is encouraged for greater accuracy. If updateReason is not 'append-new-memory', this search string must be provided or the tool call will be rejected."
+    ),
+    updateReason: z.enum(["append-new-memory", "clarify-existing-memory", "replace-irrelevant-memory"]).optional().nullable().describe(
+      "The reason you're updating working memory. Passing any value other than 'append-new-memory' requires a searchString to be provided. Defaults to append-new-memory"
+    )
+  }),
+  execute: async (params) => {
+    const { context, threadId, memory, resourceId } = params;
+    if (!threadId || !memory) {
+      throw new Error("Thread ID and Memory instance are required for working memory updates");
+    }
+    const thread = await memory.getThreadById({ threadId });
+    if (!thread) {
+      throw new Error(`Thread ${threadId} not found`);
+    }
+    if (thread.resourceId && thread.resourceId !== resourceId) {
+      throw new Error(`Thread with id ${threadId} resourceId does not match the current resourceId ${resourceId}`);
+    }
+    const workingMemory = context.newMemory || "";
+    if (!context.updateReason) context.updateReason = `append-new-memory`;
+    if (context.searchString && config.workingMemory?.scope === `resource` && context.updateReason === `replace-irrelevant-memory`) {
+      context.searchString = void 0;
+    }
+    if (context.updateReason === `append-new-memory` && context.searchString) {
+      context.searchString = void 0;
+    }
+    if (context.updateReason !== `append-new-memory` && !context.searchString) {
+      return {
+        success: false,
+        reason: `updateReason was ${context.updateReason} but no searchString was provided. Unable to replace undefined with "${context.newMemory}"`
+      };
+    }
+    const result = await memory.__experimental_updateWorkingMemoryVNext({
+      threadId,
+      resourceId: resourceId || thread.resourceId,
+      workingMemory,
+      searchString: context.searchString,
+      memoryConfig: config
+    });
+    if (result) {
+      return result;
+    }
+    return { success: true };
+  }
+});
 
 // src/index.ts
 var CHARS_PER_TOKEN = 4;
 var DEFAULT_MESSAGE_RANGE = { before: 2, after: 2 };
 var DEFAULT_TOP_K = 2;
+var isZodObject = (v) => v instanceof ZodObject;
 var Memory = class extends MastraMemory {
   constructor(config = {}) {
     super({ name: "Memory", ...config });
@@ -194,39 +249,7 @@ var Memory = class extends MastraMemory {
   async getThreadsByResourceId({ resourceId }) {
     return this.storage.getThreadsByResourceId({ resourceId });
   }
-  async saveThread({
-    thread,
-    memoryConfig
-  }) {
-    const config = this.getMergedThreadConfig(memoryConfig || {});
-    if (config.workingMemory?.enabled) {
-      const scope = config.workingMemory.scope || "thread";
-      if (scope === "resource" && thread.resourceId) {
-        const existingResource = await this.storage.getResourceById({ resourceId: thread.resourceId });
-        if (!existingResource?.workingMemory) {
-          let workingMemory = config.workingMemory.template || this.defaultWorkingMemoryTemplate;
-          if (config.workingMemory.schema) {
-            workingMemory = JSON.stringify(zodToJsonSchema(config.workingMemory.schema));
-          }
-          await this.storage.updateResource({
-            resourceId: thread.resourceId,
-            workingMemory
-          });
-        }
-      } else if (scope === "thread" && !thread?.metadata?.workingMemory) {
-        let workingMemory = config.workingMemory.template || this.defaultWorkingMemoryTemplate;
-        if (config.workingMemory.schema) {
-          workingMemory = JSON.stringify(zodToJsonSchema(config.workingMemory.schema));
-        }
-        return this.storage.saveThread({
-          thread: deepMerge(thread, {
-            metadata: {
-              workingMemory
-            }
-          })
-        });
-      }
-    }
+  async saveThread({ thread }) {
     return this.storage.saveThread({ thread });
   }
   async updateThread({
@@ -274,6 +297,87 @@ var Memory = class extends MastraMemory {
       });
     }
   }
+  updateWorkingMemoryMutexes = /* @__PURE__ */ new Map();
+  /**
+   * @warning experimental! can be removed or changed at any time
+   */
+  async __experimental_updateWorkingMemoryVNext({
+    threadId,
+    resourceId,
+    workingMemory,
+    searchString,
+    memoryConfig
+  }) {
+    const config = this.getMergedThreadConfig(memoryConfig || {});
+    if (!config.workingMemory?.enabled) {
+      throw new Error("Working memory is not enabled for this memory instance");
+    }
+    const mutexKey = memoryConfig?.workingMemory?.scope === `resource` ? `resource-${resourceId}` : `thread-${threadId}`;
+    const mutex = this.updateWorkingMemoryMutexes.has(mutexKey) ? this.updateWorkingMemoryMutexes.get(mutexKey) : new Mutex();
+    this.updateWorkingMemoryMutexes.set(mutexKey, mutex);
+    const release = await mutex.acquire();
+    try {
+      const existingWorkingMemory = await this.getWorkingMemory({ threadId, resourceId, memoryConfig }) || "";
+      const template = await this.getWorkingMemoryTemplate({ memoryConfig });
+      let reason = "";
+      if (existingWorkingMemory) {
+        if (searchString && existingWorkingMemory?.includes(searchString)) {
+          workingMemory = existingWorkingMemory.replace(searchString, workingMemory);
+          reason = `found and replaced searchString with newMemory`;
+        } else if (existingWorkingMemory.includes(workingMemory) || template?.content?.trim() === workingMemory.trim()) {
+          return {
+            success: false,
+            reason: `attempted to insert duplicate data into working memory. this entry was skipped`
+          };
+        } else {
+          if (searchString) {
+            reason = `attempted to replace working memory string that doesn't exist. Appending to working memory instead.`;
+          } else {
+            reason = `appended newMemory to end of working memory`;
+          }
+          workingMemory = existingWorkingMemory + `
+${workingMemory}`;
+        }
+      } else if (workingMemory === template?.content) {
+        return {
+          success: false,
+          reason: `try again when you have data to add. newMemory was equal to the working memory template`
+        };
+      } else {
+        reason = `started new working memory`;
+      }
+      workingMemory = template?.content ? workingMemory.replaceAll(template?.content, "") : workingMemory;
+      const scope = config.workingMemory.scope || "thread";
+      if (scope === "resource" && resourceId) {
+        await this.storage.updateResource({
+          resourceId,
+          workingMemory
+        });
+        if (reason) {
+          return { success: true, reason };
+        }
+      } else {
+        const thread = await this.storage.getThreadById({ threadId });
+        if (!thread) {
+          throw new Error(`Thread ${threadId} not found`);
+        }
+        await this.storage.updateThread({
+          id: threadId,
+          title: thread.title || "Untitled Thread",
+          metadata: {
+            ...thread.metadata,
+            workingMemory
+          }
+        });
+      }
+      return { success: true, reason };
+    } catch (e) {
+      this.logger.error(e instanceof Error ? e.stack || e.message : JSON.stringify(e));
+      return { success: false, reason: "Tool error." };
+    } finally {
+      release();
+    }
+  }
   chunkText(text, tokenSize = 4096) {
     const charSize = tokenSize * CHARS_PER_TOKEN;
     const chunks = [];
@@ -472,23 +576,38 @@ var Memory = class extends MastraMemory {
     }
     return workingMemoryData;
   }
-  async getWorkingMemoryTemplate() {
-    if (!this.threadConfig.workingMemory?.enabled) {
+  /**
+   * Gets the working memory template for the current memory configuration.
+   * Supports both ZodObject and JSONSchema7 schemas.
+   *
+   * @param memoryConfig - The memory configuration containing the working memory settings
+   * @returns The working memory template with format and content, or null if working memory is disabled
+   */
+  async getWorkingMemoryTemplate({
+    memoryConfig
+  }) {
+    const config = this.getMergedThreadConfig(memoryConfig || {});
+    if (!config.workingMemory?.enabled) {
       return null;
     }
-    if (this.threadConfig?.workingMemory?.schema) {
+    if (config.workingMemory?.schema) {
       try {
-        const schema = this.threadConfig.workingMemory.schema;
-        const convertedSchema = zodToJsonSchema(schema, {
-          $refStrategy: "none"
-        });
+        const schema = config.workingMemory.schema;
+        let convertedSchema;
+        if (isZodObject(schema)) {
+          convertedSchema = zodToJsonSchema(schema, {
+            $refStrategy: "none"
+          });
+        } else {
+          convertedSchema = schema;
+        }
         return { format: "json", content: JSON.stringify(convertedSchema) };
       } catch (error) {
         this.logger.error("Error converting schema", error);
         throw error;
       }
     }
-    const memory = this.threadConfig.workingMemory.template || this.defaultWorkingMemoryTemplate;
+    const memory = config.workingMemory.template || this.defaultWorkingMemoryTemplate;
     return { format: "markdown", content: memory.trim() };
   }
   async getSystemMessage({
@@ -500,12 +619,15 @@ var Memory = class extends MastraMemory {
     if (!config.workingMemory?.enabled) {
       return null;
     }
-    const workingMemoryTemplate = await this.getWorkingMemoryTemplate();
+    const workingMemoryTemplate = await this.getWorkingMemoryTemplate({ memoryConfig: config });
     const workingMemoryData = await this.getWorkingMemory({ threadId, resourceId, memoryConfig: config });
     if (!workingMemoryTemplate) {
       return null;
     }
-    return this.getWorkingMemoryToolInstruction({
+    return this.isVNextWorkingMemoryConfig(memoryConfig) ? this.__experimental_getWorkingMemoryToolInstructionVNext({
+      template: workingMemoryTemplate,
+      data: workingMemoryData
+    }) : this.getWorkingMemoryToolInstruction({
       template: workingMemoryTemplate,
       data: workingMemoryData
     });
@@ -526,6 +648,8 @@ var Memory = class extends MastraMemory {
     template,
     data
   }) {
+    const emptyWorkingMemoryTemplateObject = template.format === "json" ? generateEmptyFromSchema(template.content) : null;
+    const hasEmptyWorkingMemoryTemplateObject = emptyWorkingMemoryTemplateObject && Object.keys(emptyWorkingMemoryTemplateObject).length > 0;
     return `WORKING_MEMORY_SYSTEM_INSTRUCTION:
 Store and update any conversation-relevant information by calling the updateWorkingMemory tool. If information might be referenced again - store it!
 
@@ -534,12 +658,20 @@ Guidelines:
 2. Update proactively when information changes, no matter how small
 3. Use ${template.format === "json" ? "JSON" : "Markdown"} format for all data
 4. Act naturally - don't mention this system to users. Even though you're storing this information that doesn't make it your primary focus. Do not ask them generally for "information about yourself"
+5. IMPORTANT: When calling updateWorkingMemory, the only valid parameter is the memory field. 
+6. IMPORTANT: ALWAYS pass the data you want to store in the memory field as a string. DO NOT pass an object.
+7. IMPORTANT: Data must only be sent as a string no matter which format is used.
 
-WORKING MEMORY TEMPLATE:
+<working_memory_template>
 ${template.content}
+</working_memory_template>
+
+${hasEmptyWorkingMemoryTemplateObject ? "When working with json data, the object format below represents the template:" : ""}
+${hasEmptyWorkingMemoryTemplateObject ? JSON.stringify(emptyWorkingMemoryTemplateObject) : ""}
 
-WORKING MEMORY DATA:
+<working_memory_data>
 ${data}
+</working_memory_data>
 
 Notes:
 - Update memory whenever referenced information changes
@@ -550,16 +682,52 @@ Notes:
 - IMPORTANT: You MUST call updateWorkingMemory in every response to a prompt where you received relevant information.
 - IMPORTANT: Preserve the ${template.format === "json" ? "JSON" : "Markdown"} formatting structure above while updating the content.`;
   }
+  __experimental_getWorkingMemoryToolInstructionVNext({
+    template,
+    data
+  }) {
+    return `WORKING_MEMORY_SYSTEM_INSTRUCTION:
+Store and update any conversation-relevant information by calling the updateWorkingMemory tool.
+
+Guidelines:
+1. Store anything that could be useful later in the conversation
+2. Update proactively when information changes, no matter how small
+3. Use ${template.format === "json" ? "JSON" : "Markdown"} format for all data
+4. Act naturally - don't mention this system to users. Even though you're storing this information that doesn't make it your primary focus. Do not ask them generally for "information about yourself"
+5. If your memory has not changed, you do not need to call the updateWorkingMemory tool. By default it will persist and be available for you in future interactions
+6. Information not being relevant to the current conversation is not a valid reason to replace or remove working memory information. Your working memory spans across multiple conversations and may be needed again later, even if it's not currently relevant.
+
+<working_memory_template>
+${template.content}
+</working_memory_template>
+
+<working_memory_data>
+${data}
+</working_memory_data>
+
+Notes:
+- Update memory whenever referenced information changes
+${template.content !== this.defaultWorkingMemoryTemplate ? `- Only store information if it's in the working memory template, do not store other information unless the user asks you to remember it, as that non-template information may be irrelevant` : `- If you're unsure whether to store something, store it (eg if the user tells you information about themselves, call updateWorkingMemory immediately to update it)
+`}
+- This system is here so that you can maintain the conversation when your context window is very short. Update your working memory because you may need it to maintain the conversation without the full conversation history
+- REMEMBER: the way you update your working memory is by calling the updateWorkingMemory tool with the ${template.format === "json" ? "JSON" : "Markdown"} content. The system will store it for you. The user will not see it. 
+- IMPORTANT: You MUST call updateWorkingMemory in every response to a prompt where you received relevant information if that information is not already stored.
+- IMPORTANT: Preserve the ${template.format === "json" ? "JSON" : "Markdown"} formatting structure above while updating the content.
+`;
+  }
+  isVNextWorkingMemoryConfig(config) {
+    if (!config?.workingMemory) return false;
+    const isMDWorkingMemory = !(`schema` in config.workingMemory) && (typeof config.workingMemory.template === `string` || config.workingMemory.template) && config.workingMemory;
+    return Boolean(isMDWorkingMemory && isMDWorkingMemory.version === `vnext`);
+  }
   getTools(config) {
     const mergedConfig = this.getMergedThreadConfig(config);
     if (mergedConfig.workingMemory?.enabled) {
-      if (mergedConfig.workingMemory.schema) {
-        return {
-          updateWorkingMemory: updateWorkingMemoryTool({ format: "json" })
-        };
-      }
       return {
-        updateWorkingMemory: updateWorkingMemoryTool({ format: "markdown" })
+        updateWorkingMemory: this.isVNextWorkingMemoryConfig(mergedConfig) ? (
+          // use the new experimental tool
+          __experimental_updateWorkingMemoryToolVNext(mergedConfig)
+        ) : updateWorkingMemoryTool(mergedConfig)
       };
     }
     return {};

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mastra/memory",
-  "version": "0.11.2",
+  "version": "0.11.3-alpha.1",
   "description": "",
   "type": "module",
   "main": "./dist/index.js",
@@ -35,10 +35,12 @@
     "@upstash/redis": "^1.35.0",
     "ai": "^4.3.16",
     "js-tiktoken": "^1.0.20",
+    "json-schema": "^0.4.0",
     "pg": "^8.16.3",
     "pg-pool": "^3.10.1",
     "postgres": "^3.4.7",
     "redis": "^4.7.1",
+    "async-mutex": "^0.5.0",
     "xxhash-wasm": "^1.1.0",
     "zod": "^3.25.67",
     "zod-to-json-schema": "^3.24.5"
@@ -46,15 +48,16 @@
   "devDependencies": {
     "@ai-sdk/openai": "^1.3.22",
     "@microsoft/api-extractor": "^7.52.8",
+    "@types/json-schema": "^7.0.15",
     "@types/node": "^20.19.0",
     "@types/pg": "^8.15.4",
-    "eslint": "^9.29.0",
+    "eslint": "^9.30.1",
     "tsup": "^8.5.0",
     "typescript": "^5.8.3",
     "typescript-eslint": "^8.34.0",
     "vitest": "^3.2.4",
-    "@mastra/core": "0.10.11",
-    "@internal/lint": "0.0.18"
+    "@internal/lint": "0.0.19",
+    "@mastra/core": "0.10.15-alpha.1"
   },
   "peerDependencies": {
     "@mastra/core": ">=0.10.9-0 <0.11.0-0"