@mastra/memory 0.0.1 → 0.0.2-alpha.1

package/src/redis/index.test.ts

@@ -1,245 +0,0 @@
-import { ThreadType, MessageType } from '@mastra/core';
-
-import { LocalRedisProvider, RedisMemory } from './';
-
-describe('RedisMemory', () => {
-  let memory: RedisMemory;
-  let provider: LocalRedisProvider;
-
-  beforeAll(async () => {
-    provider = new LocalRedisProvider();
-    memory = new RedisMemory(provider);
-  });
-
-  afterAll(async () => {
-    await memory.cleanup();
-    await provider.quit();
-  });
-
-  beforeEach(async () => {
-    await provider.flushall();
-  });
-
-  describe('Thread Operations', () => {
-    it('should create a thread with metadata', async () => {
-      const thread = await memory.createThread('Test Thread', {
-        testKey: 'testValue',
-      });
-
-      expect(thread).toEqual({
-        id: expect.any(String),
-        title: 'Test Thread',
-        createdAt: expect.any(Date),
-        updatedAt: expect.any(Date),
-        metadata: { testKey: 'testValue' },
-      });
-    });
-
-    it('should retrieve a created thread', async () => {
-      const created = await memory.createThread('Test Thread');
-      const retrieved = await memory.getThreadById(created.id);
-
-      expect(retrieved).toEqual(created);
-    });
-
-    it('should return null for non-existent thread', async () => {
-      const thread = await memory.getThreadById('nonexistent');
-      expect(thread).toBeNull();
-    });
-
-    it('should update thread title', async () => {
-      const thread = await memory.createThread('Initial Title');
-      thread.title = 'Updated Title';
-
-      await memory.saveThread(thread);
-      const retrieved = await memory.getThreadById(thread.id);
-
-      expect(retrieved?.title).toBe('Updated Title');
-    });
-  });
-
-  describe('Message Operations', () => {
-    let testThread: ThreadType;
-
-    beforeEach(async () => {
-      testThread = await memory.createThread('Test Thread');
-    });
-
-    it('should add a single message', async () => {
-      const message = await memory.addMessage(testThread.id, 'Hello World', 'user');
-
-      const messages = await memory.getMessages(testThread.id);
-      expect(messages).toHaveLength(1);
-      expect(messages[0]).toEqual(message);
-    });
-
-    it('should maintain message order', async () => {
-      await memory.addMessage(testThread.id, 'First', 'user');
-      await memory.addMessage(testThread.id, 'Second', 'assistant');
-      await memory.addMessage(testThread.id, 'Third', 'user');
-
-      const messages = await memory.getMessages(testThread.id);
-      expect(messages).toHaveLength(3);
-      expect(messages.map(m => m.content)).toEqual(['First', 'Second', 'Third']);
-    });
-
-    it('should handle bulk message saves', async () => {
-      const messages: MessageType[] = [
-        {
-          id: 'msg1',
-          threadId: testThread.id,
-          content: 'Message 1',
-          role: 'user',
-          createdAt: new Date(),
-        },
-        {
-          id: 'msg2',
-          threadId: testThread.id,
-          content: 'Message 2',
-          role: 'assistant',
-          createdAt: new Date(),
-        },
-      ];
-
-      await memory.saveMessages(messages);
-      const retrieved = await memory.getMessages(testThread.id);
-      expect(retrieved).toHaveLength(2);
-    });
-
-    it('should prevent duplicate messages', async () => {
-      const message = await memory.addMessage(testThread.id, 'Test', 'user');
-      await memory.saveMessages([message]); // Try to save the same message again
-
-      const messages = await memory.getMessages(testThread.id);
-      expect(messages).toHaveLength(1);
-    });
-  });
-
-  describe('Thread Management', () => {
-    it('should list all thread IDs', async () => {
-      const thread1 = await memory.createThread('Thread 1');
-      const thread2 = await memory.createThread('Thread 2');
-
-      const ids = await memory.getAllThreadIds();
-      expect(ids).toContain(thread1.id);
-      expect(ids).toContain(thread2.id);
-    });
-
-    it('should delete thread and its messages', async () => {
-      const thread = await memory.createThread('Delete Test');
-      await memory.addMessage(thread.id, 'Test Message', 'user');
-
-      await memory.deleteThread(thread.id);
-
-      const deletedThread = await memory.getThreadById(thread.id);
-      const messages = await memory.getMessages(thread.id);
-
-      expect(deletedThread).toBeNull();
-      expect(messages).toHaveLength(0);
-    });
-
-    it('should retrieve multiple threads', async () => {
-      const thread1 = await memory.createThread('Thread 1');
-      const thread2 = await memory.createThread('Thread 2');
-
-      const threads = await memory.getThreads([thread1.id, thread2.id]);
-      expect(threads).toHaveLength(2);
-      expect(threads.map(t => t.title)).toEqual(['Thread 1', 'Thread 2']);
-    });
-  });
-
-  describe('Error Handling', () => {
-    it('should handle saving messages to non-existent thread', async () => {
-      const message: MessageType = {
-        id: 'test',
-        threadId: 'nonexistent',
-        content: 'Test',
-        role: 'user',
-        createdAt: new Date(),
-      };
-
-      await memory.saveMessages([message]);
-      const messages = await memory.getMessages('nonexistent');
-      expect(messages).toHaveLength(1);
-    });
-
-    it('should handle concurrent message saves', async () => {
-      const thread = await memory.createThread('Concurrent Test');
-
-      // Simulate concurrent saves
-      await Promise.all([
-        memory.addMessage(thread.id, 'Message 1', 'user'),
-        memory.addMessage(thread.id, 'Message 2', 'user'),
-        memory.addMessage(thread.id, 'Message 3', 'user'),
-      ]);
-
-      const messages = await memory.getMessages(thread.id);
-      expect(messages).toHaveLength(3);
-    });
-  });
-
-  describe('Performance', () => {
-    it('should handle large number of messages', async () => {
-      const thread = await memory.createThread('Bulk Test');
-      const messageCount = 100;
-
-      // Create messages array first
-      const messages = Array.from({ length: messageCount }, (_, i) => ({
-        id: memory['generateId'](),
-        content: `Message ${i}`,
-        role: i % 2 === 0 ? ('user' as const) : ('assistant' as const),
-        createdAt: new Date(),
-        threadId: thread.id,
-      }));
-
-      // Save all messages in one operation
-      await memory.saveMessages(messages);
-
-      const retrievedMessages = await memory.getMessages(thread.id);
-      expect(retrievedMessages).toHaveLength(messageCount);
-
-      // Verify message content and ordering
-      const sortedMessages = retrievedMessages.sort((a, b) => {
-        const aNum = parseInt(a.content.split(' ')[1] || '0');
-        const bNum = parseInt(b.content.split(' ')[1] || '0');
-        return aNum - bNum;
-      });
-
-      sortedMessages.forEach((msg, i) => {
-        expect(msg.content).toBe(`Message ${i}`);
-      });
-    });
-
-    // Add a test for truly concurrent operations
-    it('should handle concurrent batch saves', async () => {
-      const thread = await memory.createThread('Concurrent Batch Test');
-      const batchSize = 20;
-      const numberOfBatches = 5;
-
-      // Create batches of messages
-      const batches = Array.from({ length: numberOfBatches }, (_, batchIndex) =>
-        Array.from({ length: batchSize }, (_, i) => ({
-          id: memory['generateId'](),
-          content: `Batch ${batchIndex} Message ${i}`,
-          role: i % 2 === 0 ? ('user' as const) : ('assistant' as const),
-          createdAt: new Date(),
-          threadId: thread.id,
-        })),
-      );
-
-      // Save batches concurrently
-      await Promise.all(batches.map(batch => memory.saveMessages(batch)));
-
-      const messages = await memory.getMessages(thread.id);
-      expect(messages).toHaveLength(batchSize * numberOfBatches);
-
-      // Verify all messages are present
-      const messageSet = new Set(messages.map(m => m.content));
-      for (let batchIndex = 0; batchIndex < numberOfBatches; batchIndex++) {
-        for (let i = 0; i < batchSize; i++) {
-          expect(messageSet.has(`Batch ${batchIndex} Message ${i}`)).toBeTruthy();
-        }
-      }
-    });
-  });
-});

