@mastra/memory 0.10.1 → 0.10.2-alpha.1

package/src/index.ts

@@ -1,14 +1,15 @@
 import { deepMerge } from '@mastra/core';
-import type { AiMessageType, CoreMessage, CoreTool } from '@mastra/core';
+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 { MessageType, MemoryConfig, SharedMemoryConfig, StorageThreadType } from '@mastra/core/memory';
+import type { MemoryConfig, SharedMemoryConfig, StorageThreadType } from '@mastra/core/memory';
 import type { StorageGetMessagesArg } from '@mastra/core/storage';
 import { embedMany } from 'ai';
-import type { TextPart } from 'ai';
+import type { TextPart, UIMessage } from 'ai';
 
 import xxhash from 'xxhash-wasm';
 import { updateWorkingMemoryTool } from './tools/working-memory';
-import { reorderToolCallsAndResults } from './utils';
 
 // Average characters per token based on OpenAI's tokenization
 const CHARS_PER_TOKEN = 4;
@@ -55,7 +56,7 @@ export class Memory extends MastraMemory {
     threadConfig,
   }: StorageGetMessagesArg & {
     threadConfig?: MemoryConfig;
-  }): Promise<{ messages: CoreMessage[]; uiMessages: AiMessageType[] }> {
+  }): Promise<{ messages: MastraMessageV1[]; uiMessages: UIMessage[] }> {
     if (resourceId) await this.validateThreadIsOwnedByResource(threadId, resourceId);
 
     const vectorResults: {
@@ -137,16 +138,28 @@ export class Memory extends MastraMemory {
       threadConfig: config,
     });
 
-    // First sort messages by date
     const orderedByDate = rawMessages.sort((a, b) => new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime());
-    // Then reorder tool calls to be directly before their results
-    const reorderedToolCalls = reorderToolCallsAndResults(orderedByDate);
 
-    // Parse and convert messages
-    const messages = this.parseMessages(reorderedToolCalls);
-    const uiMessages = this.convertToUIMessages(reorderedToolCalls);
-
-    return { messages, uiMessages };
+    const list = new MessageList({ threadId, resourceId }).add(orderedByDate, 'memory');
+    return {
+      get messages() {
+        // returning v1 messages for backwards compat! v1 messages were CoreMessages stored in the db.
+        // returning .v1() takes stored messages which may be in v2 or v1 format and converts them to v1 shape, which is a CoreMessage + id + threadId + resourceId, etc
+        // Perhaps this should be called coreRecord or something ? - for now keeping v1 since it reflects that this used to be our db storage record shape
+        const v1Messages = list.get.all.v1();
+        // the conversion from V2/UIMessage -> V1/CoreMessage can sometimes split the messages up into more messages than before
+        // so slice off the earlier messages if it'll exceed the lastMessages setting
+        if (selectBy?.last && v1Messages.length > selectBy.last) {
+          // ex: 23 (v1 messages) minus 20 (selectBy.last messages)
+          // means we will start from index 3 and keep all the later newer messages from index 3 til the end of the array
+          return v1Messages.slice(v1Messages.length - selectBy.last);
+        }
+        return v1Messages;
+      },
+      get uiMessages() {
+        return list.get.all.ui();
+      },
+    };
   }
 
   async rememberMessages({
@@ -159,19 +172,14 @@ export class Memory extends MastraMemory {
     resourceId?: string;
     vectorMessageSearch?: string;
     config?: MemoryConfig;
-  }): Promise<{
-    threadId: string;
-    messages: CoreMessage[];
-    uiMessages: AiMessageType[];
-  }> {
+  }): Promise<{ messages: MastraMessageV1[]; messagesV2: MastraMessageV2[] }> {
     if (resourceId) await this.validateThreadIsOwnedByResource(threadId, resourceId);
     const threadConfig = this.getMergedThreadConfig(config || {});
 
     if (!threadConfig.lastMessages && !threadConfig.semanticRecall) {
       return {
         messages: [],
-        uiMessages: [],
-        threadId,
+        messagesV2: [],
       };
     }
 
@@ -183,13 +191,11 @@ export class Memory extends MastraMemory {
       },
       threadConfig: config,
     });
