@mastra/memory 0.0.1

package/src/postgres/index.ts

@@ -0,0 +1,256 @@
+import { MastraMemory, MessageType, ThreadType } from '@mastra/core';
+import { randomUUID } from 'crypto';
+import pg from 'pg';
+
+const { Pool } = pg;
+
+export class PgMemory extends MastraMemory {
+  private pool: pg.Pool;
+
+  constructor(connectionString: string) {
+    super();
+    this.pool = new Pool({ connectionString });
+  }
+
+  async drop() {
+    const client = await this.pool.connect();
+    await client.query('DELETE FROM mastra_messages');
+    await client.query('DELETE FROM mastra_threads');
+    client.release();
+    await this.pool.end();
+  }
+
+  async ensureTablesExist(): Promise<void> {
+    const client = await this.pool.connect();
+    try {
+      // Check if the threads table exists
+      const threadsResult = await client.query<{ exists: boolean }>(`
+                SELECT EXISTS (
+                    SELECT 1
+                    FROM information_schema.tables
+                    WHERE table_name = 'mastra_threads'
+                );
+            `);
+
+      if (!threadsResult?.rows?.[0]?.exists) {
+        await client.query(`
+                    CREATE TABLE IF NOT EXISTS mastra_threads (
+                        id UUID PRIMARY KEY,
+                        title TEXT,
+                        created_at TIMESTAMP WITH TIME ZONE NOT NULL,
+                        updated_at TIMESTAMP WITH TIME ZONE NOT NULL,
+                        metadata JSONB
+                    );
+                `);
+      }
+
+      // Check if the messages table exists
+      const messagesResult = await client.query<{ exists: boolean }>(`
+                SELECT EXISTS (
+                    SELECT 1
+                    FROM information_schema.tables
+                    WHERE table_name = 'mastra_messages'
+                );
+            `);
+
+      if (!messagesResult?.rows?.[0]?.exists) {
+        await client.query(`
+                    CREATE TABLE IF NOT EXISTS mastra_messages (
+                        id UUID PRIMARY KEY,
+                        content TEXT NOT NULL,
+                        role VARCHAR(20) NOT NULL,
+                        created_at TIMESTAMP WITH TIME ZONE NOT NULL,
+                        thread_id UUID NOT NULL,
+                        FOREIGN KEY (thread_id) REFERENCES mastra_threads(id)
+                    );
+                `);
+      }
+    } finally {
+      client.release();
+    }
+  }
+
+  async updateThread(id: string, title: string, metadata: Record<string, unknown>): Promise<ThreadType> {
+    const client = await this.pool.connect();
+    try {
+      const result = await client.query<ThreadType>(
+        `
+                UPDATE mastra_threads
+                SET title = $1, metadata = $2, updated_at = NOW()
+                WHERE id = $3
+                RETURNING *
+                `,
+        [title, JSON.stringify(metadata), id],
+      );
+      return result?.rows?.[0]!;
+    } finally {
+      client.release();
+    }
+  }
+
+  async deleteThread(id: string): Promise<void> {
+    const client = await this.pool.connect();
+    try {
+      await client.query(
+        `
+                DELETE FROM mastra_messages
+                WHERE thread_id = $1
+                `,
+        [id],
+      );
+
+      await client.query(
+        `
+                DELETE FROM mastra_threads
+                WHERE id = $1
+                `,
+        [id],
+      );
+    } finally {
+      client.release();
+    }
+  }
+
+  async deleteMessage(id: string): Promise<void> {
+    const client = await this.pool.connect();
+    try {
+      await client.query(
+        `
+                DELETE FROM mastra_messages
+                WHERE id = $1
+                `,
+        [id],
+      );
+    } finally {
+      client.release();
+    }
+  }
+
+  async getThreadById(threadId: string): Promise<ThreadType | null> {
+    await this.ensureTablesExist();
+
+    const client = await this.pool.connect();
+    try {
+      const result = await client.query<ThreadType>(
+        `
+                SELECT id, title, created_at AS createdAt, updated_at AS updatedAt, metadata
+                FROM mastra_threads
+                WHERE id = $1
+            `,
+        [threadId],
+      );
+      return result.rows[0] || null;
+    } finally {
+      client.release();
+    }
+  }
+
+  async saveThread(thread: ThreadType): Promise<ThreadType> {
+    await this.ensureTablesExist();
+
+    const client = await this.pool.connect();
+    try {
+      const { id, title, createdAt, updatedAt, metadata } = thread;
+      const result = await client.query<ThreadType>(
+        `
+                INSERT INTO mastra_threads (id, title, created_at, updated_at, metadata)
+                VALUES ($1, $2, $3, $4, $5)
+                ON CONFLICT (id) DO UPDATE SET title = $2, updated_at = $4, metadata = $5
+                RETURNING id, title, created_at AS createdAt, updated_at AS updatedAt, metadata
+            `,
+        [id, title, createdAt, updatedAt, JSON.stringify(metadata)],
+      );
+      return result?.rows?.[0]!;
+    } finally {
+      client.release();
+    }
+  }
+
+  async saveMessages(messages: MessageType[]): Promise<MessageType[]> {
+    await this.ensureTablesExist();
+
+    const client = await this.pool.connect();
+    try {
+      await client.query('BEGIN');
+      for (const message of messages) {
+        const { id, content, role, createdAt, threadId } = message;
+        await client.query(
+          `
+                    INSERT INTO mastra_messages (id, content, role, created_at, thread_id)
+                    VALUES ($1, $2, $3, $4, $5)
+                `,
+          [id, content, role, createdAt.toISOString(), threadId],
+        );
+      }
+      await client.query('COMMIT');
+      return messages;
+    } catch (error) {
+      await client.query('ROLLBACK');
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+
+  async getMessages(threadId: string): Promise<MessageType[]> {
+    await this.ensureTablesExist();
+
+    const client = await this.pool.connect();
+    try {
+      const result = await client.query<MessageType>(
+        `
+                SELECT 
+                    id, 
+                    content, 
+                    role, 
+                    created_at AS createdAt, 
+                    thread_id AS threadId
+                FROM mastra_messages
+                WHERE thread_id = $1
+                ORDER BY created_at ASC
+            `,
+        [threadId],
+      );
+      return result.rows;
+    } finally {
+      client.release();
+    }
+  }
+
+  async createThread(title?: string, metadata?: Record<string, unknown>): Promise<ThreadType> {
+    await this.ensureTablesExist();
+
+    const id = randomUUID();
+    const now = new Date();
+    const thread: ThreadType = {
+      id,
+      title,
+      createdAt: now,
+      updatedAt: now,
+      metadata,
+    };
+    return this.saveThread(thread);
+  }
+
+  async addMessage(threadId: string, content: string, role: 'user' | 'assistant'): Promise<MessageType> {
+    await this.ensureTablesExist();
+
+    const thread = await this.getThreadById(threadId);
+    if (!thread) {
+      throw new Error(`Thread with ID ${threadId} does not exist.`);
+    }
+
+    const id = randomUUID();
+    const message: MessageType = {
+      id,
+      content,
+      role,
+      createdAt: new Date(),
+      threadId,
+    };
+
+    const [savedMessage] = await this.saveMessages([message]);
+
+    return savedMessage!;
+  }
+}

package/src/redis/index.test.ts

@@ -0,0 +1,245 @@
+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

@@ -0,0 +1,189 @@
+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();
+  }
+}