@mastra/memory 0.0.0-a2a-20250421213654

package/src/processors/index.test.ts

@@ -0,0 +1,238 @@
+import { openai } from '@ai-sdk/openai';
+import { createTool } from '@mastra/core';
+import type { CoreMessage, 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';
+import { z } from 'zod';
+import { generateConversationHistory } from '../../integration-tests/src/test-utils';
+import { TokenLimiter, ToolCallFilter } from './index';
+
+vi.setConfig({ testTimeout: 20_000, hookTimeout: 20_000 });
+
+describe('TokenLimiter', () => {
+  it('should limit messages to the specified token count', () => {
+    // Create messages with predictable token counts (approximately 25 tokens each)
+    const { messages } = generateConversationHistory({
+      threadId: '1',
+      messageCount: 5,
+      toolNames: [],
+      toolFrequency: 0,
+    });
+
+    const limiter = new TokenLimiter(200);
+    // @ts-ignore
+    const result = limiter.process(messages);
+
+    // 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');
+  });
+
+  it('should handle empty messages array', () => {
+    const limiter = new TokenLimiter(1000);
+    const result = limiter.process([]);
+    expect(result).toEqual([]);
+  });
+
+  it('should use different encodings based on configuration', () => {
+    const { messages } = generateConversationHistory({
+      threadId: '6',
+      messageCount: 1,
+      toolNames: [],
+      toolFrequency: 0,
+    });
+
+    // Create limiters with different encoding settings
+    const defaultLimiter = new TokenLimiter(1000);
+    const customLimiter = new TokenLimiter({
+      limit: 1000,
+      encoding: cl100k_base,
+    });
+
+    // 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[]);
+
+    // Each should return the same messages but with potentially different token counts
+    expect(defaultResult.length).toBe(messages.length);
+    expect(customResult.length).toBe(messages.length);
+  });
+
+  function estimateTokens(messages: MessageType[]) {
+    // Create a TokenLimiter just for counting tokens
+    const testLimiter = new TokenLimiter(Infinity);
+
+    let estimatedTokens = testLimiter.TOKENS_PER_CONVERSATION;
+
+    // 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);
+    }
+
+    return estimatedTokens;
+  }
+
+  function percentDifference(a: number, b: number) {
+    const difference = Number(((Math.abs(a - b) / b) * 100).toFixed(2));
+    console.log(`${a} and ${b} are ${difference}% different`);
+    return difference;
+  }
+
+  async function expectTokenEstimate(config: Parameters<typeof generateConversationHistory>[0], agent: Agent) {
+    const { messages, counts } = generateConversationHistory(config);
+
+    const estimate = estimateTokens(messages);
+    const used = (await agent.generate(messages.slice(0, -1) as CoreMessage[])).usage.totalTokens;
+
+    console.log(`Estimated ${estimate} tokens, used ${used} tokens.\n`, counts);
+
+    // Check if within 2% margin
+    expect(percentDifference(estimate, used)).toBeLessThanOrEqual(2);
+  }
+
+  const calculatorTool = createTool({
+    id: 'calculator',
+    description: 'Perform a simple calculation',
+    inputSchema: z.object({
+      expression: z.string().describe('The mathematical expression to calculate'),
+    }),
+    execute: async ({ context: { expression } }) => {
+      return `The result of ${expression} is ${eval(expression)}`;
+    },
+  });
+
+  const agent = new Agent({
+    name: 'token estimate agent',
+    model: openai('gpt-4o-mini'),
+    instructions: ``,
+    tools: { calculatorTool },
+  });
+
+  describe.concurrent(`98% accuracy`, () => {
+    it(`20 messages, no tools`, async () => {
+      await expectTokenEstimate(
+        {
+          messageCount: 10,
+          toolFrequency: 0,
+          threadId: '2',
+        },
+        agent,
+      );
+    });
+
+    it(`60 messages, no tools`, async () => {
+      await expectTokenEstimate(
+        {
+          messageCount: 30,
+          toolFrequency: 0,
+          threadId: '3',
+        },
+        agent,
+      );
+    });
+
+    it(`4 messages, 0 tools`, async () => {
+      await expectTokenEstimate(
+        {
+          messageCount: 2,
+          toolFrequency: 0,
+          threadId: '3',
+        },
+        agent,
+      );
+    });
+
+    it(`20 messages, 2 tool messages`, async () => {
+      await expectTokenEstimate(
+        {
+          messageCount: 10,
+          toolFrequency: 5,
+          threadId: '3',
+        },
+        agent,
+      );
+    });
+
+    it(`40 messages, 6 tool messages`, async () => {
+      await expectTokenEstimate(
+        {
+          messageCount: 20,
+          toolFrequency: 5,
+          threadId: '4',
+        },
+        agent,
+      );
+    });
+
+    it(`100 messages, 24 tool messages`, async () => {
+      await expectTokenEstimate(
+        {
+          messageCount: 50,
+          toolFrequency: 4,
+          threadId: '5',
+        },
+        agent,
+      );
+    });
+
+    it(`101 messages, 49 tool calls`, async () => {
+      await expectTokenEstimate(
+        {
+          messageCount: 50,
+          toolFrequency: 1,
+          threadId: '5',
+        },
+        agent,
+      );
+    });
+  });
+});
+
+describe.concurrent('ToolCallFilter', () => {
+  it('should exclude all tool calls when created with no arguments', () => {
+    const { messages } = generateConversationHistory({
+      threadId: '3',
+      toolNames: ['weather', 'calculator', 'search'],
+      messageCount: 1,
+    });
+    const filter = new ToolCallFilter();
+    const result = filter.process(messages as CoreMessage[]) as MessageType[];
+
+    // Should only keep the text message and assistant res
+    expect(result.length).toBe(2);
+    expect(result[0].id).toBe('message-0');
+  });
+
+  it('should exclude specific tool calls by name', () => {
+    const { messages } = generateConversationHistory({
+      threadId: '4',
+      toolNames: ['weather', 'calculator'],
+      messageCount: 2,
+    });
+    const filter = new ToolCallFilter({ exclude: ['weather'] });
+    const result = filter.process(messages as CoreMessage[]) as MessageType[];
+
+    // Should keep text message, assistant reply, calculator tool call, and calculator result
+    expect(result.length).toBe(4);
+    expect(result[0].id).toBe('message-0');
+    expect(result[1].id).toBe('message-1');
+    expect(result[2].id).toBe('message-2');
+    expect(result[3].id).toBe('message-3');
+  });
+
+  it('should keep all messages when exclude list is empty', () => {
+    const { messages } = generateConversationHistory({
+      threadId: '5',
+      toolNames: ['weather', 'calculator'],
+    });
+
+    const filter = new ToolCallFilter({ exclude: [] });
+    const result = filter.process(messages as CoreMessage[]);
+
+    // Should keep all messages
+    expect(result.length).toBe(messages.length);
+  });
+});