package/src/redis/index.ts

@@ -1,189 +0,0 @@
-import { MastraMemory, ThreadType, MessageType } from '@mastra/core';
-
-import { RedisClient } from './types';
-
-export * from './types';
-export * from './providers';
-
-export class RedisMemory extends MastraMemory {
-  private redis: RedisClient;
-  private threadPrefix = 'thread:';
-  private messagePrefix = 'messages:';
-  private lockTimeouts: Map<string, NodeJS.Timeout> = new Map();
-
-  constructor(redis: RedisClient) {
-    super();
-    this.redis = redis;
-  }
-
-  async getThreadById(threadId: string): Promise<ThreadType | null> {
-    const thread = await this.redis.get(`${this.threadPrefix}${threadId}`);
-    if (thread && typeof thread.createdAt === 'string') {
-      thread.createdAt = new Date(thread.createdAt);
-      thread.updatedAt = new Date(thread.updatedAt);
-    }
-    return thread;
-  }
-
-  async saveThread(thread: ThreadType): Promise<ThreadType> {
-    thread.updatedAt = new Date();
-    await this.redis.set(`${this.threadPrefix}${thread.id}`, thread);
-    await this.redis.sadd('threads', thread.id);
-    return thread;
-  }
-
-  private async withLock<T>(key: string, operation: () => Promise<T>): Promise<T> {
-    const lockKey = `lock:${key}`;
-    const lockTimeout = 5000;
-    let locked = false;
-    let timeoutId: NodeJS.Timeout | undefined;
-
-    try {
-      // Try to acquire lock with retries
-      for (let i = 0; i < 3; i++) {
-        locked = await this.redis.sadd(lockKey, '1');
-        if (locked) break;
-        await new Promise(resolve => setTimeout(resolve, Math.random() * 100));
-      }
-
-      if (!locked) {
-        throw new Error('Could not acquire lock');
-      }
-
-      // Set lock timeout
-      timeoutId = setTimeout(async () => {
-        try {
-          await this.redis.del(lockKey);
-        } catch {
-          // Ignore errors during cleanup
-        }
-        this.lockTimeouts.delete(lockKey);
-      }, lockTimeout);
-
-      this.lockTimeouts.set(lockKey, timeoutId);
-
-      // Execute operation
-      return await operation();
-    } finally {
-      if (timeoutId !== undefined) {
-        clearTimeout(timeoutId);
-        this.lockTimeouts.delete(lockKey);
-      }
-      if (locked) {
-        try {
-          await this.redis.del(lockKey);
-        } catch {
-          // Ignore errors during cleanup
-        }
-      }
-    }
-  }
-
-  // Add cleanup method
-  async cleanup(): Promise<void> {
-    // Clear all timeouts
-    for (const timeout of this.lockTimeouts.values()) {
-      clearTimeout(timeout);
-    }
-    this.lockTimeouts.clear();
-  }
-
-  async saveMessages(messages: MessageType[]): Promise<MessageType[]> {
-    if (!messages.length) return [];
-
-    const messagesByThread = messages.reduce(
-      (acc, message) => {
-        const key = `${this.messagePrefix}${message.threadId}`;
-        if (!acc[key]) {
-          acc[key] = [];
-        }
-        acc[key].push({
-          ...message,
-          createdAt: new Date(message.createdAt),
-        });
-        return acc;
-      },
-      {} as Record<string, MessageType[]>,
-    );
-
-    for (const [key, threadMessages] of Object.entries(messagesByThread)) {
-      await this.withLock(key, async () => {
-        const existingMessages = (await this.redis.get(key)) || [];
-
-        const messageMap = new Map<string, MessageType>();
-
-        // Process existing messages
-        existingMessages.forEach((msg: MessageType) => {
-          messageMap.set(msg.id, {
-            ...msg,
-            createdAt: new Date(msg.createdAt),
-          });
-        });
-
-        // Add new messages
-        threadMessages.forEach(msg => {
-          messageMap.set(msg.id, msg);
-        });
-
-        const updatedMessages = Array.from(messageMap.values());
-        updatedMessages.sort((a, b) => {
-          const timeCompare = a.createdAt.getTime() - b.createdAt.getTime();
-          return timeCompare === 0 ? a.id.localeCompare(b.id) : timeCompare;
-        });
-
-        await this.redis.set(key, updatedMessages);
-        if (threadMessages?.[0]?.threadId) {
-          const thread = await this.getThreadById(threadMessages?.[0]?.threadId);
-          if (thread) {
-            thread.updatedAt = new Date();
-            await this.redis.set(`${this.threadPrefix}${thread.id}`, thread);
-          }
-        }
-      });
-    }
-
-    return messages;
-  }
-
-  async addMessage(threadId: string, content: string, role: 'user' | 'assistant'): Promise<MessageType> {
-    const message: MessageType = {
-      id: this.generateId(),
-      content,
-      role,
-      createdAt: new Date(),
-      threadId,
-    };
-
-    await this.saveMessages([message]);
-    return message;
-  }
-
-  async getMessages(threadId: string): Promise<MessageType[]> {
-    const messages = (await this.redis.get(`${this.messagePrefix}${threadId}`)) || [];
-    return messages.map((msg: MessageType) => ({
-      ...msg,
-      createdAt: new Date(msg.createdAt),
-    }));
-  }
-
-  async getAllThreadIds(): Promise<string[]> {
-    return this.redis.smembers('threads');
-  }
-
-  async deleteThread(threadId: string): Promise<void> {
-    const pipeline = this.redis.pipeline();
-    pipeline.del(`${this.threadPrefix}${threadId}`);
-    pipeline.del(`${this.messagePrefix}${threadId}`);
-    pipeline.srem('threads', threadId);
-    await pipeline.exec();
-  }
-
-  async getThreads(threadIds: string[]): Promise<ThreadType[]> {
-    const threads = await Promise.all(threadIds.map(id => this.getThreadById(id)));
-    return threads.filter((t): t is ThreadType => t !== null);
-  }
-
-  protected generateId(): string {
-    return crypto.randomUUID();
-  }
-}

