diff --git a/src/memory/__tests__/knowledge-graph.test.ts b/src/memory/__tests__/knowledge-graph.test.ts index 236242413a..1197c1e57f 100644 --- a/src/memory/__tests__/knowledge-graph.test.ts +++ b/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'); + }); }); }); diff --git a/src/memory/index.ts b/src/memory/index.ts index 9865c5318e..19e438c2ee 100644 --- a/src/memory/index.ts +++ b/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 { @@ -127,6 +139,14 @@ export class KnowledgeGraphManager { async createRelations(relations: Relation[]): Promise { 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 &&