package/src/processors/index.ts

@@ -0,0 +1,2 @@
+export * from './token-limiter';
+export * from './tool-call-filter';

package/src/processors/token-limiter.ts

@@ -0,0 +1,150 @@
+import type { CoreMessage, MemoryProcessorOpts } from '@mastra/core';
+import { MemoryProcessor } from '@mastra/core/memory';
+import { Tiktoken } from 'js-tiktoken/lite';
+import type { TiktokenBPE } from 'js-tiktoken/lite';
+import o200k_base from 'js-tiktoken/ranks/o200k_base';
+
+/**
+ * Configuration options for TokenLimiter
+ */
+interface TokenLimiterOptions {
+  /** Maximum number of tokens to allow */
+  limit: number;
+  /** Optional encoding to use (defaults to o200k_base which is used by gpt-4o) */
+  encoding?: TiktokenBPE;
+}
+
+/**
+ * Limits the total number of tokens in the messages.
+ * Uses js-tiktoken with o200k_base encoding by default for accurate token counting with modern models.
+ */
+export class TokenLimiter extends MemoryProcessor {
+  private encoder: Tiktoken;
+  private maxTokens: number;
+
+  // Token overheads per OpenAI's documentation
+  // See: https://cookbook.openai.com/examples/how_to_count_tokens_with_tiktoken#6-counting-tokens-for-chat-completions-api-calls
+  // Every message follows <|start|>{role/name}\n{content}<|end|>
+  public TOKENS_PER_MESSAGE = 3; // tokens added for each message (start & end tokens)
+  public TOKENS_PER_TOOL = 2; // empirical adjustment for tool calls
+  public TOKENS_PER_CONVERSATION = 25; // fixed overhead for the conversation
+
+  /**
+   * Create a token limiter for messages.
+   * @param options Either a number (token limit) or a configuration object
+   */
+  constructor(options: number | TokenLimiterOptions) {
+    super({
+      name: 'TokenLimiter',
+    });
+
+    if (typeof options === 'number') {
+      // Simple number format - just the token limit with default encoding
+      this.maxTokens = options;
+      this.encoder = new Tiktoken(o200k_base);
+    } else {
+      // Object format with limit and optional encoding
+      this.maxTokens = options.limit;
+      this.encoder = new Tiktoken(options.encoding || o200k_base);
+    }
+  }
+
+  process(
+    messages: CoreMessage[],
+    { systemMessage, memorySystemMessage, newMessages }: MemoryProcessorOpts = {},
+  ): CoreMessage[] {
+    // Messages are already chronologically ordered - take most recent ones up to the token limit
+    let totalTokens = 0;
+
+    // Start with the conversation overhead
+    totalTokens += this.TOKENS_PER_CONVERSATION;
+
+    if (systemMessage) {
+      totalTokens += this.countTokens(systemMessage);
+      totalTokens += this.TOKENS_PER_MESSAGE; // Add message overhead for system message
+    }
+
+    if (memorySystemMessage) {
+      totalTokens += this.countTokens(memorySystemMessage);
+      totalTokens += this.TOKENS_PER_MESSAGE; // Add message overhead for memory system message
+    }
+
+    const allMessages = [...messages, ...(newMessages || [])];
+
+    const result: CoreMessage[] = [];
+
+    // Process messages in reverse (newest first)
+    for (let i = allMessages.length - 1; i >= 0; i--) {
+      const message = allMessages[i];
+
+      // Skip undefined messages (shouldn't happen, but TypeScript is concerned)
+      if (!message) continue;
+
+      const messageTokens = this.countTokens(message);
+
+      if (totalTokens + messageTokens <= this.maxTokens) {
+        // Insert at the beginning to maintain chronological order
+        result.unshift(message);
+        totalTokens += messageTokens;
+      } else {
+        this.logger.info(
+          `filtering ${allMessages.length - result.length}/${allMessages.length} messages, token limit of ${this.maxTokens} exceeded`,
+        );
+        // If we can't fit the message, we stop
+        break;
+      }
+    }
+
+    return result;
+  }
+
+  public countTokens(message: string | CoreMessage): number {
+    if (typeof message === `string`) {
+      return this.encoder.encode(message).length;
+    }
+
+    let tokenString = message.role;
+
+    if (typeof message.content === 'string') {
+      tokenString += message.content;
+    } else if (Array.isArray(message.content)) {
+      // Calculate tokens for each content part
+      for (const part of message.content) {
+        tokenString += part.type;
+        if (part.type === 'text') {
+          tokenString += part.text;
+        } else if (part.type === 'tool-call') {
+          tokenString += part.toolName as any;
+          if (part.args) {
+            tokenString += typeof part.args === 'string' ? part.args : JSON.stringify(part.args);
+          }
+        } else if (part.type === 'tool-result') {
+          // Token cost for result if present
+          if (part.result !== undefined) {
+            tokenString += typeof part.result === 'string' ? part.result : JSON.stringify(part.result);
+          }
+        } else {
+          tokenString += JSON.stringify(part);
+        }
+      }
+    }
+
+    // Ensure we account for message formatting tokens
+    // See: https://cookbook.openai.com/examples/how_to_count_tokens_with_tiktoken#6-counting-tokens-for-chat-completions-api-calls
+    const messageOverhead = this.TOKENS_PER_MESSAGE;
+
+    // Count tool calls for additional overhead
+    let toolOverhead = 0;
+    if (Array.isArray(message.content)) {
+      for (const part of message.content) {
+        if (part.type === 'tool-call' || part.type === 'tool-result') {
+          toolOverhead += this.TOKENS_PER_TOOL;
+        }
+      }
+    }
+
+    const totalMessageOverhead = messageOverhead + toolOverhead;
+
+    return this.encoder.encode(tokenString).length + totalMessageOverhead;
+  }
+}

