@mastra/memory 0.10.4-alpha.0 → 0.10.4-alpha.1

package/.turbo/turbo-build.log

@@ -1,9 +1,9 @@
 
-> @mastra/memory@0.10.4-alpha.0 build /home/runner/work/mastra/mastra/packages/memory
+> @mastra/memory@0.10.4-alpha.1 build /home/runner/work/mastra/mastra/packages/memory
 > pnpm run check && tsup --silent src/index.ts src/processors/index.ts --format esm,cjs --experimental-dts --clean --treeshake=smallest --splitting
 
 
-> @mastra/memory@0.10.4-alpha.0 check /home/runner/work/mastra/mastra/packages/memory
+> @mastra/memory@0.10.4-alpha.1 check /home/runner/work/mastra/mastra/packages/memory
 > tsc --noEmit
 
 Analysis will use the bundled TypeScript version 5.8.3

package/CHANGELOG.md

@@ -1,5 +1,17 @@
 # @mastra/memory
 
+## 0.10.4-alpha.1
+
+### Patch Changes
+
+- 4051477: dependencies updates:
+  - Added dependency [`zod-to-json-schema@^3.24.5` ↗︎](https://www.npmjs.com/package/zod-to-json-schema/v/3.24.5) (to `dependencies`)
+- d70c420: fix(core, memory): fix fetchMemory regression
+- 2a16996: Working Memory Schema and Template
+- Updated dependencies [d70c420]
+- Updated dependencies [2a16996]
+  - @mastra/core@0.10.6-alpha.3
+
 ## 0.10.4-alpha.0
 
 ### Patch Changes

package/dist/_tsup-dts-rollup.d.cts

@@ -13,6 +13,8 @@ import type { StorageGetMessagesArg } from '@mastra/core/storage';
 import type { StorageThreadType } from '@mastra/core/memory';
 import type { TiktokenBPE } from 'js-tiktoken/lite';
 import type { UIMessage } from 'ai';
+import type { WorkingMemoryFormat } from '@mastra/core/memory';
+import type { WorkingMemoryTemplate } from '@mastra/core/memory';
 
 /**
  * Concrete implementation of MastraMemory that adds support for thread configuration
@@ -72,13 +74,18 @@ export declare class Memory extends MastraMemory {
     protected updateMessageToHideWorkingMemory(message: MastraMessageV1): MastraMessageV1 | null;
     protected updateMessageToHideWorkingMemoryV2(message: MastraMessageV2): MastraMessageV2 | null;
     protected parseWorkingMemory(text: string): string | null;
-    getWorkingMemory({ threadId }: {
+    getWorkingMemory({ threadId, format, }: {
         threadId: string;
+        format?: WorkingMemoryFormat;
     }): Promise<string | null>;
+    getWorkingMemoryTemplate(): Promise<WorkingMemoryTemplate | null>;
     getSystemMessage({ threadId, memoryConfig, }: {
         threadId: string;
         memoryConfig?: MemoryConfig;
     }): Promise<string | null>;
+    getUserContextMessage({ threadId }: {
+        threadId: string;
+    }): Promise<string | null>;
     defaultWorkingMemoryTemplate: string;
     private getWorkingMemoryToolInstruction;
     getTools(config?: MemoryConfig): Record<string, CoreTool>;
@@ -134,6 +141,8 @@ declare class ToolCallFilter extends MemoryProcessor_2 {
 export { ToolCallFilter }
 export { ToolCallFilter as ToolCallFilter_alias_1 }
 
-export declare const updateWorkingMemoryTool: CoreTool;
+export declare const updateWorkingMemoryTool: ({ format }: {
+    format: WorkingMemoryFormat;
+}) => CoreTool;
 
 export { }

package/dist/_tsup-dts-rollup.d.ts

@@ -13,6 +13,8 @@ import type { StorageGetMessagesArg } from '@mastra/core/storage';
 import type { StorageThreadType } from '@mastra/core/memory';
 import type { TiktokenBPE } from 'js-tiktoken/lite';
 import type { UIMessage } from 'ai';
+import type { WorkingMemoryFormat } from '@mastra/core/memory';
+import type { WorkingMemoryTemplate } from '@mastra/core/memory';
 
 /**
  * Concrete implementation of MastraMemory that adds support for thread configuration
@@ -72,13 +74,18 @@ export declare class Memory extends MastraMemory {
     protected updateMessageToHideWorkingMemory(message: MastraMessageV1): MastraMessageV1 | null;
     protected updateMessageToHideWorkingMemoryV2(message: MastraMessageV2): MastraMessageV2 | null;
     protected parseWorkingMemory(text: string): string | null;
-    getWorkingMemory({ threadId }: {
+    getWorkingMemory({ threadId, format, }: {
         threadId: string;
+        format?: WorkingMemoryFormat;
     }): Promise<string | null>;
+    getWorkingMemoryTemplate(): Promise<WorkingMemoryTemplate | null>;
     getSystemMessage({ threadId, memoryConfig, }: {
         threadId: string;
         memoryConfig?: MemoryConfig;
     }): Promise<string | null>;
+    getUserContextMessage({ threadId }: {
+        threadId: string;
+    }): Promise<string | null>;
     defaultWorkingMemoryTemplate: string;
     private getWorkingMemoryToolInstruction;
     getTools(config?: MemoryConfig): Record<string, CoreTool>;
@@ -134,6 +141,8 @@ declare class ToolCallFilter extends MemoryProcessor_2 {
 export { ToolCallFilter }
 export { ToolCallFilter as ToolCallFilter_alias_1 }
 
-export declare const updateWorkingMemoryTool: CoreTool;
+export declare const updateWorkingMemoryTool: ({ format }: {
+    format: WorkingMemoryFormat;
+}) => CoreTool;
 
 export { }

package/dist/index.cjs

@@ -5,20 +5,22 @@ var agent = require('@mastra/core/agent');
 var memory = require('@mastra/core/memory');
 var ai = require('ai');
 var xxhash = require('xxhash-wasm');
+var zodToJsonSchema = require('zod-to-json-schema');
 var zod = require('zod');
 
 function _interopDefault (e) { return e && e.__esModule ? e : { default: e }; }
 
 var xxhash__default = /*#__PURE__*/_interopDefault(xxhash);
+var zodToJsonSchema__default = /*#__PURE__*/_interopDefault(zodToJsonSchema);
 
 // src/index.ts
-var updateWorkingMemoryTool = {
+var updateWorkingMemoryTool = ({ format }) => ({
   description: "Update the working memory with new information",
   parameters: zod.z.object({
-    memory: zod.z.string().describe("The Markdown-formatted working memory content to store")
+    memory: zod.z.string().describe(`The ${format === "json" ? "JSON" : "Markdown"} formatted working memory content to store`)
   }),
   execute: async (params) => {
-    const { context, threadId, memory } = params;
+    const { context, threadId, memory, resourceId } = params;
     if (!threadId || !memory) {
       throw new Error("Thread ID and Memory instance are required for working memory updates");
     }
@@ -26,18 +28,22 @@ var updateWorkingMemoryTool = {
     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.memory;
     await memory.saveThread({
       thread: {
         ...thread,
         metadata: {
           ...thread.metadata,
-          workingMemory: context.memory
+          workingMemory
         }
       }
     });
     return { success: true };
   }
-};
+});
 
 // src/index.ts
 var CHARS_PER_TOKEN = 4;
@@ -200,10 +206,14 @@ var Memory = class extends memory.MastraMemory {
   }) {
     const config = this.getMergedThreadConfig(memoryConfig || {});
     if (config.workingMemory?.enabled && !thread?.metadata?.workingMemory) {
+      let workingMemory = config.workingMemory.template || this.defaultWorkingMemoryTemplate;
+      if (config.workingMemory.schema) {
+        workingMemory = JSON.stringify(zodToJsonSchema__default.default(config.workingMemory.schema));
+      }
       return this.storage.saveThread({
         thread: core.deepMerge(thread, {
           metadata: {
-            workingMemory: config.workingMemory.template || this.defaultWorkingMemoryTemplate
+            workingMemory
           }
         })
       });
@@ -399,12 +409,41 @@ var Memory = class extends memory.MastraMemory {
     }
     return null;
   }
-  async getWorkingMemory({ threadId }) {
-    if (!this.threadConfig.workingMemory?.enabled) return null;
+  async getWorkingMemory({
+    threadId,
+    format
+  }) {
+    if (!this.threadConfig.workingMemory?.enabled) {
+      return null;
+    }
     const thread = await this.storage.getThreadById({ threadId });
-    if (!thread) return this.threadConfig?.workingMemory?.template || this.defaultWorkingMemoryTemplate;
-    const memory = thread.metadata?.workingMemory || this.threadConfig.workingMemory.template || this.defaultWorkingMemoryTemplate;
-    return memory.trim();
+    if (format === "json") {
+      try {
+        return JSON.parse(thread?.metadata?.workingMemory) || null;
+      } catch (e) {
+        this.logger.error("Unable to parse working memory as JSON. Returning string.", e);
+      }
+    }
+    return thread?.metadata?.workingMemory ? JSON.stringify(thread?.metadata?.workingMemory) : null;
+  }
+  async getWorkingMemoryTemplate() {
+    if (!this.threadConfig.workingMemory?.enabled) {
+      return null;
+    }
+    if (this.threadConfig?.workingMemory?.schema) {
+      try {
+        const schema = this.threadConfig.workingMemory.schema;
+        const convertedSchema = zodToJsonSchema__default.default(schema, {
+          $refStrategy: "none"
+        });
+        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;
+    return { format: "markdown", content: memory.trim() };
   }
   async getSystemMessage({
     threadId,
@@ -414,11 +453,28 @@ var Memory = class extends memory.MastraMemory {
     if (!config.workingMemory?.enabled) {
       return null;
     }
+    const workingMemoryTemplate = await this.getWorkingMemoryTemplate();
+    const workingMemoryData = await this.getWorkingMemory({ threadId });
+    if (!workingMemoryTemplate) {
+      return null;
+    }
+    return this.getWorkingMemoryToolInstruction({
+      template: workingMemoryTemplate,
+      data: workingMemoryData
+    });
+  }
+  async getUserContextMessage({ threadId }) {
     const workingMemory = await this.getWorkingMemory({ threadId });
     if (!workingMemory) {
       return null;
     }
-    return this.getWorkingMemoryToolInstruction(workingMemory);
+    return `The following is the most up-to-date information about the user's state and context:
+${JSON.stringify(workingMemory)}
+Use this information as the source of truth when generating responses. 
+Do not reference or mention this memory directly to the user. 
+If conversation history shows information that is not in the working memory, use the working memory as the source of truth.
+If there is a discrepancy between this information and conversation history, always rely on this information unless the user explicitly asks for an update.
+`;
   }
   defaultWorkingMemoryTemplate = `
 # User Information
@@ -432,33 +488,44 @@ var Memory = class extends memory.MastraMemory {
 - **Facts**: 
 - **Projects**: 
 `;
-  getWorkingMemoryToolInstruction(workingMemoryBlock) {
+  getWorkingMemoryToolInstruction({
+    template,
+    data
+  }) {
     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!
 
 Guidelines:
 1. Store anything that could be useful later in the conversation
 2. Update proactively when information changes, no matter how small
-3. Use Markdown format for all data
+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"
 
-Memory Structure:
-${workingMemoryBlock}
+WORKING MEMORY TEMPLATE:
+${template.content}
+
+WORKING MEMORY DATA:
+${data}
 
 Notes:
 - Update memory whenever referenced information changes
 - 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
 - Do not remove empty sections - you must include the empty sections along with the ones you're filling in
-- REMEMBER: the way you update your working memory is by calling the updateWorkingMemory tool with the entire Markdown content. The system will store it for you. The user will not see it.
+- REMEMBER: the way you update your working memory is by calling the updateWorkingMemory tool with the entire ${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.
-- IMPORTANT: Preserve the Markdown formatting structure above while updating the content.`;
+- IMPORTANT: Preserve the ${template.format === "json" ? "JSON" : "Markdown"} formatting structure above while updating the content.`;
   }
   getTools(config) {
     const mergedConfig = this.getMergedThreadConfig(config);
     if (mergedConfig.workingMemory?.enabled) {
+      if (mergedConfig.workingMemory.schema) {
+        return {
+          updateWorkingMemory: updateWorkingMemoryTool({ format: "json" })
+        };
+      }
       return {
-        updateWorkingMemory: updateWorkingMemoryTool
+        updateWorkingMemory: updateWorkingMemoryTool({ format: "markdown" })
       };
     }
     return {};

package/dist/index.js

@@ -3,16 +3,17 @@ import { MessageList } from '@mastra/core/agent';
 import { MastraMemory } from '@mastra/core/memory';
 import { embedMany } from 'ai';
 import xxhash from 'xxhash-wasm';
+import zodToJsonSchema from 'zod-to-json-schema';
 import { z } from 'zod';
 
 // src/index.ts
-var updateWorkingMemoryTool = {
+var updateWorkingMemoryTool = ({ format }) => ({
   description: "Update the working memory with new information",
   parameters: z.object({
-    memory: z.string().describe("The Markdown-formatted working memory content to store")
+    memory: z.string().describe(`The ${format === "json" ? "JSON" : "Markdown"} formatted working memory content to store`)
   }),
   execute: async (params) => {
-    const { context, threadId, memory } = params;
+    const { context, threadId, memory, resourceId } = params;
     if (!threadId || !memory) {
       throw new Error("Thread ID and Memory instance are required for working memory updates");
     }
@@ -20,18 +21,22 @@ var updateWorkingMemoryTool = {
     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.memory;
     await memory.saveThread({
       thread: {
         ...thread,
         metadata: {
           ...thread.metadata,
-          workingMemory: context.memory
+          workingMemory
         }
       }
     });
     return { success: true };
   }
-};
+});
 
 // src/index.ts
 var CHARS_PER_TOKEN = 4;
@@ -194,10 +199,14 @@ var Memory = class extends MastraMemory {
   }) {
     const config = this.getMergedThreadConfig(memoryConfig || {});
     if (config.workingMemory?.enabled && !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: config.workingMemory.template || this.defaultWorkingMemoryTemplate
+            workingMemory
           }
         })
       });
@@ -393,12 +402,41 @@ var Memory = class extends MastraMemory {
     }
     return null;
   }
-  async getWorkingMemory({ threadId }) {
-    if (!this.threadConfig.workingMemory?.enabled) return null;
+  async getWorkingMemory({
+    threadId,
+    format
+  }) {
+    if (!this.threadConfig.workingMemory?.enabled) {
+      return null;
+    }
     const thread = await this.storage.getThreadById({ threadId });
-    if (!thread) return this.threadConfig?.workingMemory?.template || this.defaultWorkingMemoryTemplate;
-    const memory = thread.metadata?.workingMemory || this.threadConfig.workingMemory.template || this.defaultWorkingMemoryTemplate;
-    return memory.trim();
+    if (format === "json") {
+      try {
+        return JSON.parse(thread?.metadata?.workingMemory) || null;
+      } catch (e) {
+        this.logger.error("Unable to parse working memory as JSON. Returning string.", e);
+      }
+    }
+    return thread?.metadata?.workingMemory ? JSON.stringify(thread?.metadata?.workingMemory) : null;
+  }
+  async getWorkingMemoryTemplate() {
+    if (!this.threadConfig.workingMemory?.enabled) {
+      return null;
+    }
+    if (this.threadConfig?.workingMemory?.schema) {
+      try {
+        const schema = this.threadConfig.workingMemory.schema;
+        const convertedSchema = zodToJsonSchema(schema, {
+          $refStrategy: "none"
+        });
+        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;
+    return { format: "markdown", content: memory.trim() };
   }
   async getSystemMessage({
     threadId,
@@ -408,11 +446,28 @@ var Memory = class extends MastraMemory {
     if (!config.workingMemory?.enabled) {
       return null;
     }
+    const workingMemoryTemplate = await this.getWorkingMemoryTemplate();
+    const workingMemoryData = await this.getWorkingMemory({ threadId });
+    if (!workingMemoryTemplate) {
+      return null;
+    }
+    return this.getWorkingMemoryToolInstruction({
+      template: workingMemoryTemplate,
+      data: workingMemoryData
+    });
+  }
+  async getUserContextMessage({ threadId }) {
     const workingMemory = await this.getWorkingMemory({ threadId });
     if (!workingMemory) {
       return null;
     }
-    return this.getWorkingMemoryToolInstruction(workingMemory);
+    return `The following is the most up-to-date information about the user's state and context:
+${JSON.stringify(workingMemory)}
+Use this information as the source of truth when generating responses. 
+Do not reference or mention this memory directly to the user. 
+If conversation history shows information that is not in the working memory, use the working memory as the source of truth.
+If there is a discrepancy between this information and conversation history, always rely on this information unless the user explicitly asks for an update.
+`;
   }
   defaultWorkingMemoryTemplate = `
 # User Information
@@ -426,33 +481,44 @@ var Memory = class extends MastraMemory {
 - **Facts**: 
 - **Projects**: 
 `;
-  getWorkingMemoryToolInstruction(workingMemoryBlock) {
+  getWorkingMemoryToolInstruction({
+    template,
+    data
+  }) {
     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!
 
 Guidelines:
 1. Store anything that could be useful later in the conversation
 2. Update proactively when information changes, no matter how small
-3. Use Markdown format for all data
+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"
 
-Memory Structure:
-${workingMemoryBlock}
+WORKING MEMORY TEMPLATE:
+${template.content}
+
+WORKING MEMORY DATA:
+${data}
 
 Notes:
 - Update memory whenever referenced information changes
 - 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
 - Do not remove empty sections - you must include the empty sections along with the ones you're filling in
-- REMEMBER: the way you update your working memory is by calling the updateWorkingMemory tool with the entire Markdown content. The system will store it for you. The user will not see it.
+- REMEMBER: the way you update your working memory is by calling the updateWorkingMemory tool with the entire ${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.
-- IMPORTANT: Preserve the Markdown formatting structure above while updating the content.`;
+- IMPORTANT: Preserve the ${template.format === "json" ? "JSON" : "Markdown"} formatting structure above while updating the content.`;
   }
   getTools(config) {
     const mergedConfig = this.getMergedThreadConfig(config);
     if (mergedConfig.workingMemory?.enabled) {
+      if (mergedConfig.workingMemory.schema) {
+        return {
+          updateWorkingMemory: updateWorkingMemoryTool({ format: "json" })
+        };
+      }
       return {
-        updateWorkingMemory: updateWorkingMemoryTool
+        updateWorkingMemory: updateWorkingMemoryTool({ format: "markdown" })
       };
     }
     return {};

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mastra/memory",
-  "version": "0.10.4-alpha.0",
+  "version": "0.10.4-alpha.1",
   "description": "",
   "type": "module",
   "main": "./dist/index.js",
@@ -40,7 +40,8 @@
     "postgres": "^3.4.7",
     "redis": "^4.7.1",
     "xxhash-wasm": "^1.1.0",
-    "zod": "^3.25.57"
+    "zod": "^3.25.57",
+    "zod-to-json-schema": "^3.24.5"
   },
   "devDependencies": {
     "@ai-sdk/openai": "^1.3.22",
@@ -53,7 +54,7 @@
     "typescript-eslint": "^8.34.0",
     "vitest": "^3.2.3",
     "@internal/lint": "0.0.12",
-    "@mastra/core": "0.10.6-alpha.0"
+    "@mastra/core": "0.10.6-alpha.3"
   },
   "peerDependencies": {
     "@mastra/core": ">=0.10.4-0 <0.11.0"

package/src/index.ts

@@ -3,12 +3,19 @@ import type { CoreTool, MastraMessageV1 } from '@mastra/core';
 import { MessageList } from '@mastra/core/agent';
 import type { MastraMessageV2 } from '@mastra/core/agent';
 import { MastraMemory } from '@mastra/core/memory';
-import type { MemoryConfig, SharedMemoryConfig, StorageThreadType } from '@mastra/core/memory';
+import type {
+  MemoryConfig,
+  SharedMemoryConfig,
+  StorageThreadType,
+  WorkingMemoryFormat,
+  WorkingMemoryTemplate,
+} from '@mastra/core/memory';
 import type { StorageGetMessagesArg } from '@mastra/core/storage';
 import { embedMany } from 'ai';
 import type { CoreMessage, TextPart, UIMessage } from 'ai';
 
 import xxhash from 'xxhash-wasm';
+import zodToJsonSchema from 'zod-to-json-schema';
 import { updateWorkingMemoryTool } from './tools/working-memory';
 
 // Average characters per token based on OpenAI's tokenization
@@ -247,10 +254,16 @@ export class Memory extends MastraMemory {
 
     if (config.workingMemory?.enabled && !thread?.metadata?.workingMemory) {
       // if working memory is enabled but the thread doesn't have it, we need to set it
+      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: config.workingMemory.template || this.defaultWorkingMemoryTemplate,
+            workingMemory,
           },
         }),
       });
@@ -540,20 +553,54 @@ export class Memory extends MastraMemory {
     return null;
   }
 
-  public async getWorkingMemory({ threadId }: { threadId: string }): Promise<string | null> {
-    if (!this.threadConfig.workingMemory?.enabled) return null;
+  public async getWorkingMemory({
+    threadId,
+    format,
+  }: {
+    threadId: string;
+    format?: WorkingMemoryFormat;
+  }): Promise<string | null> {
+    if (!this.threadConfig.workingMemory?.enabled) {
+      return null;
+    }
 
-    // Get thread from storage
     const thread = await this.storage.getThreadById({ threadId });
-    if (!thread) return this.threadConfig?.workingMemory?.template || this.defaultWorkingMemoryTemplate;
+
+    if (format === 'json') {
+      try {
+        return JSON.parse(thread?.metadata?.workingMemory as string) || null;
+      } catch (e) {
+        this.logger.error('Unable to parse working memory as JSON. Returning string.', e);
+      }
+    }
+
+    return thread?.metadata?.workingMemory ? JSON.stringify(thread?.metadata?.workingMemory) : null;
+  }
+
+  public async getWorkingMemoryTemplate(): Promise<WorkingMemoryTemplate | null> {
+    if (!this.threadConfig.workingMemory?.enabled) {
+      return null;
+    }
+
+    // Get thread from storage
+    if (this.threadConfig?.workingMemory?.schema) {
+      try {
+        const schema = this.threadConfig.workingMemory.schema;
+        const convertedSchema = zodToJsonSchema(schema, {
+          $refStrategy: 'none',
+        });
+
+        return { format: 'json', content: JSON.stringify(convertedSchema) };
+      } catch (error) {
+        this.logger.error('Error converting schema', error);
+        throw error;
+      }
+    }
 
     // Return working memory from metadata
-    const memory =
-      (thread.metadata?.workingMemory as string) ||
-      this.threadConfig.workingMemory.template ||
-      this.defaultWorkingMemoryTemplate;
+    const memory = this.threadConfig.workingMemory.template || this.defaultWorkingMemoryTemplate;
 
-    return memory.trim();
+    return { format: 'markdown', content: memory.trim() };
   }
 
   public async getSystemMessage({
@@ -568,12 +615,32 @@ export class Memory extends MastraMemory {
       return null;
     }
 
+    const workingMemoryTemplate = await this.getWorkingMemoryTemplate();
+    const workingMemoryData = await this.getWorkingMemory({ threadId });
+
+    if (!workingMemoryTemplate) {
+      return null;
+    }
+
+    return this.getWorkingMemoryToolInstruction({
+      template: workingMemoryTemplate,
+      data: workingMemoryData,
+    });
+  }
+
+  public async getUserContextMessage({ threadId }: { threadId: string }) {
     const workingMemory = await this.getWorkingMemory({ threadId });
     if (!workingMemory) {
       return null;
     }
 
-    return this.getWorkingMemoryToolInstruction(workingMemory);
+    return `The following is the most up-to-date information about the user's state and context:
+${JSON.stringify(workingMemory)}
+Use this information as the source of truth when generating responses. 
+Do not reference or mention this memory directly to the user. 
+If conversation history shows information that is not in the working memory, use the working memory as the source of truth.
+If there is a discrepancy between this information and conversation history, always rely on this information unless the user explicitly asks for an update.
+`;
   }
 
   public defaultWorkingMemoryTemplate = `
@@ -589,34 +656,49 @@ export class Memory extends MastraMemory {
 - **Projects**: 
 `;
 
-  private getWorkingMemoryToolInstruction(workingMemoryBlock: string) {
+  private getWorkingMemoryToolInstruction({
+    template,
+    data,
+  }: {
+    template: WorkingMemoryTemplate;
+    data: string | null;
+  }) {
     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!
 
 Guidelines:
 1. Store anything that could be useful later in the conversation
 2. Update proactively when information changes, no matter how small
-3. Use Markdown format for all data
+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"
 
-Memory Structure:
-${workingMemoryBlock}
+WORKING MEMORY TEMPLATE:
+${template.content}
+
+WORKING MEMORY DATA:
+${data}
 
 Notes:
 - Update memory whenever referenced information changes
 - 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
 - Do not remove empty sections - you must include the empty sections along with the ones you're filling in
-- REMEMBER: the way you update your working memory is by calling the updateWorkingMemory tool with the entire Markdown content. The system will store it for you. The user will not see it.
+- REMEMBER: the way you update your working memory is by calling the updateWorkingMemory tool with the entire ${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.
-- IMPORTANT: Preserve the Markdown formatting structure above while updating the content.`;
+- IMPORTANT: Preserve the ${template.format === 'json' ? 'JSON' : 'Markdown'} formatting structure above while updating the content.`;
   }
 
   public getTools(config?: MemoryConfig): Record<string, CoreTool> {
     const mergedConfig = this.getMergedThreadConfig(config);
     if (mergedConfig.workingMemory?.enabled) {
+      if (mergedConfig.workingMemory.schema) {
+        return {
+          updateWorkingMemory: updateWorkingMemoryTool({ format: 'json' }),
+        };
+      }
+
       return {
-        updateWorkingMemory: updateWorkingMemoryTool,
+        updateWorkingMemory: updateWorkingMemoryTool({ format: 'markdown' }),
       };
     }
     return {};

package/src/tools/working-memory.ts

@@ -1,33 +1,43 @@
 import type { CoreTool } from '@mastra/core';
+import type { WorkingMemoryFormat } from '@mastra/core/memory';
 import { z } from 'zod';
 
-export const updateWorkingMemoryTool: CoreTool = {
+export const updateWorkingMemoryTool = ({ format }: { format: WorkingMemoryFormat }): CoreTool => ({
   description: 'Update the working memory with new information',
   parameters: z.object({
-    memory: z.string().describe('The Markdown-formatted working memory content to store'),
+    memory: z
+      .string()
+      .describe(`The ${format === 'json' ? 'JSON' : 'Markdown'} formatted working memory content to store`),
   }),
   execute: async (params: any) => {
-    const { context, threadId, memory } = 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.memory;
+
     // Update thread metadata with new working memory
     await memory.saveThread({
       thread: {
         ...thread,
         metadata: {
           ...thread.metadata,
-          workingMemory: context.memory,
+          workingMemory: workingMemory,
         },
       },
     });
 
     return { success: true };
   },
-};
+});