flynn/src/memory/store.test.ts

import { describe, it, expect, beforeEach, afterEach } from 'vitest';
import { MemoryStore } from './store.js';
import { mkdtempSync, mkdirSync, writeFileSync, rmSync, readFileSync } from 'fs';
import { join } from 'path';
import { tmpdir } from 'os';

describe('MemoryStore', () => {
  let dir: string;
  let store: MemoryStore;

  beforeEach(() => {
    dir = mkdtempSync(join(tmpdir(), 'flynn-memory-test-'));
    store = new MemoryStore({ dir, maxContextTokens: 4096 });
  });

  afterEach(() => {
    rmSync(dir, { recursive: true, force: true });
  });

  describe('read', () => {
    it('returns empty string for non-existent namespace', () => {
      const result = store.read('nonexistent');
      expect(result).toBe('');
    });

    it('returns file contents for existing namespace', () => {
      // Pre-populate a memory file directly on disk
      const filePath = join(dir, 'notes.md');
      writeFileSync(filePath, 'Remember to check logs daily');

      const result = store.read('notes');
      expect(result).toBe('Remember to check logs daily');
    });

    it('handles nested namespaces (sessions/abc123)', () => {
      // Pre-populate a nested memory file
      const nestedDir = join(dir, 'sessions');
      mkdirSync(nestedDir, { recursive: true });
      writeFileSync(join(nestedDir, 'abc123.md'), 'Session context data');

      const result = store.read('sessions/abc123');
      expect(result).toBe('Session context data');
    });
  });

  describe('write', () => {
    it('creates file and directories on first write', () => {
      store.write('projects/alpha', 'Initial content', 'replace');

      const filePath = join(dir, 'projects', 'alpha.md');
      const content = readFileSync(filePath, 'utf-8');
      expect(content).toContain('Initial content');
    });

    it('appends content with newline separator in append mode', () => {
      store.write('journal', 'First entry', 'replace');
      store.write('journal', 'Second entry', 'append');

      const result = store.read('journal');
      expect(result).toContain('First entry');
      expect(result).toContain('Second entry');

      // Verify entries are separated by newline
      const lines = result.split('\n').filter((l: string) => l.trim().length > 0);
      expect(lines.length).toBeGreaterThanOrEqual(2);
    });

    it('replaces content in replace mode', () => {
      store.write('config', 'Old content', 'replace');
      store.write('config', 'New content', 'replace');

      const result = store.read('config');
      expect(result).toContain('New content');
      expect(result).not.toContain('Old content');
    });

    it('handles nested namespace paths', () => {
      store.write('deep/nested/path', 'Deep content', 'replace');

      const result = store.read('deep/nested/path');
      expect(result).toContain('Deep content');
    });
  });

  describe('category APIs', () => {
    it('reads and writes category namespaces', () => {
      store.writeCategory('user', 'facts', 'User lives in Berlin', 'replace');
      expect(store.readCategory('user', 'facts')).toBe('User lives in Berlin');
    });

    it('supports append and replace modes in category writes', () => {
      store.writeCategory('user', 'preferences', 'Prefers short answers', 'replace');
      store.writeCategory('user', 'preferences', 'Likes numbered lists', 'append');
      expect(store.readCategory('user', 'preferences')).toContain('Prefers short answers');
      expect(store.readCategory('user', 'preferences')).toContain('Likes numbered lists');

      store.writeCategory('user', 'preferences', 'Only this remains', 'replace');
      const content = store.readCategory('user', 'preferences');
      expect(content).toContain('Only this remains');
      expect(content).not.toContain('Prefers short answers');
    });

    it('lists only categories that exist under a base namespace', () => {
      store.writeCategory('user', 'facts', 'Fact', 'replace');
      store.writeCategory('user', 'projects', 'Project', 'replace');
      store.writeCategory('global', 'decisions', 'Decision', 'replace');

      expect(store.listCategories('user')).toEqual(['facts', 'projects']);
      expect(store.listCategories('global')).toEqual(['decisions']);
      expect(store.listCategories('sessions/abc')).toEqual([]);
    });

    it('reads all existing categories under a base namespace', () => {
      store.writeCategory('user', 'facts', 'Fact content', 'replace');
      store.writeCategory('user', 'decisions', 'Decision content', 'replace');

      expect(store.readAllCategories('user')).toEqual({
        facts: 'Fact content',
        decisions: 'Decision content',
      });
    });
  });

  describe('search', () => {
    beforeEach(() => {
      store.write('notes', 'The quick brown fox jumps over the lazy dog\nAnother line of text\nFox sightings are common here', 'replace');
      store.write('journal', 'Today I saw a fox in the garden\nIt was a beautiful day\nThe weather was sunny', 'replace');
      store.write('config', 'No relevant animals mentioned\nJust some settings here', 'replace');
    });

    it('finds matching lines case-insensitively', () => {
      const results = store.search('fox');

      // Should match lines containing "fox" regardless of case
      expect(results.length).toBeGreaterThanOrEqual(1);

      const contents = results.map((r) => r.content.toLowerCase());
      for (const content of contents) {
        expect(content).toContain('fox');
      }
    });

    it('returns context lines around matches', () => {
      const results = store.search('fox');

      // Each result should have a context field with surrounding text
      for (const result of results) {
        expect(result.context).toBeDefined();
        expect(result.context.length).toBeGreaterThan(0);
      }
    });

    it('searches across multiple namespaces', () => {
      const results = store.search('fox');

      // Should find matches in both 'notes' and 'journal' namespaces
      const namespaces = new Set(results.map((r) => r.namespace));
      expect(namespaces.size).toBeGreaterThanOrEqual(2);
    });

    it('returns empty array when no matches', () => {
      const results = store.search('xyznonexistent');
      expect(results).toEqual([]);
    });

    it('supports filtering by category', () => {
      store.writeCategory('user', 'facts', 'fox factual statement', 'replace');
      store.writeCategory('user', 'preferences', 'prefers fox metaphors', 'replace');

      const factOnly = store.search('fox', { categories: ['facts'] });
      expect(factOnly.length).toBeGreaterThan(0);
      expect(factOnly.every(result => result.namespace.endsWith('/facts'))).toBe(true);
    });

    it('supports filtering by base namespace prefix', () => {
      store.writeCategory('user', 'facts', 'fox in user facts', 'replace');
      store.writeCategory('global', 'facts', 'fox in global facts', 'replace');

      const userOnly = store.search('fox', { baseNamespacePrefix: 'user/' });
      expect(userOnly.length).toBeGreaterThan(0);
      expect(userOnly.every(result => result.namespace.startsWith('user/'))).toBe(true);
    });
  });

  describe('listNamespaces', () => {
    it('returns empty array when no files exist', () => {
      const namespaces = store.listNamespaces();
      expect(namespaces).toEqual([]);
    });

    it('lists all namespaces from existing files', () => {
      store.write('notes', 'Some notes', 'replace');
      store.write('journal', 'A journal entry', 'replace');
      store.write('config', 'Configuration data', 'replace');

      const namespaces = store.listNamespaces();

      expect(namespaces).toContain('notes');
      expect(namespaces).toContain('journal');
      expect(namespaces).toContain('config');
      expect(namespaces).toHaveLength(3);
    });

    it('includes nested namespaces', () => {
      store.write('sessions/abc123', 'Session data', 'replace');
      store.write('sessions/def456', 'Another session', 'replace');
      store.write('global', 'Global notes', 'replace');

      const namespaces = store.listNamespaces();

      expect(namespaces).toContain('sessions/abc123');
      expect(namespaces).toContain('sessions/def456');
      expect(namespaces).toContain('global');
    });
  });

  describe('getContextForPrompt', () => {
    it('returns empty string when no memory files exist', () => {
      const context = store.getContextForPrompt();
      expect(context).toBe('');
    });

    it('includes user and global memory under headings', () => {
      store.write('user', 'User prefers concise answers', 'replace');
      store.write('global', 'System-wide knowledge base', 'replace');
      store.writeCategory('user', 'facts', 'User timezone is UTC', 'replace');
      store.writeCategory('global', 'decisions', 'Adopt pnpm workspace', 'replace');

      const context = store.getContextForPrompt();

      // Should contain both memory sections with some form of heading
      expect(context).toContain('User Memory');
      expect(context).toContain('User prefers concise answers');
      expect(context).toContain('Global Memory');
      expect(context).toContain('System-wide knowledge base');
      expect(context).toContain('User Facts');
      expect(context).toContain('User timezone is UTC');
      expect(context).toContain('Global Decisions');
      expect(context).toContain('Adopt pnpm workspace');
    });

    it('truncates content to stay within maxContextTokens', () => {
      // Create a store with a very small token limit (50 tokens ≈ 200 chars)
      const smallStore = new MemoryStore({ dir, maxContextTokens: 50 });

      // Write a large amount of content to user and global (which getContextForPrompt reads)
      const longContent = 'A'.repeat(5000);
      smallStore.write('user', longContent, 'replace');
      smallStore.write('global', longContent, 'replace');

      const context = smallStore.getContextForPrompt();

      // 50 tokens * 4 chars/token = 200 chars max
      expect(context.length).toBeLessThanOrEqual(200);
      expect(context.length).toBeLessThan(longContent.length);
    });
  });
});