package/src/redis/providers.ts

@@ -1,191 +0,0 @@
-import { Redis as UpstashRedis } from '@upstash/redis';
-import { createClient } from 'redis';
-
-import { RedisClient, RedisPipeline } from './types';
-
-// Helper functions for date handling
-function serializeData(data: any): any {
-  if (data === null || data === undefined) return data;
-
-  if (data instanceof Date) {
-    return { __type: 'Date', value: data.toISOString() };
-  }
-
-  if (Array.isArray(data)) {
-    return data.map(item => serializeData(item));
-  }
-
-  if (typeof data === 'object') {
-    return Object.keys(data).reduce((acc, key) => {
-      acc[key] = serializeData(data[key]);
-      return acc;
-    }, {} as any);
-  }
-
-  return data;
-}
-
-function deserializeData(data: any): any {
-  if (data === null || data === undefined) return data;
-
-  if (typeof data === 'object' && data.__type === 'Date') {
-    return new Date(data.value);
-  }
-
-  if (Array.isArray(data)) {
-    return data.map(item => deserializeData(item));
-  }
-
-  if (typeof data === 'object') {
-    return Object.keys(data).reduce((acc, key) => {
-      acc[key] = deserializeData(data[key]);
-      return acc;
-    }, {} as any);
-  }
-
-  return data;
-}
-
-// Provider for Upstash/Vercel KV
-export class UpstashProvider implements RedisClient {
-  private client: UpstashRedis;
-
-  constructor(url: string, token: string) {
-    this.client = new UpstashRedis({
-      url,
-      token,
-    });
-  }
-  async get(key: string) {
-    const data = await this.client.get(key);
-    return deserializeData(data);
-  }
-
-  async set(key: string, value: any) {
-    return this.client.set(key, serializeData(value));
-  }
-  async del(key: string) {
-    return this.client.del(key);
-  }
-  async sadd(key: string, value: string) {
-    return this.client.sadd(key, value);
-  }
-  async srem(key: string, value: string) {
-    return this.client.srem(key, value);
-  }
-  async smembers(key: string) {
-    return this.client.smembers(key);
-  }
-  async flushall() {
-    return this.client.flushall();
-  }
-  pipeline() {
-    const multi = this.client.multi();
-    const pipeline: RedisPipeline = {
-      get: (key: string) => {
-        multi.get(key);
-        return pipeline;
-      },
-      set: (key: string, value: any) => {
-        multi.set(key, JSON.stringify(serializeData(value)));
-        return pipeline;
-      },
-      del: (key: string) => {
-        multi.del(key);
-        return pipeline;
-      },
-      srem: (key: string, value: string) => {
-        multi.srem(key, value);
-        return pipeline;
-      },
-      exec: async () => {
-        const results = await multi.exec();
-        return results.map(result => {
-          try {
-            // For get operations that return string data
-            if (typeof result === 'string') {
-              return deserializeData(JSON.parse(result));
-            }
-            // For array results (like from lists/sets)
-            if (Array.isArray(result)) {
-              return result.map(item => {
-                try {
-                  return typeof item === 'string' ? deserializeData(JSON.parse(item)) : item;
-                } catch {
-                  return item;
-                }
-              });
-            }
-            return result;
-          } catch {
-            return result;
-          }
-        });
-      },
-    };
-    return pipeline;
-  }
-}
-
-export class LocalRedisProvider implements RedisClient {
-  private client: any;
-
-  constructor() {
-    this.client = createClient({ url: 'redis://localhost:6379' });
-    this.client.connect();
-  }
-
-  async get(key: string) {
-    const data = await this.client.get(key);
-    return data ? deserializeData(JSON.parse(data)) : null;
-  }
-
-  async set(key: string, value: any) {
-    return this.client.set(key, JSON.stringify(serializeData(value)));
-  }
-  async del(key: string) {
-    return this.client.del(key);
-  }
-  async sadd(key: string, value: string) {
-    return this.client.sAdd(key, value);
-  }
-  async srem(key: string, value: string) {
-    return this.client.sRem(key, value);
-  }
-  async smembers(key: string) {
-    return this.client.sMembers(key);
-  }
-  async flushall() {
-    return this.client.flushAll();
-  }
-
-  pipeline() {
-    const multi = this.client.multi();
-    const pipeline: RedisPipeline = {
-      get: (key: string) => {
-        multi.get(key);
-        return pipeline;
-      },
-      set: (key: string, value: any) => {
-        multi.set(key, JSON.stringify(value));
-        return pipeline;
-      },
-      del: (key: string) => {
-        multi.del(key);
-        return pipeline;
-      },
-      srem: (key: string, value: string) => {
-        multi.sRem(key, value);
-        return pipeline;
-      },
-      exec: async () => {
-        return multi.exec();
-      },
-    };
-    return pipeline;
-  }
-
-  async quit() {
-    await this.client.quit();
-  }
-}

package/src/redis/types.ts

@@ -1,18 +0,0 @@
-export interface RedisClient {
-  get(key: string): Promise<any>;
-  set(key: string, value: any): Promise<any>;
-  del(key: string): Promise<any>;
-  sadd(key: string, ...values: string[]): Promise<any>;
-  srem(key: string, ...values: string[]): Promise<any>;
-  smembers(key: string): Promise<string[]>;
-  flushall(): Promise<any>;
-  pipeline(): RedisPipeline;
-}
-
-export interface RedisPipeline {
-  get(key: string): RedisPipeline;
-  set(key: string, value: any): RedisPipeline;
-  del(key: string): RedisPipeline;
-  srem(key: string, value: string): RedisPipeline;
-  exec(): Promise<any>;
-}