diff --git a/src/memory/README.md b/src/memory/README.md index fdb8c4613e..e9fb37e241 100644 --- a/src/memory/README.md +++ b/src/memory/README.md @@ -125,6 +125,14 @@ Example: - Relations between requested entities - Silently skips non-existent nodes +### Resources + +- **knowledge-graph** (`memory://knowledge-graph`) + - The full knowledge graph as a readable MCP Resource + - MIME type: `application/json` + - Returns the same shape as `read_graph` (entities and relations) + - Mutation tools (`create_entities`, `create_relations`, `add_observations`, `delete_entities`, `delete_observations`, `delete_relations`) emit `notifications/resources/updated` for this URI, so subscribed clients see live changes + # Usage with Claude Desktop ### Setup diff --git a/src/memory/__tests__/resource.test.ts b/src/memory/__tests__/resource.test.ts new file mode 100644 index 0000000000..432633050b --- /dev/null +++ b/src/memory/__tests__/resource.test.ts @@ -0,0 +1,97 @@ +import { describe, it, expect, vi } from 'vitest'; +import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; +import { SubscribeRequestSchema, UnsubscribeRequestSchema } from '@modelcontextprotocol/sdk/types.js'; +import { + KnowledgeGraphManager, + registerKnowledgeGraphResource, + registerKnowledgeGraphSubscriptions, +} from '../index.js'; + +describe('knowledge-graph resource', () => { + it('registers with kebab-case name, correct URI, and JSON mime type', () => { + const mockServer = { registerResource: vi.fn() } as unknown as McpServer; + const manager = {} as KnowledgeGraphManager; + + registerKnowledgeGraphResource(mockServer, manager); + + expect(mockServer.registerResource).toHaveBeenCalledWith( + 'knowledge-graph', + 'memory://knowledge-graph', + expect.objectContaining({ + title: 'Knowledge Graph', + mimeType: 'application/json', + }), + expect.any(Function), + ); + }); + + it('handler returns the graph as JSON in the contents array', async () => { + const mockServer = { registerResource: vi.fn() } as unknown as McpServer; + const fakeGraph = { + entities: [{ name: 'Alice', entityType: 'person', observations: ['engineer'] }], + relations: [{ from: 'Alice', to: 'Acme', relationType: 'works_at' }], + }; + const manager = { + readGraph: vi.fn().mockResolvedValue(fakeGraph), + } as unknown as KnowledgeGraphManager; + + registerKnowledgeGraphResource(mockServer, manager); + + const handler = (mockServer.registerResource as ReturnType).mock.calls[0][3]; + const result = await handler(new URL('memory://knowledge-graph')); + + expect(result.contents).toHaveLength(1); + expect(result.contents[0].uri).toBe('memory://knowledge-graph'); + expect(result.contents[0].mimeType).toBe('application/json'); + expect(JSON.parse(result.contents[0].text)).toEqual(fakeGraph); + expect(manager.readGraph).toHaveBeenCalledOnce(); + }); +}); + +describe('knowledge-graph resource subscriptions', () => { + function makeMockServer() { + const inner = { + registerCapabilities: vi.fn(), + setRequestHandler: vi.fn(), + sendResourceUpdated: vi.fn(), + }; + const mockServer = { server: inner } as unknown as McpServer; + return { mockServer, inner }; + } + + function handlerFor(inner: ReturnType['inner'], schema: unknown) { + const call = inner.setRequestHandler.mock.calls.find((c) => c[0] === schema); + if (!call) throw new Error('handler not registered'); + return call[1] as (request: { params: { uri: string } }) => Promise; + } + + it('declares the resources.subscribe capability', () => { + const { mockServer, inner } = makeMockServer(); + + registerKnowledgeGraphSubscriptions(mockServer); + + expect(inner.registerCapabilities).toHaveBeenCalledWith({ + resources: { subscribe: true }, + }); + }); + + it('registers subscribe and unsubscribe request handlers', () => { + const { mockServer, inner } = makeMockServer(); + + registerKnowledgeGraphSubscriptions(mockServer); + + const schemas = inner.setRequestHandler.mock.calls.map((c) => c[0]); + expect(schemas).toContain(SubscribeRequestSchema); + expect(schemas).toContain(UnsubscribeRequestSchema); + }); + + it('subscribe and unsubscribe handlers acknowledge with an empty result', async () => { + const { mockServer, inner } = makeMockServer(); + + registerKnowledgeGraphSubscriptions(mockServer); + + const req = { params: { uri: 'memory://knowledge-graph' } }; + await expect(handlerFor(inner, SubscribeRequestSchema)(req)).resolves.toEqual({}); + await expect(handlerFor(inner, UnsubscribeRequestSchema)(req)).resolves.toEqual({}); + }); +}); diff --git a/src/memory/index.ts b/src/memory/index.ts index 7b4c683300..9865c5318e 100644 --- a/src/memory/index.ts +++ b/src/memory/index.ts @@ -2,6 +2,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 { promises as fs } from 'fs'; import path from 'path'; @@ -258,6 +259,20 @@ const server = new McpServer({ version: "0.6.3", }); +const RESOURCE_URI = "memory://knowledge-graph"; + +// Track which resource URIs the connected client has subscribed to, so we only +// emit notifications/resources/updated to a client that asked for them. +const resourceSubscribers = new Set(); + +// Notify subscribers that the knowledge graph resource changed. No-op when the +// client has not subscribed. +function notifyGraphUpdated() { + if (resourceSubscribers.has(RESOURCE_URI)) { + server.server.sendResourceUpdated({ uri: RESOURCE_URI }); + } +} + // Register create_entities tool server.registerTool( "create_entities", @@ -279,6 +294,7 @@ server.registerTool( }, async ({ entities }) => { const result = await knowledgeGraphManager.createEntities(entities); + notifyGraphUpdated(); return { content: [{ type: "text" as const, text: JSON.stringify(result, null, 2) }], structuredContent: { entities: result } @@ -307,6 +323,7 @@ server.registerTool( }, async ({ relations }) => { const result = await knowledgeGraphManager.createRelations(relations); + notifyGraphUpdated(); return { content: [{ type: "text" as const, text: JSON.stringify(result, null, 2) }], structuredContent: { relations: result } @@ -341,6 +358,7 @@ server.registerTool( }, async ({ observations }) => { const result = await knowledgeGraphManager.addObservations(observations); + notifyGraphUpdated(); return { content: [{ type: "text" as const, text: JSON.stringify(result, null, 2) }], structuredContent: { results: result } @@ -370,6 +388,7 @@ server.registerTool( }, async ({ entityNames }) => { await knowledgeGraphManager.deleteEntities(entityNames); + notifyGraphUpdated(); return { content: [{ type: "text" as const, text: "Entities deleted successfully" }], structuredContent: { success: true, message: "Entities deleted successfully" } @@ -402,6 +421,7 @@ server.registerTool( }, async ({ deletions }) => { await knowledgeGraphManager.deleteObservations(deletions); + notifyGraphUpdated(); return { content: [{ type: "text" as const, text: "Observations deleted successfully" }], structuredContent: { success: true, message: "Observations deleted successfully" } @@ -431,6 +451,7 @@ server.registerTool( }, async ({ relations }) => { await knowledgeGraphManager.deleteRelations(relations); + notifyGraphUpdated(); return { content: [{ type: "text" as const, text: "Relations deleted successfully" }], structuredContent: { success: true, message: "Relations deleted successfully" } @@ -523,12 +544,52 @@ server.registerTool( } ); +export function registerKnowledgeGraphResource( + server: McpServer, + manager: KnowledgeGraphManager, +) { + server.registerResource( + "knowledge-graph", + RESOURCE_URI, + { + title: "Knowledge Graph", + description: "The full knowledge graph with all entities and relations", + mimeType: "application/json", + }, + async (uri) => { + const graph = await manager.readGraph(); + return { + contents: [ + { + uri: uri.href, + mimeType: "application/json", + text: JSON.stringify(graph, null, 2), + }, + ], + }; + }, + ); +} + +// Enable clients to subscribe to the knowledge-graph resource and receive +// notifications/resources/updated when mutation tools change the graph. +export function registerKnowledgeGraphSubscriptions(server: McpServer) { + server.server.registerCapabilities({ resources: { subscribe: true } }); + server.server.setRequestHandler(SubscribeRequestSchema, async (request) => { + resourceSubscribers.add(request.params.uri); + return {}; + }); + server.server.setRequestHandler(UnsubscribeRequestSchema, async (request) => { + resourceSubscribers.delete(request.params.uri); + return {}; + }); +} + async function main() { - // Initialize memory file path with backward compatibility MEMORY_FILE_PATH = await ensureMemoryFilePath(); - - // Initialize knowledge graph manager with the memory file path knowledgeGraphManager = new KnowledgeGraphManager(MEMORY_FILE_PATH); + registerKnowledgeGraphResource(server, knowledgeGraphManager); + registerKnowledgeGraphSubscriptions(server); const transport = new StdioServerTransport(); await server.connect(transport);