diff --git a/src/memory/__tests__/knowledge-graph.test.ts b/src/memory/__tests__/knowledge-graph.test.ts index 236242413a..cb1ab99a04 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); @@ -396,6 +398,35 @@ describe('KnowledgeGraphManager', () => { }); describe('file persistence', () => { + 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); + + 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'); + + await expect(fs.readFile(testFilePath, 'utf-8')).resolves.toBe( + originalFileContent + ); + + const files = await fs.readdir(path.dirname(testFilePath)); + expect( + files.filter(file => file.startsWith(`${path.basename(testFilePath)}.`)) + ).toEqual([]); + }); + it('should persist data across manager instances', async () => { await manager.createEntities([ { name: 'Alice', entityType: 'person', observations: ['persistent data'] }, diff --git a/src/memory/index.ts b/src/memory/index.ts index 9865c5318e..35360477e0 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,17 @@ export class KnowledgeGraphManager { relationType: r.relationType })), ]; - await fs.writeFile(this.memoryFilePath, lines.join("\n")); + const temporaryFilePath = `${this.memoryFilePath}.${randomBytes(16).toString('hex')}.tmp`; + + try { + await fs.writeFile(temporaryFilePath, lines.join("\n")); + await fs.rename(temporaryFilePath, this.memoryFilePath); + } catch (error) { + try { + await fs.unlink(temporaryFilePath); + } catch {} + throw error; + } } async createEntities(entities: Entity[]): Promise {