@mastra/memory 0.0.1

package/src/cloudflare/index.test.ts

@@ -0,0 +1,230 @@
+import { ThreadType } from '@mastra/core';
+
+import { CloudflareKVMemory } from './';
+import { MockKV } from './kv';
+
+describe('CloudflareKVMemory', () => {
+  let memory: CloudflareKVMemory;
+  let mockKV: MockKV;
+
+  beforeEach(() => {
+    mockKV = new MockKV();
+    memory = new CloudflareKVMemory(mockKV);
+  });
+
+  afterEach(async () => {
+    await memory.cleanup();
+  });
+
+  describe('Thread Operations', () => {
+    it('should create and retrieve a thread', 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' },
+      });
+
+      const retrieved = await memory.getThreadById(thread.id);
+      expect(retrieved).toEqual(thread);
+    });
+
+    it('should return null for non-existent thread', async () => {
+      const thread = await memory.getThreadById('nonexistent');
+      expect(thread).toBeNull();
+    });
+
+    it('should update thread timestamps', async () => {
+      const thread = await memory.createThread('Initial Title');
+      const originalUpdatedAt = thread.updatedAt;
+
+      // Wait a bit to ensure timestamp difference
+      await new Promise(resolve => setTimeout(resolve, 100));
+
+      thread.title = 'Updated Title';
+      const updated = await memory.saveThread(thread);
+
+      expect(updated.updatedAt.getTime()).toBeGreaterThan(originalUpdatedAt.getTime());
+    });
+  });
+
+  describe('Message Operations', () => {
+    let testThread: ThreadType;
+
+    beforeEach(async () => {
+      testThread = await memory.createThread('Message Test Thread');
+    });
+
+    it('should add and retrieve messages', async () => {
+      // Add multiple messages
+      await memory.addMessage(testThread.id, 'Hello', 'user');
+      await memory.addMessage(testThread.id, 'Hi there', 'assistant');
+
+      // Retrieve all messages
+      const messages = await memory.getMessages(testThread.id);
+
+      expect(messages).toHaveLength(2);
+
+      // Verify both messages are present without assuming order
+      expect(messages).toEqual(
+        expect.arrayContaining([
+          expect.objectContaining({
+            content: 'Hello',
+            role: 'user',
+            threadId: testThread.id,
+          }),
+          expect.objectContaining({
+            content: 'Hi there',
+            role: 'assistant',
+            threadId: testThread.id,
+          }),
+        ]),
+      );
+
+      // Also verify that each message has all required fields
+      messages.forEach(msg => {
+        expect(msg).toEqual(
+          expect.objectContaining({
+            id: expect.any(String),
+            content: expect.any(String),
+            role: expect.stringMatching(/^(user|assistant)$/),
+            createdAt: expect.any(Date),
+            threadId: testThread.id,
+          }),
+        );
+      });
+    });
+
+    it('should handle message ordering', async () => {
+      const messages = [
+        {
+          id: memory['generateId'](),
+          content: 'Message 1',
+          role: 'user' as const,
+          createdAt: new Date('2024-01-01'),
+          threadId: testThread.id,
+        },
+        {
+          id: memory['generateId'](),
+          content: 'Message 2',
+          role: 'assistant' as const,
+          createdAt: new Date('2024-01-02'),
+          threadId: testThread.id,
+        },
+      ];
+
+      await memory.saveMessages(messages);
+      const retrieved = await memory.getMessages(testThread.id);
+
+      expect(retrieved).toHaveLength(2);
+      expect(retrieved?.[0]?.content).toBe('Message 1');
+      expect(retrieved?.[1]?.content).toBe('Message 2');
+    });
+
+    it('should handle duplicate message saves', async () => {
+      const message = await memory.addMessage(testThread.id, 'Duplicate', 'user');
+
+      // Try to save the same message again
+      await memory.saveMessages([message]);
+
+      const messages = await memory.getMessages(testThread.id);
+      expect(messages).toHaveLength(1);
+    });
+  });
+
+  describe('Batch Operations', () => {
+    it('should handle multiple threads', async () => {
+      const thread1 = await memory.createThread('Thread 1');
+      const thread2 = await memory.createThread('Thread 2');
+
+      const allThreadIds = await memory.getAllThreadIds();
+      expect(allThreadIds).toContain(thread1.id);
+      expect(allThreadIds).toContain(thread2.id);
+
+      const threads = await memory.getThreads([thread1.id, thread2.id]);
+      expect(threads).toHaveLength(2);
+    });
+
+    it('should delete thread and associated 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);
+    });
+  });
+
+  describe.only('Error Handling', () => {
+    it('should handle concurrent message saves', async () => {
+      const thread = await memory.createThread('Concurrent Test');
+      console.log('Created thread:', thread.id);
+
+      const messagesToSave = [
+        {
+          id: memory['generateId'](),
+          content: 'Message 1',
+          role: 'user' as const,
+          createdAt: new Date(),
+          threadId: thread.id,
+        },
+        {
+          id: memory['generateId'](),
+          content: 'Message 2',
+          role: 'user' as const,
+          createdAt: new Date(),
+          threadId: thread.id,
+        },
+        {
+          id: memory['generateId'](),
+          content: 'Message 3',
+          role: 'user' as const,
+          createdAt: new Date(),
+          threadId: thread.id,
+        },
+      ];
+
+      console.log('Messages to save:', messagesToSave);
+
+      // Save messages concurrently but in batches
+      if (messagesToSave?.length && messagesToSave?.[0] && messagesToSave?.[2]) {
+        await Promise.all([memory.saveMessages(messagesToSave).then(() => console.log('Batch 2 saved'))]);
+      }
+
+      const messages = await memory.getMessages(thread.id);
+      console.log('Retrieved messages:', messages);
+
+      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 many messages
+      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,
+      }));
+
+      await memory.saveMessages(messages);
+      const retrieved = await memory.getMessages(thread.id);
+
+      expect(retrieved).toHaveLength(messageCount);
+    });
+  });
+});

