Add Vitest testing framework and implement tests for memory management features

package-lock.json

@@ -4284,8 +4284,10 @@
       },
       "devDependencies": {
         "@types/node": "^22",
+        "@vitest/coverage-v8": "^2.1.8",
         "shx": "^0.3.4",
-        "typescript": "^5.6.2"
+        "typescript": "^5.6.2",
+        "vitest": "^2.1.8"
       }
     },
     "src/postgres": {

src/memory/__tests__/file-path.test.ts

@@ -0,0 +1,156 @@
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
+import { promises as fs } from 'fs';
+import path from 'path';
+import { fileURLToPath } from 'url';
+import { ensureMemoryFilePath, defaultMemoryPath } from '../index.js';
+
+describe('ensureMemoryFilePath', () => {
+  const testDir = path.dirname(fileURLToPath(import.meta.url));
+  const oldMemoryPath = path.join(testDir, '..', 'memory.json');
+  const newMemoryPath = path.join(testDir, '..', 'memory.jsonl');
+
+  let originalEnv: string | undefined;
+
+  beforeEach(() => {
+    // Save original environment variable
+    originalEnv = process.env.MEMORY_FILE_PATH;
+    // Delete environment variable
+    delete process.env.MEMORY_FILE_PATH;
+  });
+
+  afterEach(async () => {
+    // Restore original environment variable
+    if (originalEnv !== undefined) {
+      process.env.MEMORY_FILE_PATH = originalEnv;
+    } else {
+      delete process.env.MEMORY_FILE_PATH;
+    }
+
+    // Clean up test files
+    try {
+      await fs.unlink(oldMemoryPath);
+    } catch {
+      // Ignore if file doesn't exist
+    }
+    try {
+      await fs.unlink(newMemoryPath);
+    } catch {
+      // Ignore if file doesn't exist
+    }
+  });
+
+  describe('with MEMORY_FILE_PATH environment variable', () => {
+    it('should return absolute path when MEMORY_FILE_PATH is absolute', async () => {
+      const absolutePath = '/tmp/custom-memory.jsonl';
+      process.env.MEMORY_FILE_PATH = absolutePath;
+
+      const result = await ensureMemoryFilePath();
+
+      expect(result).toBe(absolutePath);
+    });
+
+    it('should convert relative path to absolute when MEMORY_FILE_PATH is relative', async () => {
+      const relativePath = 'custom-memory.jsonl';
+      process.env.MEMORY_FILE_PATH = relativePath;
+
+      const result = await ensureMemoryFilePath();
+
+      expect(path.isAbsolute(result)).toBe(true);
+      expect(result).toContain('custom-memory.jsonl');
+    });
+
+    it('should handle Windows absolute paths', async () => {
+      const windowsPath = 'C:\\temp\\memory.jsonl';
+      process.env.MEMORY_FILE_PATH = windowsPath;
+
+      const result = await ensureMemoryFilePath();
+
+      // On Windows, should return as-is; on Unix, will be treated as relative
+      if (process.platform === 'win32') {
+        expect(result).toBe(windowsPath);
+      } else {
+        expect(path.isAbsolute(result)).toBe(true);
+      }
+    });
+  });
+
+  describe('without MEMORY_FILE_PATH environment variable', () => {
+    it('should return default path when no files exist', async () => {
+      const result = await ensureMemoryFilePath();
+
+      expect(result).toBe(defaultMemoryPath);
+    });
+
+    it('should migrate from memory.json to memory.jsonl when only old file exists', async () => {
+      // Create old memory.json file
+      await fs.writeFile(oldMemoryPath, '{"test":"data"}');
+
+      const consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
+
+      const result = await ensureMemoryFilePath();
+
+      expect(result).toBe(defaultMemoryPath);
+
+      // Verify migration happened
+      const newFileExists = await fs.access(newMemoryPath).then(() => true).catch(() => false);
+      const oldFileExists = await fs.access(oldMemoryPath).then(() => true).catch(() => false);
+
+      expect(newFileExists).toBe(true);
+      expect(oldFileExists).toBe(false);
+
+      // Verify console messages
+      expect(consoleErrorSpy).toHaveBeenCalledWith(
+        expect.stringContaining('DETECTED: Found legacy memory.json file')
+      );
+      expect(consoleErrorSpy).toHaveBeenCalledWith(
+        expect.stringContaining('COMPLETED: Successfully migrated')
+      );
+
+      consoleErrorSpy.mockRestore();
+    });
+
+    it('should use new file when both old and new files exist', async () => {
+      // Create both files
+      await fs.writeFile(oldMemoryPath, '{"old":"data"}');
+      await fs.writeFile(newMemoryPath, '{"new":"data"}');
+
+      const consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
+
+      const result = await ensureMemoryFilePath();
+
+      expect(result).toBe(defaultMemoryPath);
+
+      // Verify no migration happened (both files should still exist)
+      const newFileExists = await fs.access(newMemoryPath).then(() => true).catch(() => false);
+      const oldFileExists = await fs.access(oldMemoryPath).then(() => true).catch(() => false);
+
+      expect(newFileExists).toBe(true);
+      expect(oldFileExists).toBe(true);
+
+      // Verify no console messages about migration
+      expect(consoleErrorSpy).not.toHaveBeenCalled();
+
+      consoleErrorSpy.mockRestore();
+    });
+
+    it('should preserve file content during migration', async () => {
+      const testContent = '{"entities": [{"name": "test", "type": "person"}]}';
+      await fs.writeFile(oldMemoryPath, testContent);
+
+      await ensureMemoryFilePath();
+
+      const migratedContent = await fs.readFile(newMemoryPath, 'utf-8');
+      expect(migratedContent).toBe(testContent);
+    });
+  });
+
+  describe('defaultMemoryPath', () => {
+    it('should end with memory.jsonl', () => {
+      expect(defaultMemoryPath).toMatch(/memory\.jsonl$/);
+    });
+
+    it('should be an absolute path', () => {
+      expect(path.isAbsolute(defaultMemoryPath)).toBe(true);
+    });
+  });
+});

src/memory/__tests__/knowledge-graph.test.ts

@@ -0,0 +1,394 @@
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
+import { promises as fs } from 'fs';
+import path from 'path';
+import { fileURLToPath } from 'url';
+import { KnowledgeGraphManager, Entity, Relation, KnowledgeGraph } from '../index.js';
+
+describe('KnowledgeGraphManager', () => {
+  let manager: KnowledgeGraphManager;
+  let testFilePath: string;
+
+  beforeEach(async () => {
+    // Create a temporary test file path
+    testFilePath = path.join(
+      path.dirname(fileURLToPath(import.meta.url)),
+      `test-memory-${Date.now()}.jsonl`
+    );
+    manager = new KnowledgeGraphManager(testFilePath);
+  });
+
+  afterEach(async () => {
+    // Clean up test file
+    try {
+      await fs.unlink(testFilePath);
+    } catch (error) {
+      // Ignore errors if file doesn't exist
+    }
+  });
+
+  describe('createEntities', () => {
+    it('should create new entities', async () => {
+      const entities: Entity[] = [
+        { name: 'Alice', entityType: 'person', observations: ['works at Acme Corp'] },
+        { name: 'Bob', entityType: 'person', observations: ['likes programming'] },
+      ];
+
+      const newEntities = await manager.createEntities(entities);
+      expect(newEntities).toHaveLength(2);
+      expect(newEntities).toEqual(entities);
+
+      const graph = await manager.readGraph();
+      expect(graph.entities).toHaveLength(2);
+    });
+
+    it('should not create duplicate entities', async () => {
+      const entities: Entity[] = [
+        { name: 'Alice', entityType: 'person', observations: ['works at Acme Corp'] },
+      ];
+
+      await manager.createEntities(entities);
+      const newEntities = await manager.createEntities(entities);
+
+      expect(newEntities).toHaveLength(0);
+
+      const graph = await manager.readGraph();
+      expect(graph.entities).toHaveLength(1);
+    });
+
+    it('should handle empty entity arrays', async () => {
+      const newEntities = await manager.createEntities([]);
+      expect(newEntities).toHaveLength(0);
+    });
+  });
+
+  describe('createRelations', () => {
+    it('should create new relations', async () => {
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: [] },
+        { name: 'Bob', entityType: 'person', observations: [] },
+      ]);
+
+      const relations: Relation[] = [
+        { from: 'Alice', to: 'Bob', relationType: 'knows' },
+      ];
+
+      const newRelations = await manager.createRelations(relations);
+      expect(newRelations).toHaveLength(1);
+      expect(newRelations).toEqual(relations);
+
+      const graph = await manager.readGraph();
+      expect(graph.relations).toHaveLength(1);
+    });
+
+    it('should not create duplicate relations', async () => {
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: [] },
+        { name: 'Bob', entityType: 'person', observations: [] },
+      ]);
+
+      const relations: Relation[] = [
+        { from: 'Alice', to: 'Bob', relationType: 'knows' },
+      ];
+
+      await manager.createRelations(relations);
+      const newRelations = await manager.createRelations(relations);
+
+      expect(newRelations).toHaveLength(0);
+
+      const graph = await manager.readGraph();
+      expect(graph.relations).toHaveLength(1);
+    });
+
+    it('should handle empty relation arrays', async () => {
+      const newRelations = await manager.createRelations([]);
+      expect(newRelations).toHaveLength(0);
+    });
+  });
+
+  describe('addObservations', () => {
+    it('should add observations to existing entities', async () => {
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: ['works at Acme Corp'] },
+      ]);
+
+      const results = await manager.addObservations([
+        { entityName: 'Alice', contents: ['likes coffee', 'has a dog'] },
+      ]);
+
+      expect(results).toHaveLength(1);
+      expect(results[0].entityName).toBe('Alice');
+      expect(results[0].addedObservations).toHaveLength(2);
+
+      const graph = await manager.readGraph();
+      const alice = graph.entities.find(e => e.name === 'Alice');
+      expect(alice?.observations).toHaveLength(3);
+    });
+
+    it('should not add duplicate observations', async () => {
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: ['works at Acme Corp'] },
+      ]);
+
+      await manager.addObservations([
+        { entityName: 'Alice', contents: ['likes coffee'] },
+      ]);
+
+      const results = await manager.addObservations([
+        { entityName: 'Alice', contents: ['likes coffee', 'has a dog'] },
+      ]);
+
+      expect(results[0].addedObservations).toHaveLength(1);
+      expect(results[0].addedObservations).toContain('has a dog');
+
+      const graph = await manager.readGraph();
+      const alice = graph.entities.find(e => e.name === 'Alice');
+      expect(alice?.observations).toHaveLength(3);
+    });
+
+    it('should throw error for non-existent entity', async () => {
+      await expect(
+        manager.addObservations([
+          { entityName: 'NonExistent', contents: ['some observation'] },
+        ])
+      ).rejects.toThrow('Entity with name NonExistent not found');
+    });
+  });
+
+  describe('deleteEntities', () => {
+    it('should delete entities', async () => {
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: [] },
+        { name: 'Bob', entityType: 'person', observations: [] },
+      ]);
+
+      await manager.deleteEntities(['Alice']);
+
+      const graph = await manager.readGraph();
+      expect(graph.entities).toHaveLength(1);
+      expect(graph.entities[0].name).toBe('Bob');
+    });
+
+    it('should cascade delete relations when deleting entities', async () => {
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: [] },
+        { name: 'Bob', entityType: 'person', observations: [] },
+        { name: 'Charlie', entityType: 'person', observations: [] },
+      ]);
+
+      await manager.createRelations([
+        { from: 'Alice', to: 'Bob', relationType: 'knows' },
+        { from: 'Bob', to: 'Charlie', relationType: 'knows' },
+      ]);
+
+      await manager.deleteEntities(['Bob']);
+
+      const graph = await manager.readGraph();
+      expect(graph.entities).toHaveLength(2);
+      expect(graph.relations).toHaveLength(0);
+    });
+
+    it('should handle deleting non-existent entities', async () => {
+      await manager.deleteEntities(['NonExistent']);
+      const graph = await manager.readGraph();
+      expect(graph.entities).toHaveLength(0);
+    });
+  });
+
+  describe('deleteObservations', () => {
+    it('should delete observations from entities', async () => {
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: ['works at Acme Corp', 'likes coffee'] },
+      ]);
+
+      await manager.deleteObservations([
+        { entityName: 'Alice', observations: ['likes coffee'] },
+      ]);
+
+      const graph = await manager.readGraph();
+      const alice = graph.entities.find(e => e.name === 'Alice');
+      expect(alice?.observations).toHaveLength(1);
+      expect(alice?.observations).toContain('works at Acme Corp');
+    });
+
+    it('should handle deleting from non-existent entities', async () => {
+      await manager.deleteObservations([
+        { entityName: 'NonExistent', observations: ['some observation'] },
+      ]);
+      // Should not throw error
+      const graph = await manager.readGraph();
+      expect(graph.entities).toHaveLength(0);
+    });
+  });
+
+  describe('deleteRelations', () => {
+    it('should delete specific relations', async () => {
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: [] },
+        { name: 'Bob', entityType: 'person', observations: [] },
+      ]);
+
+      await manager.createRelations([
+        { from: 'Alice', to: 'Bob', relationType: 'knows' },
+        { from: 'Alice', to: 'Bob', relationType: 'works_with' },
+      ]);
+
+      await manager.deleteRelations([
+        { from: 'Alice', to: 'Bob', relationType: 'knows' },
+      ]);
+
+      const graph = await manager.readGraph();
+      expect(graph.relations).toHaveLength(1);
+      expect(graph.relations[0].relationType).toBe('works_with');
+    });
+  });
+
+  describe('readGraph', () => {
+    it('should return empty graph when file does not exist', async () => {
+      const graph = await manager.readGraph();
+      expect(graph.entities).toHaveLength(0);
+      expect(graph.relations).toHaveLength(0);
+    });
+
+    it('should return complete graph with entities and relations', async () => {
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: ['works at Acme Corp'] },
+      ]);
+
+      await manager.createRelations([
+        { from: 'Alice', to: 'Alice', relationType: 'self' },
+      ]);
+
+      const graph = await manager.readGraph();
+      expect(graph.entities).toHaveLength(1);
+      expect(graph.relations).toHaveLength(1);
+    });
+  });
+
+  describe('searchNodes', () => {
+    beforeEach(async () => {
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: ['works at Acme Corp', 'likes programming'] },
+        { name: 'Bob', entityType: 'person', observations: ['works at TechCo'] },
+        { name: 'Acme Corp', entityType: 'company', observations: ['tech company'] },
+      ]);
+
+      await manager.createRelations([
+        { from: 'Alice', to: 'Acme Corp', relationType: 'works_at' },
+        { from: 'Bob', to: 'Acme Corp', relationType: 'competitor' },
+      ]);
+    });
+
+    it('should search by entity name', async () => {
+      const result = await manager.searchNodes('Alice');
+      expect(result.entities).toHaveLength(1);
+      expect(result.entities[0].name).toBe('Alice');
+    });
+
+    it('should search by entity type', async () => {
+      const result = await manager.searchNodes('company');
+      expect(result.entities).toHaveLength(1);
+      expect(result.entities[0].name).toBe('Acme Corp');
+    });
+
+    it('should search by observation content', async () => {
+      const result = await manager.searchNodes('programming');
+      expect(result.entities).toHaveLength(1);
+      expect(result.entities[0].name).toBe('Alice');
+    });
+
+    it('should be case insensitive', async () => {
+      const result = await manager.searchNodes('ALICE');
+      expect(result.entities).toHaveLength(1);
+      expect(result.entities[0].name).toBe('Alice');
+    });
+
+    it('should include relations between matched entities', async () => {
+      const result = await manager.searchNodes('Acme');
+      expect(result.entities).toHaveLength(2); // Alice and Acme Corp
+      expect(result.relations).toHaveLength(1); // Only Alice -> Acme Corp relation
+    });
+
+    it('should return empty graph for no matches', async () => {
+      const result = await manager.searchNodes('NonExistent');
+      expect(result.entities).toHaveLength(0);
+      expect(result.relations).toHaveLength(0);
+    });
+  });
+
+  describe('openNodes', () => {
+    beforeEach(async () => {
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: [] },
+        { name: 'Bob', entityType: 'person', observations: [] },
+        { name: 'Charlie', entityType: 'person', observations: [] },
+      ]);
+
+      await manager.createRelations([
+        { from: 'Alice', to: 'Bob', relationType: 'knows' },
+        { from: 'Bob', to: 'Charlie', relationType: 'knows' },
+      ]);
+    });
+
+    it('should open specific nodes by name', async () => {
+      const result = await manager.openNodes(['Alice', 'Bob']);
+      expect(result.entities).toHaveLength(2);
+      expect(result.entities.map(e => e.name)).toContain('Alice');
+      expect(result.entities.map(e => e.name)).toContain('Bob');
+    });
+
+    it('should include relations between opened nodes', async () => {
+      const result = await manager.openNodes(['Alice', 'Bob']);
+      expect(result.relations).toHaveLength(1);
+      expect(result.relations[0].from).toBe('Alice');
+      expect(result.relations[0].to).toBe('Bob');
+    });
+
+    it('should exclude relations to unopened nodes', async () => {
+      const result = await manager.openNodes(['Bob']);
+      expect(result.relations).toHaveLength(0);
+    });
+
+    it('should handle opening non-existent nodes', async () => {
+      const result = await manager.openNodes(['NonExistent']);
+      expect(result.entities).toHaveLength(0);
+    });
+
+    it('should handle empty node list', async () => {
+      const result = await manager.openNodes([]);
+      expect(result.entities).toHaveLength(0);
+      expect(result.relations).toHaveLength(0);
+    });
+  });
+
+  describe('file persistence', () => {
+    it('should persist data across manager instances', async () => {
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: ['persistent data'] },
+      ]);
+
+      // Create new manager instance with same file path
+      const manager2 = new KnowledgeGraphManager(testFilePath);
+      const graph = await manager2.readGraph();
+
+      expect(graph.entities).toHaveLength(1);
+      expect(graph.entities[0].name).toBe('Alice');
+    });
+
+    it('should handle JSONL format correctly', async () => {
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: [] },
+      ]);
+      await manager.createRelations([
+        { from: 'Alice', to: 'Alice', relationType: 'self' },
+      ]);
+
+      // Read file directly
+      const fileContent = await fs.readFile(testFilePath, 'utf-8');
+      const lines = fileContent.split('\n').filter(line => line.trim());
+
+      expect(lines).toHaveLength(2);
+      expect(JSON.parse(lines[0])).toHaveProperty('type', 'entity');
+      expect(JSON.parse(lines[1])).toHaveProperty('type', 'relation');
+    });
+  });
+});

