From d8d57b85c19bd0a3a982e9a1e46049da8c412cb2 Mon Sep 17 00:00:00 2001 From: aayushbaluni <73417844+aayushbaluni@users.noreply.github.com> Date: Sat, 2 May 2026 00:00:38 +0530 Subject: [PATCH] fix(memory): include type field in output schemas to fix validation error The memory server stores entities with a 'type' field in JSONL, but the tool output schemas did not include this property. When structured output validation is enabled, this causes MCP error -32602 for read_graph, search_nodes, and open_nodes tools. Fixes modelcontextprotocol/servers#3074 --- src/memory/__tests__/knowledge-graph.test.ts | 23 +++++++++----------- src/memory/index.ts | 6 +++++ 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/src/memory/__tests__/knowledge-graph.test.ts b/src/memory/__tests__/knowledge-graph.test.ts index 236242413a..1e544783b9 100644 --- a/src/memory/__tests__/knowledge-graph.test.ts +++ b/src/memory/__tests__/knowledge-graph.test.ts @@ -426,7 +426,7 @@ describe('KnowledgeGraphManager', () => { expect(JSON.parse(lines[1])).toHaveProperty('type', 'relation'); }); - it('should strip type field from entities when loading from file', async () => { + it('should preserve JSONL type discriminator when loading from file', async () => { // Create entities and relations (these get saved with type field) await manager.createEntities([ { name: 'Alice', entityType: 'person', observations: ['test observation'] }, @@ -451,26 +451,25 @@ describe('KnowledgeGraphManager', () => { const manager2 = new KnowledgeGraphManager(testFilePath); const graph = await manager2.readGraph(); - // Verify loaded entities don't have type field + // Loaded graph mirrors JSONL discriminators for structured-output validation (#3074) expect(graph.entities).toHaveLength(2); graph.entities.forEach(entity => { - expect(entity).not.toHaveProperty('type'); + expect(entity).toHaveProperty('type', 'entity'); expect(entity).toHaveProperty('name'); expect(entity).toHaveProperty('entityType'); expect(entity).toHaveProperty('observations'); }); - // Verify loaded relations don't have type field expect(graph.relations).toHaveLength(1); graph.relations.forEach(relation => { - expect(relation).not.toHaveProperty('type'); + expect(relation).toHaveProperty('type', 'relation'); expect(relation).toHaveProperty('from'); expect(relation).toHaveProperty('to'); expect(relation).toHaveProperty('relationType'); }); }); - it('should strip type field from searchNodes results', async () => { + it('should include type discriminator in searchNodes results after reload', async () => { await manager.createEntities([ { name: 'Alice', entityType: 'person', observations: ['works at Acme'] }, ]); @@ -482,17 +481,16 @@ describe('KnowledgeGraphManager', () => { const manager2 = new KnowledgeGraphManager(testFilePath); const result = await manager2.searchNodes('Alice'); - // Verify search results don't have type field expect(result.entities).toHaveLength(1); - expect(result.entities[0]).not.toHaveProperty('type'); + expect(result.entities[0]).toHaveProperty('type', 'entity'); expect(result.entities[0].name).toBe('Alice'); expect(result.relations).toHaveLength(1); - expect(result.relations[0]).not.toHaveProperty('type'); + expect(result.relations[0]).toHaveProperty('type', 'relation'); expect(result.relations[0].from).toBe('Alice'); }); - it('should strip type field from openNodes results', async () => { + it('should include type discriminator in openNodes results after reload', async () => { await manager.createEntities([ { name: 'Alice', entityType: 'person', observations: [] }, { name: 'Bob', entityType: 'person', observations: [] }, @@ -505,14 +503,13 @@ describe('KnowledgeGraphManager', () => { const manager2 = new KnowledgeGraphManager(testFilePath); const result = await manager2.openNodes(['Alice', 'Bob']); - // Verify open results don't have type field expect(result.entities).toHaveLength(2); result.entities.forEach(entity => { - expect(entity).not.toHaveProperty('type'); + expect(entity).toHaveProperty('type', 'entity'); }); expect(result.relations).toHaveLength(1); - expect(result.relations[0]).not.toHaveProperty('type'); + expect(result.relations[0]).toHaveProperty('type', 'relation'); }); }); }); diff --git a/src/memory/index.ts b/src/memory/index.ts index b560bf1e53..3e483df8a4 100644 --- a/src/memory/index.ts +++ b/src/memory/index.ts @@ -48,12 +48,14 @@ let MEMORY_FILE_PATH: string; // We are storing our memory using entities, relations, and observations in a graph structure export interface Entity { + type?: "entity"; name: string; entityType: string; observations: string[]; } export interface Relation { + type?: "relation"; from: string; to: string; relationType: string; @@ -76,6 +78,7 @@ export class KnowledgeGraphManager { const item = JSON.parse(line); if (item.type === "entity") { graph.entities.push({ + type: "entity", name: item.name, entityType: item.entityType, observations: item.observations @@ -83,6 +86,7 @@ export class KnowledgeGraphManager { } if (item.type === "relation") { graph.relations.push({ + type: "relation", from: item.from, to: item.to, relationType: item.relationType @@ -241,12 +245,14 @@ let knowledgeGraphManager: KnowledgeGraphManager; // Zod schemas for entities and relations const EntitySchema = z.object({ + type: z.literal("entity").optional().describe('Discriminator when parsed from JSONL; omit when creating entities'), name: z.string().describe("The name of the entity"), entityType: z.string().describe("The type of the entity"), observations: z.array(z.string()).describe("An array of observation contents associated with the entity") }); const RelationSchema = z.object({ + type: z.literal("relation").optional().describe('Discriminator when parsed from JSONL; omit when creating relations'), from: z.string().describe("The name of the entity where the relation starts"), to: z.string().describe("The name of the entity where the relation ends"), relationType: z.string().describe("The type of the relation")