diff --git a/src/memory/__tests__/knowledge-graph.test.ts b/src/memory/__tests__/knowledge-graph.test.ts index 236242413a..a8e777784f 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'; @@ -18,6 +18,8 @@ describe('KnowledgeGraphManager', () => { }); afterEach(async () => { + vi.restoreAllMocks(); + // Clean up test file try { await fs.unlink(testFilePath); @@ -409,6 +411,38 @@ describe('KnowledgeGraphManager', () => { expect(graph.entities[0].name).toBe('Alice'); }); + it('should preserve existing data when a write is interrupted', async () => { + await manager.createEntities([ + { name: 'Alice', entityType: 'person', observations: ['persistent data'] }, + ]); + + const originalFileContent = await fs.readFile(testFilePath, 'utf-8'); + const writeFile = fs.writeFile.bind(fs); + + // Simulate the process being killed mid-write: partial bytes written, then error + vi.spyOn(fs, 'writeFile').mockImplementation(async (file, data) => { + await writeFile(file, data.toString().slice(0, 1)); + throw new Error('interrupted write'); + }); + + await expect( + manager.createEntities([ + { name: 'Bob', entityType: 'person', observations: [] }, + ]) + ).rejects.toThrow('interrupted write'); + + // The original memory file must be untouched + await expect(fs.readFile(testFilePath, 'utf-8')).resolves.toBe( + originalFileContent + ); + + // No temp files may leak from the failed write + const files = await fs.readdir(path.dirname(testFilePath)); + expect( + files.filter(file => file.startsWith(`${path.basename(testFilePath)}.`)) + ).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..149391e19e 100644 --- a/src/memory/index.ts +++ b/src/memory/index.ts @@ -4,6 +4,7 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; import { SubscribeRequestSchema, UnsubscribeRequestSchema } from "@modelcontextprotocol/sdk/types.js"; import { z } from "zod"; +import { randomBytes } from 'crypto'; import { promises as fs } from 'fs'; import path from 'path'; import { fileURLToPath } from 'url'; @@ -114,7 +115,26 @@ export class KnowledgeGraphManager { relationType: r.relationType })), ]; - await fs.writeFile(this.memoryFilePath, lines.join("\n")); + // Write atomically: write to a unique temp file, then rename into place, so + // an interruption mid-write can never leave a truncated or corrupted file. + // The random suffix avoids collisions when saves overlap. + // + // Known limitation: a hard kill (SIGKILL) between the write and the rename + // leaves a stray .tmp file behind. This is unavoidable without a journal; + // the guarantee we provide is that memory.jsonl itself is never corrupted. + 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) { + // Clean up the temp file so a failed save never leaks artifacts. + try { + await fs.unlink(tmpPath); + } catch { + /* best-effort cleanup: never mask the original error */ + } + throw error; + } } async createEntities(entities: Entity[]): Promise {