src/memory/index.ts

@@ -11,10 +11,10 @@ import path from 'path';
 import { fileURLToPath } from 'url';
 
 // Define memory file path using environment variable with fallback
-const defaultMemoryPath = path.join(path.dirname(fileURLToPath(import.meta.url)), 'memory.jsonl');
+export const defaultMemoryPath = path.join(path.dirname(fileURLToPath(import.meta.url)), 'memory.jsonl');
 
 // Handle backward compatibility: migrate memory.json to memory.jsonl if needed
-async function ensureMemoryFilePath(): Promise<string> {
+export async function ensureMemoryFilePath(): Promise<string> {
   if (process.env.MEMORY_FILE_PATH) {
     // Custom path provided, use it as-is (with absolute path resolution)
     return path.isAbsolute(process.env.MEMORY_FILE_PATH)
@@ -50,28 +50,30 @@ async function ensureMemoryFilePath(): Promise<string> {
 let MEMORY_FILE_PATH: string;
 
 // We are storing our memory using entities, relations, and observations in a graph structure
-interface Entity {
+export interface Entity {
   name: string;
   entityType: string;
   observations: string[];
 }
 
-interface Relation {
+export interface Relation {
   from: string;
   to: string;
   relationType: string;
 }
 
-interface KnowledgeGraph {
+export interface KnowledgeGraph {
   entities: Entity[];
   relations: Relation[];
 }
 
 // The KnowledgeGraphManager class contains all operations to interact with the knowledge graph
-class KnowledgeGraphManager {
+export class KnowledgeGraphManager {
+  constructor(private memoryFilePath: string) {}
+
   private async loadGraph(): Promise<KnowledgeGraph> {
     try {
-      const data = await fs.readFile(MEMORY_FILE_PATH, "utf-8");
+      const data = await fs.readFile(this.memoryFilePath, "utf-8");
       const lines = data.split("\n").filter(line => line.trim() !== "");
       return lines.reduce((graph: KnowledgeGraph, line) => {
         const item = JSON.parse(line);
@@ -89,20 +91,20 @@ class KnowledgeGraphManager {
 
   private async saveGraph(graph: KnowledgeGraph): Promise<void> {
     const lines = [
-      ...graph.entities.map(e => JSON.stringify({ 
-        type: "entity", 
-        name: e.name, 
-        entityType: e.entityType, 
-        observations: e.observations 
+      ...graph.entities.map(e => JSON.stringify({
+        type: "entity",
+        name: e.name,
+        entityType: e.entityType,
+        observations: e.observations
       })),
-      ...graph.relations.map(r => JSON.stringify({ 
-        type: "relation", 
-        from: r.from, 
-        to: r.to, 
-        relationType: r.relationType 
+      ...graph.relations.map(r => JSON.stringify({
+        type: "relation",
+        from: r.from,
+        to: r.to,
+        relationType: r.relationType
       })),
     ];
-    await fs.writeFile(MEMORY_FILE_PATH, lines.join("\n"));
+    await fs.writeFile(this.memoryFilePath, lines.join("\n"));
   }
 
   async createEntities(entities: Entity[]): Promise<Entity[]> {
@@ -222,7 +224,7 @@ class KnowledgeGraphManager {
   }
 }
 
-const knowledgeGraphManager = new KnowledgeGraphManager();
+let knowledgeGraphManager: KnowledgeGraphManager;
 
 
 // The server instance and tools exposed to Claude
@@ -465,7 +467,10 @@ server.setRequestHandler(CallToolRequestSchema, async (request) => {
 async function main() {
   // Initialize memory file path with backward compatibility
   MEMORY_FILE_PATH = await ensureMemoryFilePath();
-  
+
+  // Initialize knowledge graph manager with the memory file path
+  knowledgeGraphManager = new KnowledgeGraphManager(MEMORY_FILE_PATH);
+
   const transport = new StdioServerTransport();
   await server.connect(transport);
   console.error("Knowledge Graph MCP Server running on stdio");

src/memory/package.json

@@ -16,14 +16,17 @@
   "scripts": {
     "build": "tsc && shx chmod +x dist/*.js",
     "prepare": "npm run build",
-    "watch": "tsc --watch"
+    "watch": "tsc --watch",
+    "test": "vitest run --coverage"
   },
   "dependencies": {
     "@modelcontextprotocol/sdk": "^1.19.1"
   },
   "devDependencies": {
     "@types/node": "^22",
+    "@vitest/coverage-v8": "^2.1.8",
     "shx": "^0.3.4",
-    "typescript": "^5.6.2"
+    "typescript": "^5.6.2",
+    "vitest": "^2.1.8"
   }
 }

src/memory/vitest.config.ts

@@ -0,0 +1,14 @@
+import { defineConfig } from 'vitest/config';
+
+export default defineConfig({
+  test: {
+    globals: true,
+    environment: 'node',
+    include: ['**/__tests__/**/*.test.ts'],
+    coverage: {
+      provider: 'v8',
+      include: ['**/*.ts'],
+      exclude: ['**/__tests__/**', '**/dist/**'],
+    },
+  },
+});