Skip to content

Commit ae737a3

Browse files
committed
feat: fixed mcp issue
1 parent d5364d4 commit ae737a3

86 files changed

Lines changed: 1075 additions & 1521 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

packages/openmemory-js/src/ai/mcp.ts

Lines changed: 42 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import type { sector_type, mem_row, rpc_err_code } from "../core/types";
1717
import { update_user_summary } from "../memory/user_summary";
1818
import { insert_fact } from "../temporal_graph/store";
1919
import { query_facts_at_time } from "../temporal_graph/query";
20+
import { ToolRegistry } from "./mcp_tools";
2021

2122
const sec_enum = z.enum([
2223
"episodic",
@@ -87,7 +88,9 @@ export const create_mcp_srv = () => {
8788
{ capabilities: { tools: {}, resources: {}, logging: {} } },
8889
);
8990

90-
srv.tool(
91+
const registry = new ToolRegistry();
92+
93+
registry.tool(
9194
"openmemory_query",
9295
"Query OpenMemory for contextual memories (HSG) and/or temporal facts",
9396
{
@@ -164,13 +167,15 @@ export const create_mcp_srv = () => {
164167
const results: any = { type, query };
165168
const at_date = at ? new Date(at) : new Date();
166169

167-
// Query HSG if contextual or unified
170+
168171
if (type === "contextual" || type === "unified") {
169172
const flt =
170173
sector || min_salience !== undefined || u
171174
? {
172175
...(sector ? { sectors: [sector as sector_type] } : {}),
173-
...(min_salience !== undefined ? { minSalience: min_salience } : {}),
176+
...(min_salience !== undefined
177+
? { minSalience: min_salience }
178+
: {}),
174179
...(u ? { user_id: u } : {}),
175180
}
176181
: undefined;
@@ -189,14 +194,14 @@ export const create_mcp_srv = () => {
189194
}));
190195
}
191196

192-
// Query temporal facts if factual or unified
197+
193198
if (type === "factual" || type === "unified") {
194199
const facts = await query_facts_at_time(
195200
fact_pattern?.subject,
196201
fact_pattern?.predicate,
197202
fact_pattern?.object,
198203
at_date,
199-
0.0, // min_confidence
204+
0.0,
200205
);
201206

202207
results.factual = facts.map((f: any) => ({
@@ -212,7 +217,7 @@ export const create_mcp_srv = () => {
212217
}));
213218
}
214219

215-
// Format text summary
220+
216221
let summ = "";
217222
if (type === "contextual") {
218223
summ = results.contextual.length
@@ -230,7 +235,7 @@ export const create_mcp_srv = () => {
230235
.join("\n\n");
231236
}
232237
} else {
233-
// unified
238+
234239
const ctx_count = results.contextual?.length || 0;
235240
const fact_count = results.factual?.length || 0;
236241
summ = `Found ${ctx_count} contextual memories and ${fact_count} temporal facts.\n\n`;
@@ -267,7 +272,7 @@ export const create_mcp_srv = () => {
267272
},
268273
);
269274

270-
srv.tool(
275+
registry.tool(
271276
"openmemory_store",
272277
"Persist new content into OpenMemory (HSG contextual memory and/or temporal facts)",
273278
{
@@ -283,7 +288,10 @@ export const create_mcp_srv = () => {
283288
.array(
284289
z.object({
285290
subject: z.string().min(1).describe("Fact subject (entity)"),
286-
predicate: z.string().min(1).describe("Fact predicate (relationship)"),
291+
predicate: z
292+
.string()
293+
.min(1)
294+
.describe("Fact predicate (relationship)"),
287295
object: z.string().min(1).describe("Fact object (value)"),
288296
confidence: z
289297
.number()
@@ -294,14 +302,19 @@ export const create_mcp_srv = () => {
294302
valid_from: z
295303
.string()
296304
.optional()
297-
.describe("ISO date string for fact validity start (default: now)"),
305+
.describe(
306+
"ISO date string for fact validity start (default: now)",
307+
),
298308
}),
299309
)
300310
.optional()
301311
.describe(
302312
"Array of facts to store in temporal graph. Required when type is 'factual' or 'both'",
303313
),
304-
tags: z.array(z.string()).optional().describe("Optional tag list (for HSG storage)"),
314+
tags: z
315+
.array(z.string())
316+
.optional()
317+
.describe("Optional tag list (for HSG storage)"),
305318
metadata: z
306319
.record(z.any())
307320
.optional()
@@ -319,14 +332,17 @@ export const create_mcp_srv = () => {
319332
const u = uid(user_id);
320333
const results: any = { type };
321334

322-
// Validate facts are provided when needed
323-
if ((type === "factual" || type === "both") && (!facts || facts.length === 0)) {
335+
336+
if (
337+
(type === "factual" || type === "both") &&
338+
(!facts || facts.length === 0)
339+
) {
324340
throw new Error(
325341
`Facts array is required when type is '${type}'. Please provide at least one fact.`,
326342
);
327343
}
328344

329-
// Store in HSG if contextual or both
345+
330346
if (type === "contextual" || type === "both") {
331347
const res = await add_hsg_memory(
332348
content,
@@ -339,23 +355,23 @@ export const create_mcp_srv = () => {
339355
primary_sector: res.primary_sector,
340356
sectors: res.sectors,
341357
};
342-
358+
343359
if (u) {
344360
update_user_summary(u).catch((err) =>
345361
console.error("[MCP] user summary update failed:", err),
346362
);
347363
}
348364
}
349365

350-
// Store in temporal graph if factual or both
366+
351367
if ((type === "factual" || type === "both") && facts) {
352368
const temporal_results = [];
353369
for (const fact of facts) {
354370
const valid_from = fact.valid_from
355371
? new Date(fact.valid_from)
356372
: new Date();
357373
const confidence = fact.confidence ?? 1.0;
358-
374+
359375
const fact_id = await insert_fact(
360376
fact.subject,
361377
fact.predicate,
@@ -364,7 +380,7 @@ export const create_mcp_srv = () => {
364380
confidence,
365381
metadata,
366382
);
367-
383+
368384
temporal_results.push({
369385
id: fact_id,
370386
subject: fact.subject,
@@ -377,7 +393,7 @@ export const create_mcp_srv = () => {
377393
results.temporal = temporal_results;
378394
}
379395

380-
// Format response
396+
381397
let txt = "";
382398
if (type === "contextual") {
383399
txt = `Stored memory ${results.hsg.id} (primary=${results.hsg.primary_sector}) across sectors: ${results.hsg.sectors.join(", ")}${u ? ` [user=${u}]` : ""}`;
@@ -403,7 +419,7 @@ export const create_mcp_srv = () => {
403419
},
404420
);
405421

406-
srv.tool(
422+
registry.tool(
407423
"openmemory_reinforce",
408424
"Boost salience for an existing memory",
409425
{
@@ -428,7 +444,7 @@ export const create_mcp_srv = () => {
428444
},
429445
);
430446

431-
srv.tool(
447+
registry.tool(
432448
"openmemory_list",
433449
"List recent memories for quick inspection",
434450
{
@@ -439,9 +455,7 @@ export const create_mcp_srv = () => {
439455
.max(50)
440456
.default(10)
441457
.describe("Number of memories to return"),
442-
sector: sec_enum
443-
.optional()
444-
.describe("Optionally limit to a sector"),
458+
sector: sec_enum.optional().describe("Optionally limit to a sector"),
445459
user_id: z
446460
.string()
447461
.trim()
@@ -483,7 +497,7 @@ export const create_mcp_srv = () => {
483497
},
484498
);
485499

486-
srv.tool(
500+
registry.tool(
487501
"openmemory_get",
488502
"Fetch a single memory by identifier",
489503
{
@@ -543,6 +557,7 @@ export const create_mcp_srv = () => {
543557
};
544558
},
545559
);
560+
registry.apply(srv);
546561

547562
srv.resource(
548563
"openmemory-config",
@@ -582,7 +597,7 @@ export const create_mcp_srv = () => {
582597
);
583598

584599
srv.server.oninitialized = () => {
585-
// Use stderr for debug output, not stdout
600+
586601
console.error(
587602
"[MCP] initialization completed with client:",
588603
srv.server.getClientVersion(),
@@ -683,7 +698,7 @@ export const start_mcp_stdio = async () => {
683698
const srv = create_mcp_srv();
684699
const trans = new StdioServerTransport();
685700
await srv.connect(trans);
686-
// console.error("[MCP] STDIO transport connected"); // Use stderr for debug output, not stdout
701+
687702
};
688703

689704
if (typeof require !== "undefined" && require.main === module) {
Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
import { z } from "zod";
2+
import { zodToJsonSchema } from "zod-to-json-schema";
3+
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
4+
import { CallToolRequestSchema, ListToolsRequestSchema, McpError, ErrorCode } from "@modelcontextprotocol/sdk/types.js";
5+
6+
7+
type ToolCallback = (args: any, extra?: any) => Promise<any> | any;
8+
9+
interface ToolDef {
10+
name: string;
11+
description: string;
12+
inputSchema: z.ZodType<any>;
13+
callback: ToolCallback;
14+
}
15+
16+
export class ToolRegistry {
17+
private tools: Map<string, ToolDef> = new Map();
18+
19+
tool(name: string, description: string, inputSchema: any, callback: ToolCallback) {
20+
this.tools.set(name, {
21+
name,
22+
description,
23+
inputSchema: z.object(inputSchema),
24+
callback
25+
});
26+
}
27+
28+
apply(server: McpServer) {
29+
30+
31+
const srv = server.server;
32+
33+
srv.setRequestHandler(ListToolsRequestSchema, async () => {
34+
return {
35+
tools: Array.from(this.tools.values()).map(t => {
36+
const jsonSchema = zodToJsonSchema(t.inputSchema, {
37+
target: "jsonSchema2019-09"
38+
}) as any;
39+
40+
41+
if (jsonSchema && typeof jsonSchema === 'object') {
42+
43+
44+
45+
46+
delete jsonSchema.$schema;
47+
}
48+
49+
return {
50+
name: t.name,
51+
description: t.description,
52+
inputSchema: jsonSchema
53+
};
54+
})
55+
};
56+
});
57+
58+
srv.setRequestHandler(CallToolRequestSchema, async (req: any, extra: any) => {
59+
const name = req.params.name;
60+
const tool = this.tools.get(name);
61+
if (!tool) {
62+
throw new McpError(ErrorCode.MethodNotFound, `Tool not found: ${name}`);
63+
}
64+
65+
66+
const args = req.params.arguments || {};
67+
const parse = await tool.inputSchema.safeParseAsync(args);
68+
if (!parse.success) {
69+
throw new McpError(ErrorCode.InvalidParams, `Invalid arguments: ${parse.error.message}`);
70+
}
71+
72+
try {
73+
return await tool.callback(parse.data, extra);
74+
} catch (err: any) {
75+
return {
76+
content: [{ type: "text", text: `Error: ${err.message}` }],
77+
isError: true
78+
};
79+
}
80+
});
81+
}
82+
}

0 commit comments

Comments
 (0)