Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
173 changes: 138 additions & 35 deletions packages/openmemory-js/src/ai/mcp.ts
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,55 @@ const sec_enum = z.enum([
const trunc = (val: string, max = 200) =>
val.length <= max ? val : `${val.slice(0, max).trimEnd()}...`;

/** Cap HSG / unbounded graph work so a hung embed or huge scan cannot drop Streamable HTTP. */
const query_timeout_ms = () => {
const n = Number(process.env.OM_MCP_QUERY_TIMEOUT_MS);
return Number.isFinite(n) && n > 0 ? n : 12_000;
};

const with_timeout = async <T>(
p: Promise<T>,
ms: number,
label: string,
): Promise<T> => {
let timer: ReturnType<typeof setTimeout> | undefined;
try {
return await Promise.race([
p,
new Promise<T>((_, reject) => {
timer = setTimeout(
() =>
reject(
new Error(
`${label} timed out after ${ms}ms; try a tighter query or type=factual with fact_pattern`,
),
),
ms,
);
if (typeof (timer as NodeJS.Timeout).unref === "function") {
(timer as NodeJS.Timeout).unref();
}
}),
]);
} finally {
if (timer) clearTimeout(timer);
// Timed-out work may still reject later; swallow so the HTTP
// session is not taken down by an unhandled rejection.
p.catch(() => {});
Comment on lines +64 to +68

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Cancel HSG work when its timeout wins

When an embedding or vector-search promise never settles, Promise.race returns the timeout error but leaves p running; the added .catch only observes a later rejection and does not cancel it. Because the underlying hsg_query therefore never reaches its finally block, its active_queries slot remains occupied, and after env.max_active timed-out calls every subsequent contextual query is rejected by the rate limiter until the hung provider recovers or the process restarts.

Useful? React with 👍 / 👎.

}
};

const tool_error = (msg: string) => ({
content: [{ type: "text" as const, text: msg }],
isError: true as const,
});

const has_fact_pattern = (fp?: {
subject?: string;
predicate?: string;
object?: string;
}) => Boolean(fp?.subject || fp?.predicate || fp?.object);

const build_mem_snap = (row: mem_row) => ({
id: row.id,
primary_sector: row.primary_sector,
Expand Down Expand Up @@ -140,7 +189,7 @@ export const create_mcp_srv = (tenant?: string) => {
.optional()
.default("contextual")
.describe(
"Query type: 'contextual' for HSG semantic search (default), 'factual' for temporal fact queries, 'unified' for both",
"Query type: 'contextual' for HSG semantic search (default; fails soft on embed/search errors), 'factual' for temporal fact queries (pass fact_pattern or results are capped at k), 'unified' for both",
),
fact_pattern: z
.object({
Expand Down Expand Up @@ -179,7 +228,9 @@ export const create_mcp_srv = (tenant?: string) => {
.min(1)
.max(32)
.default(8)
.describe("Maximum results to return (for HSG queries)"),
.describe(
"Maximum results to return. Applies to HSG queries and to factual queries that omit fact_pattern",
),
sector: sec_enum
.optional()
.describe(
Expand Down Expand Up @@ -221,6 +272,8 @@ export const create_mcp_srv = (tenant?: string) => {
const proj = uid(project_id);
const results: any = { type, query };
const at_date = at ? new Date(at) : new Date();
const top_k = k ?? 8;
const warnings: string[] = [];

if (type === "contextual" || type === "unified") {
const flt =
Expand All @@ -237,44 +290,87 @@ export const create_mcp_srv = (tenant?: string) => {
}
: undefined;

const matches = await hsg_query(query, k ?? 8, flt);
results.contextual = matches.map((m: any) => ({
source: "hsg",
id: m.id,
score: Number(m.score.toFixed(4)),
primary_sector: m.primary_sector,
sectors: m.sectors,
salience: Number(m.salience.toFixed(4)),
last_seen_at: m.last_seen_at,
path: m.path,
content: m.content,
}));
try {
const matches = await with_timeout(
hsg_query(query, top_k, flt),
query_timeout_ms(),
"openmemory_query HSG search",
);
results.contextual = matches.map((m: any) => ({
source: "hsg",
id: m.id,
score: Number(m.score.toFixed(4)),
primary_sector: m.primary_sector,
sectors: m.sectors,
salience: Number(m.salience.toFixed(4)),
last_seen_at: m.last_seen_at,
path: m.path,
content: m.content,
}));
} catch (err: any) {
const msg = err?.message || "HSG contextual search failed";
// Default type is contextual: never let a hung/failed
// embed drop the Streamable HTTP POST. Unified can still
// return facts; pure contextual surfaces a tool error.
if (type === "contextual") {
return tool_error(`openmemory_query failed: ${msg}`);
}
warnings.push(`contextual search failed: ${msg}`);
results.contextual = [];
results.contextual_error = msg;
}
}

if (type === "factual" || type === "unified") {
const facts = await query_facts_at_time({
user_id: u ?? "anonymous",
project_id: proj,
subject: fact_pattern?.subject,
predicate: fact_pattern?.predicate,
object: fact_pattern?.object,
at: at_date,
min_confidence: 0.0,
});

results.factual = facts.map((f: any) => ({
source: "temporal",
id: f.id,
subject: f.subject,
predicate: f.predicate,
object: f.object,
valid_from: f.valid_from,
valid_to: f.valid_to,
confidence: Number(f.confidence.toFixed(4)),
content: `${f.subject} ${f.predicate} ${f.object}`,
}));
const patterned = has_fact_pattern(fact_pattern);
try {
const facts = await with_timeout(
query_facts_at_time({
user_id: u ?? "anonymous",
project_id: proj,
subject: fact_pattern?.subject,
predicate: fact_pattern?.predicate,
object: fact_pattern?.object,
at: at_date,
min_confidence: 0.0,
// Unpatterned dumps can exceed 1MB and stall the
// MCP session; always cap, tighter without a pattern.
limit: patterned ? Math.max(top_k, 32) : top_k,
}),
query_timeout_ms(),
"openmemory_query factual search",
);

results.factual = facts.map((f: any) => ({
source: "temporal",
id: f.id,
subject: f.subject,
predicate: f.predicate,
object: f.object,
valid_from: f.valid_from,
valid_to: f.valid_to,
confidence: Number(f.confidence.toFixed(4)),
content: `${f.subject} ${f.predicate} ${f.object}`,
}));
if (!patterned) {
results.factual_capped = top_k;
warnings.push(
`factual results capped at k=${top_k}; pass fact_pattern to narrow the graph scan`,
);
}
} catch (err: any) {
const msg = err?.message || "temporal fact query failed";
if (type === "factual") {
return tool_error(`openmemory_query failed: ${msg}`);
}
warnings.push(`factual search failed: ${msg}`);
results.factual = [];
results.factual_error = msg;
}
}

if (warnings.length) results.warnings = warnings;

let summ = "";
if (type === "contextual") {
summ = results.contextual.length
Expand Down Expand Up @@ -316,6 +412,13 @@ export const create_mcp_srv = (tenant?: string) => {
}
}

if (warnings.length) {
summ =
warnings.map((w) => `Warning: ${w}`).join("\n") +
"\n\n" +
summ;
}

return {
content: [
{ type: "text", text: summ },
Expand Down
148 changes: 87 additions & 61 deletions packages/openmemory-js/src/memory/hsg.ts
Original file line number Diff line number Diff line change
Expand Up @@ -835,6 +835,88 @@ const get_sal = async (id: string, def_sal: number): Promise<number> => {
sal_cache.set(id, { s, t: Date.now() });
return s;
};
/**
* Post-hit reinforcement: salience/feedback bump, waypoint + associative
* reinforcement, and regeneration/re-embed. These are maintenance work, not
* part of producing the ranked hits.
*
* They were previously awaited on the request path of `hsg_query`, which meant
* a handful of extra DB round-trips (and, when regeneration is on, an embed
* network call per hit) stacked on top of the 5-sector embed + vector search.
* That tail did not fit inside the MCP fail-soft query budget, so a perfectly
* good search would surface as "HSG search timed out after 8000ms".
*
* `hsg_query` now returns as soon as ranking + caching is done and runs this
* fire-and-forget. Each write is independently guarded so a failure in one hit
* cannot surface on the request path.
*/
async function reinforce_query_hits(
top: hsg_q_result[],
tids: string[],
): Promise<void> {
for (const r of top) {
const cur_fb = (await q.get_mem.get(r.id))?.feedback_score || 0;
const new_fb = cur_fb * 0.9 + r.score * 0.1;
await q.upd_feedback.run(r.id, new_fb);
}

for (let i = 0; i < tids.length; i++) {
for (let j = i + 1; j < tids.length; j++) {
const [a, b] = [tids[i], tids[j]].sort();
coact_buf.push([a, b]);
}
}

for (const r of top) {
const rsal = await applyRetrievalTraceReinforcementToMemory(
r.id,
r.salience,
);
await q.upd_seen.run(r.id, Date.now(), rsal, Date.now());
if (r.path.length > 1) {
await reinforce_waypoints(r.path);
const wps = await q.get_waypoints_by_src.all(r.id);
const lns = wps.map((wp: any) => ({
target_id: wp.dst_id,
weight: wp.weight,
}));
const pru =
await propagateAssociativeReinforcementToLinkedNodes(
r.id,
rsal,
lns,
);
for (const u of pru) {
const linked_mem = await q.get_mem.get(u.node_id);
if (linked_mem) {
const time_diff =
(Date.now() - linked_mem.last_seen_at) / 86400000;
const decay_fact = Math.exp(-0.02 * time_diff);
const ctx_boost =
hybrid_params.gamma *
(rsal - linked_mem.salience) *
decay_fact;
const new_sal = Math.max(
0,
Math.min(1, linked_mem.salience + ctx_boost),
);
await q.upd_seen.run(
u.node_id,
Date.now(),
new_sal,
Date.now(),
);
}
}
}
}

for (const r of top) {
await on_query_hit(r.id, r.primary_sector, (text) =>
embedForSector(text, r.primary_sector),
);
}
}
export async function hsg_query(
qt: string,
k = 10,
Expand Down Expand Up @@ -1043,67 +1125,11 @@ export async function hsg_query(
const top = top_cands.slice(0, k);
const tids = top.map((r) => r.id);

for (const r of top) {
const cur_fb = (await q.get_mem.get(r.id))?.feedback_score || 0;
const new_fb = cur_fb * 0.9 + r.score * 0.1;
await q.upd_feedback.run(r.id, new_fb);
}

for (let i = 0; i < tids.length; i++) {
for (let j = i + 1; j < tids.length; j++) {
const [a, b] = [tids[i], tids[j]].sort();
coact_buf.push([a, b]);
}
}
for (const r of top) {
const rsal = await applyRetrievalTraceReinforcementToMemory(
r.id,
r.salience,
);
await q.upd_seen.run(r.id, Date.now(), rsal, Date.now());
if (r.path.length > 1) {
await reinforce_waypoints(r.path);
const wps = await q.get_waypoints_by_src.all(r.id);
const lns = wps.map((wp: any) => ({
target_id: wp.dst_id,
weight: wp.weight,
}));
const pru =
await propagateAssociativeReinforcementToLinkedNodes(
r.id,
rsal,
lns,
);
for (const u of pru) {
const linked_mem = await q.get_mem.get(u.node_id);
if (linked_mem) {
const time_diff =
(Date.now() - linked_mem.last_seen_at) / 86400000;
const decay_fact = Math.exp(-0.02 * time_diff);
const ctx_boost =
hybrid_params.gamma *
(rsal - linked_mem.salience) *
decay_fact;
const new_sal = Math.max(
0,
Math.min(1, linked_mem.salience + ctx_boost),
);
await q.upd_seen.run(
u.node_id,
Date.now(),
new_sal,
Date.now(),
);
}
}
}
}

for (const r of top) {
on_query_hit(r.id, r.primary_sector, (text) =>
embedForSector(text, r.primary_sector),
).catch(() => {});
}
// Reinforcement is off the request path: return the hits as soon as
// ranking + caching is done. The tail work (feedback/salience bumps,
// waypoint + associative propagation, regeneration re-embeds) runs in
// the background and must never be able to blow the MCP query budget.
reinforce_query_hits(top, tids).catch(() => {});

cache.set(h, { r: top, t: Date.now() });
return top;
Expand Down
Loading
Loading