package/src/processors/tool-call-filter.ts

@@ -0,0 +1,77 @@
+import type { CoreMessage } from '@mastra/core';
+import { MemoryProcessor } from '@mastra/core';
+
+/**
+ * Filters out tool calls and results from messages.
+ * By default (with no arguments), excludes all tool calls and their results.
+ * Can be configured to exclude only specific tools by name.
+ */
+export class ToolCallFilter extends MemoryProcessor {
+  private exclude: string[] | 'all';
+
+  /**
+   * Create a filter for tool calls and results.
+   * @param options Configuration options
+   * @param options.exclude List of specific tool names to exclude. If not provided, all tool calls are excluded.
+   */
+  constructor(options: { exclude?: string[] } = {}) {
+    super({ name: 'ToolCallFilter' });
+    // If no options or exclude is provided, exclude all tools
+    if (!options || !options.exclude) {
+      this.exclude = 'all'; // Exclude all tools
+    } else {
+      // Exclude specific tools
+      this.exclude = Array.isArray(options.exclude) ? options.exclude : [];
+    }
+  }
+
+  process(messages: CoreMessage[]): CoreMessage[] {
+    // Case 1: Exclude all tool calls and tool results
+    if (this.exclude === 'all') {
+      return messages.filter(message => {
+        if (Array.isArray(message.content)) {
+          return !message.content.some(part => part.type === 'tool-call' || part.type === 'tool-result');
+        }
+        return true;
+      });
+    }
+
+    // Case 2: Exclude specific tools by name
+    if (this.exclude.length > 0) {
+      // Single pass approach - track excluded tool call IDs while filtering
+      const excludedToolCallIds = new Set<string>();
+
+      return messages.filter(message => {
+        if (!Array.isArray(message.content)) return true;
+
+        // For assistant messages, check for excluded tool calls and track their IDs
+        if (message.role === 'assistant') {
+          let shouldExclude = false;
+
+          for (const part of message.content) {
+            if (part.type === 'tool-call' && this.exclude.includes(part.toolName)) {
+              excludedToolCallIds.add(part.toolCallId);
+              shouldExclude = true;
+            }
+          }
+
+          return !shouldExclude;
+        }
+
+        // For tool messages, filter out results for excluded tool calls
+        if (message.role === 'tool') {
+          const shouldExclude = message.content.some(
+            part => part.type === 'tool-result' && excludedToolCallIds.has(part.toolCallId),
+          );
+
+          return !shouldExclude;
+        }
+
+        return true;
+      });
+    }
+
+    // Case 3: Empty exclude array, return original messages
+    return messages;
+  }
+}

package/src/tools/working-memory.ts

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

package/src/utils/index.ts

@@ -0,0 +1,88 @@
+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;
+}

package/tsconfig.json

@@ -0,0 +1,5 @@
+{
+  "extends": "../../tsconfig.node.json",
+  "include": ["src/**/*"],
+  "exclude": ["node_modules", "**/*.test.ts"]
+}

package/vitest.config.ts

@@ -0,0 +1,8 @@
+import { defineConfig } from 'vitest/config';
+
+export default defineConfig({
+  test: {
+    environment: 'node',
+    include: ['src/**/*.test.ts'],
+  },
+});