+    // Using MessageList here just to convert mixed input messages to single type output messages
+    const list = new MessageList({ threadId, resourceId }).add(messagesResult.messages, 'memory');
 
     this.logger.debug(`Remembered message history includes ${messagesResult.messages.length} messages.`);
-    return {
-      threadId,
-      messages: messagesResult.messages,
-      uiMessages: messagesResult.uiMessages,
-    };
+    return { messages: list.get.all.v1(), messagesV2: list.get.all.mastra() };
   }
 
   async getThreadById({ threadId }: { threadId: string }): Promise<StorageThreadType | null> {
@@ -324,14 +330,20 @@ export class Memory extends MastraMemory {
     messages,
     memoryConfig,
   }: {
-    messages: MessageType[];
+    messages: (MastraMessageV1 | MastraMessageV2)[];
     memoryConfig?: MemoryConfig;
-  }): Promise<MessageType[]> {
-    // First save working memory from any messages
-    await this.saveWorkingMemory(messages);
-
+  }): Promise<MastraMessageV2[]> {
     // Then strip working memory tags from all messages
-    const updatedMessages = this.updateMessagesToHideWorkingMemory(messages);
+    const updatedMessages = messages
+      .map(m => {
+        if (MessageList.isMastraMessageV1(m)) {
+          return this.updateMessageToHideWorkingMemory(m);
+        }
+        // add this to prevent "error saving undefined in the db" if a project is on an earlier storage version but new memory/storage
+        if (!m.type) m.type = `v2`;
+        return this.updateMessageToHideWorkingMemoryV2(m);
+      })
+      .filter((m): m is MastraMessageV1 | MastraMessageV2 => Boolean(m));
 
     const config = this.getMergedThreadConfig(memoryConfig);
 
@@ -343,16 +355,34 @@ export class Memory extends MastraMemory {
         updatedMessages.map(async message => {
           let textForEmbedding: string | null = null;
 
-          if (typeof message.content === 'string' && message.content.trim() !== '') {
-            textForEmbedding = message.content;
-          } else if (Array.isArray(message.content)) {
-            // Extract text from all text parts, concatenate
-            const joined = message.content
-              .filter(part => part && part.type === 'text' && typeof part.text === 'string')
-              .map(part => (part as TextPart).text)
-              .join(' ')
-              .trim();
-            if (joined) textForEmbedding = joined;
+          if (MessageList.isMastraMessageV2(message)) {
+            if (
+              message.content.content &&
+              typeof message.content.content === 'string' &&
+              message.content.content.trim() !== ''
+            ) {
+              textForEmbedding = message.content.content;
+            } else if (message.content.parts && message.content.parts.length > 0) {
+              // Extract text from all text parts, concatenate
+              const joined = message.content.parts
+                .filter(part => part.type === 'text')
+                .map(part => (part as TextPart).text)
+                .join(' ')
+                .trim();
+              if (joined) textForEmbedding = joined;
+            }
+          } else if (MessageList.isMastraMessageV1(message)) {
+            if (message.content && typeof message.content === 'string' && message.content.trim() !== '') {
+              textForEmbedding = message.content;
+            } else if (message.content && Array.isArray(message.content) && message.content.length > 0) {
+              // Extract text from all text parts, concatenate
+              const joined = message.content
+                .filter(part => part.type === 'text')
+                .map(part => part.text)
+                .join(' ')
+                .trim();
+              if (joined) textForEmbedding = joined;
+            }
           }
 
           if (!textForEmbedding) return;
@@ -384,47 +414,70 @@ export class Memory extends MastraMemory {
 
     return result;
   }
+  protected updateMessageToHideWorkingMemory(message: MastraMessageV1): MastraMessageV1 | null {
+    const workingMemoryRegex = /<working_memory>([^]*?)<\/working_memory>/g;
 
-  protected updateMessagesToHideWorkingMemory(messages: MessageType[]): MessageType[] {
+    if (typeof message?.content === `string`) {
+      return {
+        ...message,
+        content: message.content.replace(workingMemoryRegex, ``).trim(),
+      };
+    } else if (Array.isArray(message?.content)) {
+      // Filter out updateWorkingMemory tool-call/result content items
+      const filteredContent = message.content.filter(
+        content =>
+          (content.type !== 'tool-call' && content.type !== 'tool-result') ||
+          content.toolName !== 'updateWorkingMemory',
+      );
+      const newContent = filteredContent.map(content => {
+        if (content.type === 'text') {
+          return {
+            ...content,
+            text: content.text.replace(workingMemoryRegex, '').trim(),
+          };
+        }
+        return { ...content };
+      }) as MastraMessageV1['content'];
+      if (!newContent.length) return null;
+      return { ...message, content: newContent };
+    } else {
+      return { ...message };
+    }
+  }
+  protected updateMessageToHideWorkingMemoryV2(message: MastraMessageV2): MastraMessageV2 | null {
     const workingMemoryRegex = /<working_memory>([^]*?)<\/working_memory>/g;
 
-    const updatedMessages: MessageType[] = [];
+    const newMessage = { ...message, content: { ...message.content } }; // Deep copy message and content
 
-    for (const message of messages) {
-      if (typeof message?.content === `string`) {
-        updatedMessages.push({
-          ...message,
-          content: message.content.replace(workingMemoryRegex, ``).trim(),
-        });
-      } else if (Array.isArray(message?.content)) {
-        // Filter out updateWorkingMemory tool-call/result content items
-        const filteredContent = message.content.filter(
-          content =>
-            !(
-              (content.type === 'tool-call' || content.type === 'tool-result') &&
-              content.toolName === 'updateWorkingMemory'
-            ),
-        );
-        if (filteredContent.length === 0) {
-          // If nothing left, skip this message
-          continue;
-        }
-        const newContent = filteredContent.map(content => {
-          if (content.type === 'text') {
+    if (newMessage.content.content && typeof newMessage.content.content === 'string') {
+      newMessage.content.content = newMessage.content.content.replace(workingMemoryRegex, '').trim();
+    }
+
+    if (newMessage.content.parts) {
+      newMessage.content.parts = newMessage.content.parts
+        .filter(part => {
+          if (part.type === 'tool-invocation') {
+            return part.toolInvocation.toolName !== 'updateWorkingMemory';
+          }
+          return true;
+        })
+        .map(part => {
+          if (part.type === 'text') {
             return {
-              ...content,
-              text: content.text.replace(workingMemoryRegex, '').trim(),
+              ...part,
+              text: part.text.replace(workingMemoryRegex, '').trim(),
             };
           }
-          return { ...content };
-        }) as MessageType['content'];
-        updatedMessages.push({ ...message, content: newContent });
-      } else {
-        updatedMessages.push({ ...message });
+          return part;
+        });
+
+      // If all parts were filtered out (e.g., only contained updateWorkingMemory tool calls) we need to skip the whole message, it was only working memory tool calls/results
+      if (newMessage.content.parts.length === 0) {
+        return null;
       }
     }
 
-    return updatedMessages;
+    return newMessage;
   }
 
   protected parseWorkingMemory(text: string): string | null {
@@ -457,46 +510,6 @@ export class Memory extends MastraMemory {
     return memory.trim();
   }
 
-  private async saveWorkingMemory(messages: MessageType[]) {
-    const latestMessage = messages[messages.length - 1];
-
-    if (!latestMessage || !this.threadConfig.workingMemory?.enabled) {
-      return;
-    }
-
-    const latestContent = !latestMessage?.content
-      ? null
-      : typeof latestMessage.content === 'string'
-        ? latestMessage.content
-        : latestMessage.content
-            .filter(c => c.type === 'text')
-            .map(c => c.text)
-            .join('\n');
-
-    const threadId = latestMessage?.threadId;
-    if (!latestContent || !threadId) {
-      return;
-    }
-
-    const newMemory = this.parseWorkingMemory(latestContent);
-    if (!newMemory) {
-      return;
-    }
-
-    const thread = await this.storage.getThreadById({ threadId });
-    if (!thread) return;
-
-    // Update thread metadata with new working memory
-    await this.storage.updateThread({
-      id: thread.id,
-      title: thread.title || '',
-      metadata: deepMerge(thread.metadata || {}, {
-        workingMemory: newMemory,
-      }),
-    });
-    return newMemory;
-  }
-
   public async getSystemMessage({
     threadId,
     memoryConfig,
@@ -530,30 +543,6 @@ export class Memory extends MastraMemory {
 - **Projects**: 
 `;
 
-  private getWorkingMemoryWithInstruction(workingMemoryBlock: string) {
-    return `WORKING_MEMORY_SYSTEM_INSTRUCTION:
-Store and update any conversation-relevant information by including "<working_memory>text</working_memory>" in your responses. Updates replace existing memory while maintaining this structure. 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 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:
-<working_memory>
-${workingMemoryBlock}
-</working_memory>
-
-Notes:
-- Update memory whenever referenced information changes
-- If you're unsure whether to store something, store it (eg if the user tells you their name or other information, output the <working_memory> block 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 outputting the entire "<working_memory>text</working_memory>" block in your response. The system will pick this up and store it for you. The user will not see it.
-- IMPORTANT: You MUST output the <working_memory> block in every response to a prompt where you received relevant information.
-- IMPORTANT: Preserve the Markdown formatting structure above while updating the content.`;
-  }
-
   private getWorkingMemoryToolInstruction(workingMemoryBlock: string) {
     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!

package/src/processors/index.test.ts

@@ -1,6 +1,6 @@
 import { openai } from '@ai-sdk/openai';
 import { createTool } from '@mastra/core';
-import type { CoreMessage, MessageType } from '@mastra/core';
+import type { MessageType } from '@mastra/core';
 import { Agent } from '@mastra/core/agent';
 import cl100k_base from 'js-tiktoken/ranks/cl100k_base';
 import { describe, it, expect, vi } from 'vitest';
@@ -26,8 +26,8 @@ describe('TokenLimiter', () => {
 
     // Should prioritize newest messages (higher ids)
     expect(result.length).toBe(2);
-    expect((result[0] as MessageType).id).toBe('message-8');
-    expect((result[1] as MessageType).id).toBe('message-9');
+    expect(result[0].id).toBe('message-8');
+    expect(result[1].id).toBe('message-9');
   });
 
   it('should handle empty messages array', () => {
@@ -52,8 +52,8 @@ describe('TokenLimiter', () => {
     });
 
     // All should process messages successfully but potentially with different token counts
-    const defaultResult = defaultLimiter.process(messages as CoreMessage[]);
-    const customResult = customLimiter.process(messages as CoreMessage[]);
+    const defaultResult = defaultLimiter.process(messages);
+    const customResult = customLimiter.process(messages);
 
     // Each should return the same messages but with potentially different token counts
     expect(defaultResult.length).toBe(messages.length);
@@ -69,7 +69,7 @@ describe('TokenLimiter', () => {
     // Count tokens for each message including all overheads
     for (const message of messages) {
       // Base token count from the countTokens method
-      estimatedTokens += testLimiter.countTokens(message as CoreMessage);
+      estimatedTokens += testLimiter.countTokens(message);
     }
 
     return Number(estimatedTokens.toFixed(2));
@@ -85,7 +85,7 @@ describe('TokenLimiter', () => {
     const { messages, counts } = generateConversationHistory(config);
 
     const estimate = estimateTokens(messages);
-    const used = (await agent.generate(messages.slice(0, -1) as CoreMessage[])).usage.totalTokens;
+    const used = (await agent.generate(messages.slice(0, -1))).usage.totalTokens;
 
     console.log(`Estimated ${estimate} tokens, used ${used} tokens.\n`, counts);
 
@@ -199,7 +199,7 @@ describe.concurrent('ToolCallFilter', () => {
       messageCount: 1,
     });
     const filter = new ToolCallFilter();
-    const result = filter.process(messages as CoreMessage[]) as MessageType[];
+    const result = filter.process(messages) as MessageType[];
 
     // Should only keep the text message and assistant res
     expect(result.length).toBe(2);
@@ -213,7 +213,7 @@ describe.concurrent('ToolCallFilter', () => {
       messageCount: 2,
     });
     const filter = new ToolCallFilter({ exclude: ['weather'] });
-    const result = filter.process(messages as CoreMessage[]) as MessageType[];
+    const result = filter.process(messages);
 
     // Should keep text message, assistant reply, calculator tool call, and calculator result
     expect(result.length).toBe(4);
@@ -230,7 +230,7 @@ describe.concurrent('ToolCallFilter', () => {
     });
 
     const filter = new ToolCallFilter({ exclude: [] });
-    const result = filter.process(messages as CoreMessage[]);
+    const result = filter.process(messages);
 
     // Should keep all messages
     expect(result.length).toBe(messages.length);

package/vitest.config.ts

@@ -4,5 +4,8 @@ export default defineConfig({
   test: {
     environment: 'node',
     include: ['src/**/*.test.ts'],
+    // smaller output to save token space when LLMs run tests
+    reporters: 'dot',
+    bail: 1,
   },
 });

package/src/utils/index.ts

@@ -1,88 +0,0 @@
-import type { MessageType } from '@mastra/core/memory';
-
-const isToolCallWithId = (message: MessageType | undefined, targetToolCallId: string): boolean => {
-  if (!message || !Array.isArray(message.content)) return false;
-  return message.content.some(
-    part =>
-      part &&
-      typeof part === 'object' &&
-      'type' in part &&
-      part.type === 'tool-call' &&
-      'toolCallId' in part &&
-      part.toolCallId === targetToolCallId,
-  );
-};
-
-const getToolResultIndexById = (id: string, results: MessageType[]) =>
-  results.findIndex(message => {
-    if (!Array.isArray(message?.content)) return false;
-    return message.content.some(
-      part =>
-        part &&
-        typeof part === 'object' &&
-        'type' in part &&
-        part.type === 'tool-result' &&
-        'toolCallId' in part &&
-        part.toolCallId === id,
-    );
-  });
-
-/**
- * Self-heals message ordering to ensure tool calls are directly before their matching tool results.
- * This is needed due to a bug where messages were saved in the wrong order. That bug is fixed, but this code ensures any tool calls saved in the wrong order in the past will still be usable now.
- */
-export function reorderToolCallsAndResults(messages: MessageType[]): MessageType[] {
-  if (!messages.length) return messages;
-
-  // Create a copy of messages to avoid modifying the original
-  const results = [...messages];
-
-  const toolCallIds = new Set<string>();
-
-  // First loop: collect all tool result IDs in a set
-  for (const message of results) {
-    if (!Array.isArray(message.content)) continue;
-
-    for (const part of message.content) {
-      if (
-        part &&
-        typeof part === 'object' &&
-        'type' in part &&
-        part.type === 'tool-result' &&
-        'toolCallId' in part &&
-        part.toolCallId
-      ) {
-        toolCallIds.add(part.toolCallId);
-      }
-    }
-  }
-
-  // Second loop: for each tool ID, ensure tool calls come before tool results
-  for (const toolCallId of toolCallIds) {
-    // Find tool result index
-    const resultIndex = getToolResultIndexById(toolCallId, results);
-
-    // Check if tool call is at resultIndex - 1
-    const oneMessagePrev = results[resultIndex - 1];
-    if (isToolCallWithId(oneMessagePrev, toolCallId)) {
-      continue; // Tool call is already in the correct position
-    }
-
-    // Find the tool call anywhere in the array
-    const toolCallIndex = results.findIndex(message => isToolCallWithId(message, toolCallId));
-
-    if (toolCallIndex !== -1 && toolCallIndex !== resultIndex - 1) {
-      // Store the tool call message
-      const toolCall = results[toolCallIndex];
-      if (!toolCall) continue;
-
-      // Remove the tool call from its current position
-      results.splice(toolCallIndex, 1);
-
-      // Insert right before the tool result
-      results.splice(getToolResultIndexById(toolCallId, results), 0, toolCall);
-    }
-  }
-
-  return results;
-}