package/src/cloudflare/index.ts

@@ -0,0 +1,169 @@
+import { MastraMemory, MessageType, ThreadType } from '@mastra/core';
+
+import { CloudflareKVProvider, KVNamespace } from './kv';
+
+export class CloudflareKVMemory extends MastraMemory {
+  private kv: CloudflareKVProvider;
+  private threadPrefix = 'thread:';
+  private messagePrefix = 'messages:';
+
+  constructor(namespace: KVNamespace) {
+    super();
+    this.kv = new CloudflareKVProvider(namespace);
+  }
+
+  async getThreadById(threadId: string): Promise<ThreadType | null> {
+    const thread = await this.kv.get<ThreadType>(`${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.kv.set(`${this.threadPrefix}${thread.id}`, thread);
+    await this.kv.sadd('threads', thread.id);
+    return thread;
+  }
+
+  private async retryOperation<T>(operation: () => Promise<T>, maxRetries = 5): Promise<T> {
+    let lastError: Error | undefined;
+
+    for (let attempt = 0; attempt < maxRetries; attempt++) {
+      try {
+        return await operation();
+      } catch (error) {
+        lastError = error as Error;
+        await new Promise(resolve => setTimeout(resolve, Math.random() * 100 * Math.pow(2, attempt)));
+      }
+    }
+
+    throw lastError || new Error(`Operation failed after ${maxRetries} attempts`);
+  }
+
+  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.retryOperation(async () => {
+      const key = `${this.messagePrefix}${threadId}`;
+      const existingMessages = (await this.kv.get<MessageType[]>(key)) || [];
+
+      const messageMap = new Map<string, MessageType>();
+
+      existingMessages.forEach(msg => {
+        messageMap.set(msg.id, {
+          ...msg,
+          createdAt: new Date(msg.createdAt),
+        });
+      });
+
+      messageMap.set(message.id, message);
+
+      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.kv.set(key, updatedMessages);
+    });
+
+    return message;
+  }
+
+  async saveMessages(messages: MessageType[]): Promise<MessageType[]> {
+    if (!messages.length) return [];
+
+    const messagesByThread = new Map<string, MessageType[]>();
+
+    for (const message of messages) {
+      const key = `${this.messagePrefix}${message.threadId}`;
+      if (!messagesByThread.has(key)) {
+        messagesByThread.set(key, []);
+      }
+      messagesByThread.get(key)!.push({
+        ...message,
+        createdAt: new Date(message.createdAt),
+      });
+    }
+
+    await Promise.all(
+      Array.from(messagesByThread.entries()).map(([key, threadMessages]) =>
+        this.retryOperation(async () => {
+          let saved = false;
+          while (!saved) {
+            const { data: existingMessages, version } = await this.kv.getWithVersion<MessageType[]>(key);
+            console.log('Read version:', version, 'Messages:', existingMessages?.length || 0);
+
+            const messageMap = new Map<string, MessageType>();
+
+            (existingMessages || []).forEach(msg => {
+              messageMap.set(msg.id, {
+                ...msg,
+                createdAt: new Date(msg.createdAt),
+              });
+            });
+
+            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;
+            });
+
+            saved = await this.kv.setWithVersion(key, updatedMessages, version);
+            console.log('Save attempt with version:', version, 'Success:', saved);
+
+            if (!saved) {
+              // If save failed, someone else updated the messages, retry with new version
+              await new Promise(resolve => setTimeout(resolve, Math.random() * 50));
+            }
+          }
+        }),
+      ),
+    );
+
+    return messages;
+  }
+
+  async getMessages(threadId: string): Promise<MessageType[]> {
+    const messages = (await this.kv.get<MessageType[]>(`${this.messagePrefix}${threadId}`)) || [];
+    return messages.map(msg => ({
+      ...msg,
+      createdAt: new Date(msg.createdAt),
+    }));
+  }
+
+  async getAllThreadIds(): Promise<string[]> {
+    return this.kv.smembers('threads');
+  }
+
+  async deleteThread(threadId: string): Promise<void> {
+    await Promise.all([
+      this.kv.del(`${this.threadPrefix}${threadId}`),
+      this.kv.del(`${this.messagePrefix}${threadId}`),
+      this.kv.srem('threads', threadId),
+    ]);
+  }
+
+  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);
+  }
+
+  async cleanup(): Promise<void> {
+    // Flush all data
+    await this.kv.flushall();
+  }
+}

package/src/cloudflare/kv.ts

@@ -0,0 +1,139 @@
+// Actual Cloudflare KV interface
+export interface KVNamespace {
+  get(key: string): Promise<string | null>;
+  get(key: string, type: 'text'): Promise<string | null>;
+  get<T>(key: string, type: 'json'): Promise<T | null>;
+  put(key: string, value: string | ReadableStream | ArrayBuffer | FormData): Promise<void>;
+  delete(key: string): Promise<void>;
+  // List is actually on a different interface in real CF workers
+  list(options?: { prefix?: string; limit?: number; cursor?: string }): Promise<{
+    keys: { name: string }[];
+    list_complete: boolean;
+    cursor?: string;
+  }>;
+}
+
+// Mock implementation for testing
+export class MockKV implements KVNamespace {
+  private store = new Map<string, string>();
+
+  async get(key: string): Promise<string | null>;
+  async get(key: string, type: 'text'): Promise<string | null>;
+  async get<T>(key: string, type: 'json'): Promise<T | null>;
+  async get<T>(key: string, type?: 'text' | 'json'): Promise<string | T | null> {
+    const value = this.store.get(key);
+    if (value === undefined) return null;
+
+    if (type === 'json') {
+      return JSON.parse(value) as T;
+    }
+
+    return value;
+  }
+
+  async put(key: string, value: string | ReadableStream | ArrayBuffer | FormData): Promise<void> {
+    if (typeof value !== 'string') {
+      throw new Error('MockKV only supports string values for testing');
+    }
+    this.store.set(key, value);
+  }
+
+  async delete(key: string): Promise<void> {
+    this.store.delete(key);
+  }
+
+  async list(options?: { prefix?: string; limit?: number; cursor?: string }): Promise<{
+    keys: { name: string }[];
+    list_complete: boolean;
+    cursor?: string;
+  }> {
+    const allKeys = Array.from(this.store.keys())
+      .filter(key => !options?.prefix || key.startsWith(options.prefix))
+      .map(name => ({ name }));
+
+    const start = options?.cursor ? parseInt(options.cursor) : 0;
+    const limit = options?.limit ?? allKeys.length;
+    const end = start + limit;
+    const keys = allKeys.slice(start, end);
+    const list_complete = end >= allKeys.length;
+
+    return {
+      keys,
+      list_complete,
+      cursor: list_complete ? undefined : end.toString(),
+    };
+  }
+}
+
+export class CloudflareKVProvider {
+  private namespace: KVNamespace;
+
+  constructor(namespace: KVNamespace) {
+    this.namespace = namespace;
+  }
+
+  async get<T>(key: string): Promise<T | null> {
+    return this.namespace.get<T>(key, 'json');
+  }
+
+  async set(key: string, value: any) {
+    await this.namespace.put(key, JSON.stringify(value));
+  }
+
+  async del(key: string) {
+    await this.namespace.delete(key);
+  }
+
+  async sadd(key: string, value: string) {
+    const set = (await this.get<string[]>(key)) || [];
+    if (!set.includes(value)) {
+      set.push(value);
+      await this.set(key, set);
+      return 1;
+    }
+    return 0;
+  }
+
+  async srem(key: string, value: string) {
+    const set = (await this.get<string[]>(key)) || [];
+    const index = set.indexOf(value);
+    if (index !== -1) {
+      set.splice(index, 1);
+      await this.set(key, set);
+      return 1;
+    }
+    return 0;
+  }
+
+  async smembers(key: string): Promise<string[]> {
+    return (await this.get<string[]>(key)) || [];
+  }
+
+  async flushall() {
+    let cursor: string | undefined;
+    do {
+      const result = await this.namespace.list({ cursor });
+      await Promise.all(result.keys.map(key => this.namespace.delete(key.name)));
+      cursor = result.cursor;
+    } while (cursor);
+  }
+
+  async getWithVersion<T>(key: string): Promise<{ data: T | null; version: number }> {
+    const versionKey = `${key}:version`;
+    const [data, version] = await Promise.all([this.get<T>(key), this.get<number>(versionKey).then(v => v || 0)]);
+    return { data, version };
+  }
+
+  async setWithVersion<T>(key: string, value: T, expectedVersion: number): Promise<boolean> {
+    const versionKey = `${key}:version`;
+    const currentVersion = (await this.get<number>(versionKey)) || 0;
+
+    if (currentVersion !== expectedVersion) {
+      return false;
+    }
+
+    await Promise.all([this.set(key, value), this.set(versionKey, expectedVersion + 1)]);
+
+    return true;
+  }
+}

package/src/index.ts

@@ -0,0 +1,3 @@
+export * from './cloudflare';
+export * from './postgres';
+export * from './redis';

package/src/postgres/index.test.ts

@@ -0,0 +1,60 @@
+import dotenv from 'dotenv';
+
+import { PgMemory } from './';
+
+dotenv.config();
+
+const connectionString = process.env.DB_URL! || 'postgres://postgres:password@localhost:5434/mastra';
+
+describe('PgMastraMemory', () => {
+  let memory: PgMemory;
+
+  beforeAll(async () => {
+    memory = new PgMemory(connectionString);
+  });
+
+  afterAll(async () => {
+    await memory.drop();
+  });
+
+  it('should create and retrieve a thread', async () => {
+    const thread = await memory.createThread('Test Thread', { test: true });
+    const retrievedThread = await memory.getThreadById(thread.id);
+    expect(retrievedThread).toEqual(thread);
+  });
+
+  it('should save and retrieve messages', async () => {
+    const thread = await memory.createThread('Test Thread 2', { test: true });
+    const message1 = await memory.addMessage(thread.id, 'Hello', 'user');
+    // const message2 = await memory.addMessage(thread.id, 'World', 'assistant');
+    const messages = await memory.getMessages(thread.id);
+
+    console.log(messages);
+    expect(messages[0]?.content).toEqual(message1.content);
+  });
+
+  it('should update a thread', async () => {
+    const thread = await memory.createThread('Initial Thread Title', { test: true });
+    const updatedThread = await memory.updateThread(thread.id, 'Updated Thread Title', { test: true, updated: true });
+
+    expect(updatedThread.title).toEqual('Updated Thread Title');
+    expect(updatedThread.metadata).toEqual({ test: true, updated: true });
+  });
+
+  it('should delete a thread', async () => {
+    const thread = await memory.createThread('Thread to Delete', { test: true });
+    await memory.deleteThread(thread.id);
+
+    const retrievedThread = await memory.getThreadById(thread.id);
+    expect(retrievedThread).toBeNull();
+  });
+
+  it('should delete a message', async () => {
+    const thread = await memory.createThread('Thread with Message', { test: true });
+    const message = await memory.addMessage(thread.id, 'Message to Delete', 'user');
+    await memory.deleteMessage(message.id);
+
+    const messages = await memory.getMessages(thread.id);
+    expect(messages.length).toEqual(0);
+  });
+});