fix(memory): atomic graph persistence and validate entity existence in create_relations

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

@@ -1,4 +1,4 @@
-import { describe, it, expect, beforeEach, afterEach } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
 import { promises as fs } from 'fs';
 import path from 'path';
 import { fileURLToPath } from 'url';
@@ -103,6 +103,55 @@ describe('KnowledgeGraphManager', () => {
       const newRelations = await manager.createRelations([]);
       expect(newRelations).toHaveLength(0);
     });
+
+    it('should throw error when "from" entity does not exist', async () => {
+      await manager.createEntities([
+        { name: 'Bob', entityType: 'person', observations: [] },
+      ]);
+
+      await expect(
+        manager.createRelations([
+          { from: 'Alice', to: 'Bob', relationType: 'knows' },
+        ])
+      ).rejects.toThrow('Entity with name Alice not found');
+    });
+
+    it('should throw error when "to" entity does not exist', async () => {
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: [] },
+      ]);
+
+      await expect(
+        manager.createRelations([
+          { from: 'Alice', to: 'Bob', relationType: 'knows' },
+        ])
+      ).rejects.toThrow('Entity with name Bob not found');
+    });
+
+    it('should throw error when neither entity exists', async () => {
+      await expect(
+        manager.createRelations([
+          { from: 'Alice', to: 'Bob', relationType: 'knows' },
+        ])
+      ).rejects.toThrow('Entity with name Alice not found');
+    });
+
+    it('should fail fast and not create any relations if one relation in batch has missing entity', async () => {
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: [] },
+        { name: 'Bob', entityType: 'person', observations: [] },
+      ]);
+
+      await expect(
+        manager.createRelations([
+          { from: 'Alice', to: 'Bob', relationType: 'knows' },
+          { from: 'Alice', to: 'Charlie', relationType: 'knows' },
+        ])
+      ).rejects.toThrow('Entity with name Charlie not found');
+
+      const graph = await manager.readGraph();
+      expect(graph.relations).toHaveLength(0);
+    });
   });
 
   describe('addObservations', () => {
@@ -514,5 +563,49 @@ describe('KnowledgeGraphManager', () => {
       expect(result.relations).toHaveLength(1);
       expect(result.relations[0]).not.toHaveProperty('type');
     });
+
+    it('should write atomically using temporary file and rename', async () => {
+      const renameSpy = vi.spyOn(fs, 'rename');
+
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: ['atomic test'] },
+      ]);
+
+      expect(renameSpy).toHaveBeenCalledWith(
+        expect.stringMatching(new RegExp(`^${testFilePath.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}\\.[a-f0-9]{32}\\.tmp$`)),
+        testFilePath
+      );
+
+      renameSpy.mockRestore();
+    });
+
+    it('should clean up temp file and leave existing file intact if rename fails', async () => {
+      // Create initial entity successfully
+      await manager.createEntities([
+        { name: 'Alice', entityType: 'person', observations: ['initial'] },
+      ]);
+
+      const dir = path.dirname(testFilePath);
+      const renameSpy = vi.spyOn(fs, 'rename').mockRejectedValueOnce(new Error('Disk error during rename'));
+
+      await expect(
+        manager.createEntities([
+          { name: 'Bob', entityType: 'person', observations: ['new'] },
+        ])
+      ).rejects.toThrow('Disk error during rename');
+
+      // Verify no leftover .tmp files in test directory
+      const files = await fs.readdir(dir);
+      const tempFiles = files.filter(f => f.includes(path.basename(testFilePath)) && f.endsWith('.tmp'));
+      expect(tempFiles).toHaveLength(0);
+
+      renameSpy.mockRestore();
+
+      // Verify original file is still intact and valid
+      const manager2 = new KnowledgeGraphManager(testFilePath);
+      const graph = await manager2.readGraph();
+      expect(graph.entities).toHaveLength(1);
+      expect(graph.entities[0].name).toBe('Alice');
+    });
   });
 });

src/memory/index.ts

@@ -7,6 +7,7 @@ import { z } from "zod";
 import { promises as fs } from 'fs';
 import path from 'path';
 import { fileURLToPath } from 'url';
+import { randomBytes } from 'crypto';
 
 // Define memory file path using environment variable with fallback
 export const defaultMemoryPath = path.join(path.dirname(fileURLToPath(import.meta.url)), 'memory.jsonl');
@@ -114,7 +115,18 @@ export class KnowledgeGraphManager {
         relationType: r.relationType
       })),
     ];
-    await fs.writeFile(this.memoryFilePath, lines.join("\n"));
+    const tempFilePath = `${this.memoryFilePath}.${randomBytes(16).toString('hex')}.tmp`;
+    try {
+      await fs.writeFile(tempFilePath, lines.join("\n"), "utf-8");
+      await fs.rename(tempFilePath, this.memoryFilePath);
+    } catch (error) {
+      try {
+        await fs.unlink(tempFilePath);
+      } catch {
+        // Ignore unlink error if temp file was not created or already removed
+      }
+      throw error;
+    }
   }
 
   async createEntities(entities: Entity[]): Promise<Entity[]> {
@@ -127,6 +139,14 @@ export class KnowledgeGraphManager {
 
   async createRelations(relations: Relation[]): Promise<Relation[]> {
     const graph = await this.loadGraph();
+    for (const r of relations) {
+      if (!graph.entities.some(e => e.name === r.from)) {
+        throw new Error(`Entity with name ${r.from} not found`);
+      }
+      if (!graph.entities.some(e => e.name === r.to)) {
+        throw new Error(`Entity with name ${r.to} not found`);
+      }
+    }
     const newRelations = relations.filter(r => !graph.relations.some(existingRelation => 
       existingRelation.from === r.from && 
       existingRelation.to === r.to &&