diff --git a/src/memory/__tests__/knowledge-graph.test.ts b/src/memory/__tests__/knowledge-graph.test.ts index 236242413a..b998a58614 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'; @@ -409,6 +409,62 @@ describe('KnowledgeGraphManager', () => { expect(graph.entities[0].name).toBe('Alice'); }); + it('should persist atomically via a temporary file and rename', async () => { + const renameSpy = vi.spyOn(fs, 'rename'); + try { + await manager.createEntities([ + { name: 'Alice', entityType: 'person', observations: ['atomic'] }, + ]); + + // saveGraph writes to a temp file, then atomically renames it into place. + expect(renameSpy).toHaveBeenCalledTimes(1); + const [tmpArg, destArg] = renameSpy.mock.calls[0]; + expect(String(destArg)).toBe(testFilePath); + // The temp file must live in the target's own directory so the rename + // stays on the same filesystem (and is therefore atomic); guard that + // invariant rather than just matching any ".tmp" suffix. + expect(String(tmpArg).startsWith(`${testFilePath}.`)).toBe(true); + expect(String(tmpArg).endsWith('.tmp')).toBe(true); + } finally { + renameSpy.mockRestore(); + } + + // Data round-trips and no temporary file is left behind. + const graph = await manager.readGraph(); + expect(graph.entities.map(e => e.name)).toContain('Alice'); + + const dir = path.dirname(testFilePath); + const base = path.basename(testFilePath); + const leftovers = (await fs.readdir(dir)).filter( + f => f.startsWith(base) && f.endsWith('.tmp') + ); + expect(leftovers).toEqual([]); + }); + + it('should clean up the temp file if the atomic rename fails', async () => { + const renameSpy = vi + .spyOn(fs, 'rename') + .mockRejectedValueOnce(new Error('simulated rename failure')); + try { + await expect( + manager.createEntities([ + { name: 'Alice', entityType: 'person', observations: [] }, + ]) + ).rejects.toThrow('simulated rename failure'); + } finally { + renameSpy.mockRestore(); + } + + // The temp file written before the failed rename is unlinked (the catch + // branch), so no partial file is left behind. + const dir = path.dirname(testFilePath); + const base = path.basename(testFilePath); + const leftovers = (await fs.readdir(dir)).filter( + f => f.startsWith(base) && f.endsWith('.tmp') + ); + expect(leftovers).toEqual([]); + }); + it('should handle JSONL format correctly', async () => { await manager.createEntities([ { name: 'Alice', entityType: 'person', observations: [] }, diff --git a/src/memory/index.ts b/src/memory/index.ts index 9865c5318e..9747da2be1 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,20 @@ export class KnowledgeGraphManager { relationType: r.relationType })), ]; - await fs.writeFile(this.memoryFilePath, lines.join("\n")); + // Write to a temporary file and atomically rename it into place so an + // interrupted or concurrent write cannot leave the memory file truncated + // or partially written. Mirrors the atomic-write pattern used by the + // filesystem server (src/filesystem/lib.ts). + const tmpPath = `${this.memoryFilePath}.${randomBytes(16).toString('hex')}.tmp`; + try { + await fs.writeFile(tmpPath, lines.join("\n")); + await fs.rename(tmpPath, this.memoryFilePath); + } catch (error) { + try { + await fs.unlink(tmpPath); + } catch {} + throw error; + } } async createEntities(entities: Entity[]): Promise {