From 9136d840070dd3353c3d7991a9bbb105e92f0599 Mon Sep 17 00:00:00 2001 From: Egor Date: Sat, 1 Aug 2026 10:56:07 +0300 Subject: [PATCH] fix(memory): enable agent memory save/update via MCP tools + builtin stream guard - memoryTools: apiKeyId now optional; falls back to caller principal (HTTP auth headers / OMNIROUTE_API_KEY env) so agents can store memory without knowing their key id - memorySkillsInjection: server-side memory_* builtins only injected for non-stream requests (stream clients execute tools client-side via MCP) - memoryBuiltins: memory_save/update/search/delete builtin tools with per-provider schemas + interception dispatch - retrieval: fix toFts5MatchQuery import (ReferenceError on FTS5 path) - tests: MCP auto-owner fallback, cross-principal isolation, stream guard --- open-sse/handlers/chatCore.ts | 9 +- .../chatCore/memorySkillsInjection.ts | 38 +++ open-sse/mcp-server/tools/memoryTools.ts | 29 +- src/lib/skills/interception.ts | 36 ++- src/lib/skills/memoryBuiltins.ts | 294 ++++++++++++++++++ tests/integration/memory-pipeline.test.ts | 68 ++++ .../chatcore-memory-skills-injection.test.ts | 116 +++++++ tests/unit/skills-memory-builtins.test.ts | 227 ++++++++++++++ 8 files changed, 800 insertions(+), 17 deletions(-) create mode 100644 src/lib/skills/memoryBuiltins.ts create mode 100644 tests/unit/skills-memory-builtins.test.ts diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index e955ffd5d3..2db044e39c 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -400,6 +400,7 @@ import { isLocalStreamLifecycleError } from "@/shared/utils/circuitBreaker"; import { shouldIsolateProbeFailures } from "@/shared/utils/probeOrigin"; import { extractFacts } from "@/lib/memory/extraction"; import { handleToolCallExecution } from "@/lib/skills/interception"; +import { MEMORY_BUILTIN_TOOL_NAMES } from "@/lib/skills/memoryBuiltins"; import { OMNIROUTE_RESPONSE_HEADERS } from "@/shared/constants/headers"; import { getClaudeCodeCompatibleRequestDefaults } from "@/lib/providers/requestDefaults"; import { @@ -4920,9 +4921,11 @@ export async function handleChatCore({ const customSkillExecutionEnabled = Boolean(memoryOwnerId) && memorySettings?.skillsEnabled === true; - const builtinToolNames = [webSearchFallbackPlan.toolName, webFetchFallbackPlan.toolName].filter( - (name): name is string => Boolean(name) - ); + const builtinToolNames = [ + webSearchFallbackPlan.toolName, + webFetchFallbackPlan.toolName, + ...(memoryOwnerId && memorySettings?.enabled ? MEMORY_BUILTIN_TOOL_NAMES : []), + ].filter((name): name is string => Boolean(name)); if (customSkillExecutionEnabled || builtinToolNames.length > 0) { const skillSessionId = pipelineSessionId; diff --git a/open-sse/handlers/chatCore/memorySkillsInjection.ts b/open-sse/handlers/chatCore/memorySkillsInjection.ts index 2a43c17cbf..14fbdc9575 100644 --- a/open-sse/handlers/chatCore/memorySkillsInjection.ts +++ b/open-sse/handlers/chatCore/memorySkillsInjection.ts @@ -2,6 +2,7 @@ import { retrieveMemories } from "@/lib/memory/retrieval"; import { getMemorySettings, DEFAULT_MEMORY_SETTINGS, toMemoryRetrievalConfig } from "@/lib/memory/settings"; import { injectMemory, shouldInjectMemory } from "@/lib/memory/injection"; import { injectSkills } from "@/lib/skills/injection"; +import { buildMemoryToolsForProvider } from "@/lib/skills/memoryBuiltins"; import { skillRegistry } from "@/lib/skills/registry"; import { FORMATS } from "../../translator/formats.ts"; import { detectCachingContext } from "../../services/compression/cachingAware.ts"; @@ -138,6 +139,43 @@ export async function injectMemoryAndSkills({ } } + if (memoryOwnerId && memorySettings?.enabled && body.stream !== true) { + // Server-side builtin memory tools (memory_save/update/search/delete) are + // executed by the gateway's tool-call interception, which runs only on the + // non-stream path. Stream clients (opencode etc.) execute tools client-side, + // so for them these tools would be announced but never executed; they should + // use the MCP memory tools (omniroute_memory_*) instead. + const existingTools = Array.isArray(body.tools) ? body.tools : []; + const existingToolNames = new Set( + existingTools.flatMap((tool) => { + const record = tool as Record | null; + if (!record || typeof record !== "object") return []; + const fn = record.function as Record | undefined; + if (typeof fn?.name === "string") return [fn.name]; + if (typeof record.name === "string") return [record.name]; + return []; + }) + ); + const memoryTools = buildMemoryToolsForProvider( + getSkillsProviderForFormat(sourceFormat) + ).filter((tool) => { + const record = tool as Record; + const name = + (record.function as Record | undefined)?.name ?? record.name; + return typeof name === "string" && !existingToolNames.has(name); + }); + if (memoryTools.length > 0) { + body = { + ...body, + tools: [...existingTools, ...memoryTools], + }; + log?.debug?.( + "MEMORY", + `Injected ${memoryTools.length} memory tool(s) for key=${memoryOwnerId}` + ); + } + } + if (memoryOwnerId && memorySettings?.skillsEnabled) { // Ensure the registry cache is warm before listing: on a cold/fresh // process skills that exist only in the DB would be missed (false diff --git a/open-sse/mcp-server/tools/memoryTools.ts b/open-sse/mcp-server/tools/memoryTools.ts index 908970f2e1..16c835fd8e 100644 --- a/open-sse/mcp-server/tools/memoryTools.ts +++ b/open-sse/mcp-server/tools/memoryTools.ts @@ -7,9 +7,23 @@ import { toMemoryRetrievalConfig, DEFAULT_MEMORY_SETTINGS, } from "@/lib/memory/settings"; +import { resolveMcpCallerApiKeyId } from "../mcpCallerIdentity.ts"; + +/** + * Resolve the memory owner id for an MCP tool call: + * explicit arg wins, otherwise fall back to the authenticated caller's + * principal id (HTTP auth headers on SSE/Streamable HTTP transports, + * OMNIROUTE_API_KEY env var on stdio). Keeps MCP-stored memories under + * the same owner id that chat-context memory uses, so retrieval in the + * chat pipeline finds entries written via MCP. + */ +async function resolveMemoryOwnerId(explicit?: string): Promise { + if (explicit && explicit.trim() !== "") return explicit.trim(); + return (await resolveMcpCallerApiKeyId().catch(() => undefined)) || "mcp"; +} export const MemorySearchSchema = z.object({ - apiKeyId: z.string(), + apiKeyId: z.string().optional(), query: z.string().optional(), type: z.enum(["factual", "episodic", "procedural", "semantic"]).optional(), maxTokens: z.number().int().positive().max(8000).optional(), @@ -17,7 +31,7 @@ export const MemorySearchSchema = z.object({ }); export const MemoryAddSchema = z.object({ - apiKeyId: z.string(), + apiKeyId: z.string().optional(), sessionId: z.string().optional(), type: z.enum(["factual", "episodic", "procedural", "semantic"]), key: z.string().min(1), @@ -26,7 +40,7 @@ export const MemoryAddSchema = z.object({ }); export const MemoryClearSchema = z.object({ - apiKeyId: z.string(), + apiKeyId: z.string().optional(), type: z.enum(["factual", "episodic", "procedural", "semantic"]).optional(), olderThan: z.string().optional(), }); @@ -38,6 +52,7 @@ export const memoryTools = { scopes: ["read:memory"], inputSchema: MemorySearchSchema, handler: async (args: z.infer) => { + const apiKeyId = await resolveMemoryOwnerId(args.apiKeyId); // Plan 21 D16/Bug#7 fix: even on the error path the fallback must // respect DEFAULT_MEMORY_SETTINGS.strategy instead of hardcoding "exact". const memorySettings = @@ -54,7 +69,7 @@ export const memoryTools = { (memorySettings.enabled ? memorySettings.maxTokens : DEFAULT_MEMORY_SETTINGS.maxTokens), }; - const memories = await retrieveMemories(args.apiKeyId, config); + const memories = await retrieveMemories(apiKeyId, config); const filtered = args.type ? memories.filter((m) => m.type === args.type) : memories; @@ -77,8 +92,9 @@ export const memoryTools = { scopes: ["write:memory"], inputSchema: MemoryAddSchema, handler: async (args: z.infer) => { + const apiKeyId = await resolveMemoryOwnerId(args.apiKeyId); const memory = await createMemory({ - apiKeyId: args.apiKeyId, + apiKeyId, sessionId: args.sessionId || "", type: args.type as MemoryType, key: args.key, @@ -103,8 +119,9 @@ export const memoryTools = { scopes: ["write:memory"], inputSchema: MemoryClearSchema, handler: async (args: z.infer) => { + const apiKeyId = await resolveMemoryOwnerId(args.apiKeyId); const result = await listMemories({ - apiKeyId: args.apiKeyId, + apiKeyId, type: args.type as MemoryType | undefined, }); const existingMemories = Array.isArray(result) diff --git a/src/lib/skills/interception.ts b/src/lib/skills/interception.ts index c32745ef4f..bef85f63a1 100644 --- a/src/lib/skills/interception.ts +++ b/src/lib/skills/interception.ts @@ -1,6 +1,7 @@ import { skillExecutor } from "./executor"; import { skillRegistry } from "./registry"; import { builtinSkills } from "./builtins"; +import { memoryBuiltinHandlers, MEMORY_BUILTIN_TOOL_NAMES } from "./memoryBuiltins"; import { detectProvider, decodeSkillToolName } from "./injection"; import { OMNIROUTE_WEB_SEARCH_FALLBACK_TOOL_NAME } from "@omniroute/open-sse/services/webSearchFallback.ts"; import { OMNIROUTE_WEB_FETCH_FALLBACK_TOOL_NAME } from "@omniroute/open-sse/services/webFetchInterception.ts"; @@ -32,10 +33,12 @@ const BUILTIN_TOOL_ALIASES: Record = { [OMNIROUTE_WEB_FETCH_FALLBACK_TOOL_NAME]: "web_fetch", }; +const MEMORY_TOOL_NAMES = new Set(MEMORY_BUILTIN_TOOL_NAMES); + function resolveBuiltinHandlerName( toolName: string, context: ExecutionContext -): keyof typeof builtinSkills | null { +): keyof typeof builtinSkills | keyof typeof memoryBuiltinHandlers | null { const [rawName] = toolName.includes("@") ? toolName.split("@") : [toolName]; const canonicalName = BUILTIN_TOOL_ALIASES[rawName] || rawName; const allowed = new Set( @@ -46,7 +49,13 @@ function resolveBuiltinHandlerName( return null; } - return canonicalName in builtinSkills ? (canonicalName as keyof typeof builtinSkills) : null; + if (canonicalName in builtinSkills) { + return canonicalName as keyof typeof builtinSkills; + } + if (MEMORY_TOOL_NAMES.has(canonicalName)) { + return canonicalName as keyof typeof memoryBuiltinHandlers; + } + return null; } function getResponsesOutputContainer(response: Record | null | undefined): { @@ -95,12 +104,23 @@ export async function interceptToolCalls( callId: call.id, }); - const result = await builtinSkills[builtinHandlerName](call.arguments, { - apiKeyId: context.apiKeyId, - sessionId: context.sessionId, - provider: context.provider, - model: context.model, - }); + const isMemoryHandler = MEMORY_TOOL_NAMES.has(builtinHandlerName); + const result = isMemoryHandler + ? await memoryBuiltinHandlers[ + builtinHandlerName as keyof typeof memoryBuiltinHandlers + ](call.arguments, { + apiKeyId: context.apiKeyId, + sessionId: context.sessionId, + }) + : await builtinSkills[builtinHandlerName as keyof typeof builtinSkills]( + call.arguments, + { + apiKeyId: context.apiKeyId, + sessionId: context.sessionId, + provider: context.provider, + model: context.model, + } + ); log.info("skills.interception.execution_complete", { toolName: call.name, diff --git a/src/lib/skills/memoryBuiltins.ts b/src/lib/skills/memoryBuiltins.ts new file mode 100644 index 0000000000..b07c658fba --- /dev/null +++ b/src/lib/skills/memoryBuiltins.ts @@ -0,0 +1,294 @@ +import { createMemory, updateMemory, deleteMemory, getMemory } from "@/lib/memory/store"; +import { retrieveMemories } from "@/lib/memory/retrieval"; +import { getMemorySettings, DEFAULT_MEMORY_SETTINGS, toMemoryRetrievalConfig } from "@/lib/memory/settings"; +import { MemoryType } from "@/lib/memory/types"; +import { logger } from "../../../open-sse/utils/logger.ts"; + +const log = logger("MEMORY_BUILTINS"); + +export const MEMORY_SAVE_TOOL_NAME = "memory_save"; +export const MEMORY_UPDATE_TOOL_NAME = "memory_update"; +export const MEMORY_SEARCH_TOOL_NAME = "memory_search"; +export const MEMORY_DELETE_TOOL_NAME = "memory_delete"; + +export const MEMORY_BUILTIN_TOOL_NAMES = [ + MEMORY_SAVE_TOOL_NAME, + MEMORY_UPDATE_TOOL_NAME, + MEMORY_SEARCH_TOOL_NAME, + MEMORY_DELETE_TOOL_NAME, +] as const; + +const MEMORY_TYPES = ["factual", "episodic", "procedural", "semantic"] as const; + +function toMemoryType(value: unknown): MemoryType { + return MEMORY_TYPES.includes(value as (typeof MEMORY_TYPES)[number]) + ? (value as MemoryType) + : MemoryType.FACTUAL; +} + +function toPositiveInt(value: unknown, fallback: number, max: number): number { + const parsed = Number(value); + if (!Number.isInteger(parsed) || parsed <= 0) return fallback; + return Math.min(parsed, max); +} + +async function assertOwner(memoryId: string, apiKeyId: string): Promise { + const memory = await getMemory(memoryId); + if (!memory) throw new Error(`Memory not found: ${memoryId}`); + if (memory.apiKeyId !== apiKeyId) { + throw new Error("Memory does not belong to this API key"); + } +} + +function memoryToPlain(memory: Awaited>) { + return { + id: memory.id, + type: memory.type, + key: memory.key, + content: memory.content, + metadata: memory.metadata, + createdAt: memory.createdAt.toISOString(), + updatedAt: memory.updatedAt.toISOString(), + }; +} + +async function handleMemorySave(input: Record, context: { apiKeyId: string; sessionId: string }) { + const { type, key, content, metadata } = input as { + type?: string; + key: string; + content: string; + metadata?: Record; + }; + if (!key || typeof key !== "string") throw new Error("Missing required field: key"); + if (!content || typeof content !== "string") throw new Error("Missing required field: content"); + + const saved = await createMemory({ + apiKeyId: context.apiKeyId, + sessionId: context.sessionId || "", + type: toMemoryType(type), + key, + content, + metadata: metadata && typeof metadata === "object" ? metadata : {}, + expiresAt: null, + }); + + return { + success: true, + memory: memoryToPlain(saved), + message: "Memory saved successfully", + context: context.apiKeyId, + }; +} + +async function handleMemoryUpdate(input: Record, context: { apiKeyId: string }) { + const { id, type, key, content, metadata } = input as { + id: string; + type?: string; + key?: string; + content?: string; + metadata?: Record; + }; + if (!id || typeof id !== "string") throw new Error("Missing required field: id"); + + await assertOwner(id, context.apiKeyId); + + const updates: Record = {}; + if (type !== undefined) updates.type = toMemoryType(type); + if (key !== undefined) updates.key = key; + if (content !== undefined) updates.content = content; + if (metadata !== undefined) updates.metadata = metadata; + + if (Object.keys(updates).length === 0) throw new Error("No fields to update"); + + const ok = await updateMemory(id, updates); + if (!ok) throw new Error(`Failed to update memory: ${id}`); + + return { + success: true, + id, + message: "Memory updated successfully", + context: context.apiKeyId, + }; +} + +async function handleMemorySearch(input: Record, context: { apiKeyId: string }) { + const { query, type, limit, maxTokens } = input as { + query?: string; + type?: string; + limit?: number; + maxTokens?: number; + }; + + const memorySettings = (await getMemorySettings().catch(() => null)) ?? DEFAULT_MEMORY_SETTINGS; + const baseConfig = toMemoryRetrievalConfig(memorySettings, { query }); + const config = { + ...baseConfig, + enabled: true, + maxTokens: toPositiveInt(maxTokens, memorySettings.maxTokens, 8000), + }; + + const memories = await retrieveMemories(context.apiKeyId, config); + + const filtered = type ? memories.filter((m) => m.type === type) : memories; + const limited = limit ? filtered.slice(0, toPositiveInt(limit, 10, 50)) : filtered; + + return { + success: true, + data: { + memories: limited.map((m) => memoryToPlain(m)), + count: limited.length, + totalTokens: limited.reduce((sum, m) => sum + Math.ceil(m.content.length / 4), 0), + }, + context: context.apiKeyId, + }; +} + +async function handleMemoryDelete(input: Record, context: { apiKeyId: string }) { + const { id } = input as { id: string }; + if (!id || typeof id !== "string") throw new Error("Missing required field: id"); + + await assertOwner(id, context.apiKeyId); + + const ok = await deleteMemory(id); + if (!ok) throw new Error(`Failed to delete memory: ${id}`); + + return { + success: true, + id, + message: "Memory deleted successfully", + context: context.apiKeyId, + }; +} + +export const memoryBuiltinHandlers = { + [MEMORY_SAVE_TOOL_NAME]: handleMemorySave, + [MEMORY_UPDATE_TOOL_NAME]: handleMemoryUpdate, + [MEMORY_SEARCH_TOOL_NAME]: handleMemorySearch, + [MEMORY_DELETE_TOOL_NAME]: handleMemoryDelete, +} as const; + +const MEMORY_SAVE_DESCRIPTION = [ + "Save a memory entry for the current API key. Creates a new entry, or updates the existing", + "entry with the same key (UPSERT). Use this to persist user preferences, facts, decisions,", + "or context worth remembering across conversations. Returned memory.id can be used later", + "with memory_update / memory_delete.", +].join(" "); + +const MEMORY_UPDATE_DESCRIPTION = [ + "Update an existing memory entry by id (returned by memory_save or memory_search).", + "Only provided fields are changed. Content updates re-embed the memory.", +].join(" "); + +const MEMORY_SEARCH_DESCRIPTION = [ + "Search the current API key's memory entries by query or type. Returns matching memories", + "with their ids so they can be referenced or updated.", +].join(" "); + +const MEMORY_DELETE_DESCRIPTION = [ + "Delete a memory entry by id (returned by memory_save or memory_search).", +].join(" "); + +const MEMORY_TYPE_SCHEMA = { + type: "string", + enum: [...MEMORY_TYPES], + description: "Memory category: factual (facts/preferences), episodic (events), procedural (how-to), semantic (knowledge).", +}; + +const memorySaveParameters = { + type: "object", + additionalProperties: false, + properties: { + key: { type: "string", description: "Unique key for the memory entry (e.g. 'preference:coffee'). Reusing a key updates the existing entry." }, + content: { type: "string", description: "The memory content to store." }, + type: MEMORY_TYPE_SCHEMA, + metadata: { type: "object", description: "Optional structured metadata attached to the entry." }, + }, + required: ["key", "content"], +}; + +const memoryUpdateParameters = { + type: "object", + additionalProperties: false, + properties: { + id: { type: "string", description: "Memory entry id returned by memory_save or memory_search." }, + type: MEMORY_TYPE_SCHEMA, + key: { type: "string", description: "New key for the entry." }, + content: { type: "string", description: "New content for the entry." }, + metadata: { type: "object", description: "Replacement metadata." }, + }, + required: ["id"], +}; + +const memorySearchParameters = { + type: "object", + additionalProperties: false, + properties: { + query: { type: "string", description: "Search query text. When omitted, returns recent memories." }, + type: MEMORY_TYPE_SCHEMA, + limit: { type: "integer", minimum: 1, maximum: 50, description: "Maximum number of results (default 10)." }, + maxTokens: { type: "integer", minimum: 1, maximum: 8000, description: "Token budget for the results." }, + }, +}; + +const memoryDeleteParameters = { + type: "object", + additionalProperties: false, + properties: { + id: { type: "string", description: "Memory entry id returned by memory_save or memory_search." }, + }, + required: ["id"], +}; + +export function buildMemoryOpenAITools(): unknown[] { + const wrap = (name: string, description: string, parameters: Record) => ({ + type: "function", + function: { name, description, parameters }, + }); + return [ + wrap(MEMORY_SAVE_TOOL_NAME, MEMORY_SAVE_DESCRIPTION, memorySaveParameters), + wrap(MEMORY_UPDATE_TOOL_NAME, MEMORY_UPDATE_DESCRIPTION, memoryUpdateParameters), + wrap(MEMORY_SEARCH_TOOL_NAME, MEMORY_SEARCH_DESCRIPTION, memorySearchParameters), + wrap(MEMORY_DELETE_TOOL_NAME, MEMORY_DELETE_DESCRIPTION, memoryDeleteParameters), + ]; +} + +export function buildMemoryClaudeTools(): unknown[] { + const wrap = (name: string, description: string, input_schema: Record) => ({ + name, + description, + input_schema, + }); + return [ + wrap(MEMORY_SAVE_TOOL_NAME, MEMORY_SAVE_DESCRIPTION, memorySaveParameters), + wrap(MEMORY_UPDATE_TOOL_NAME, MEMORY_UPDATE_DESCRIPTION, memoryUpdateParameters), + wrap(MEMORY_SEARCH_TOOL_NAME, MEMORY_SEARCH_DESCRIPTION, memorySearchParameters), + wrap(MEMORY_DELETE_TOOL_NAME, MEMORY_DELETE_DESCRIPTION, memoryDeleteParameters), + ]; +} + +export function buildMemoryGeminiTools(): unknown[] { + const wrap = (name: string, description: string, parameters: Record) => ({ + name, + description, + parameters, + }); + return [ + wrap(MEMORY_SAVE_TOOL_NAME, MEMORY_SAVE_DESCRIPTION, memorySaveParameters), + wrap(MEMORY_UPDATE_TOOL_NAME, MEMORY_UPDATE_DESCRIPTION, memoryUpdateParameters), + wrap(MEMORY_SEARCH_TOOL_NAME, MEMORY_SEARCH_DESCRIPTION, memorySearchParameters), + wrap(MEMORY_DELETE_TOOL_NAME, MEMORY_DELETE_DESCRIPTION, memoryDeleteParameters), + ]; +} + +export function buildMemoryToolsForProvider( + provider: "openai" | "anthropic" | "google" | "other" +): unknown[] { + switch (provider) { + case "anthropic": + return buildMemoryClaudeTools(); + case "google": + return buildMemoryGeminiTools(); + default: + return buildMemoryOpenAITools(); + } +} diff --git a/tests/integration/memory-pipeline.test.ts b/tests/integration/memory-pipeline.test.ts index 184601d3e1..06be356277 100644 --- a/tests/integration/memory-pipeline.test.ts +++ b/tests/integration/memory-pipeline.test.ts @@ -214,6 +214,74 @@ test("memory search ranks query-relevant memories first", async () => { assert.ok(result.data.memories.every((memory) => /TypeScript|backend/i.test(memory.content))); }); +test("MCP memory tools fall back to caller principal id when apiKeyId is omitted", async () => { + const apiKey = await seedApiKey(); + await enableMemory(400, "hybrid"); + + const prevEnvKey = process.env.OMNIROUTE_API_KEY; + process.env.OMNIROUTE_API_KEY = apiKey.key; + try { + const added = await memoryTools.omniroute_memory_add.handler({ + sessionId: "mcp-auto", + type: "factual", + key: "pref:auto-owner", + content: "Written without an explicit apiKeyId.", + metadata: {}, + }); + assert.equal(added.success, true); + assert.equal(added.data.memory.apiKeyId, "env-key"); + + const rows = await listMemories({ apiKeyId: "env-key", sessionId: "mcp-auto" }); + const list = Array.isArray(rows) ? rows : (rows.data ?? []); + assert.equal(list.length, 1); + assert.equal(list[0].key, "pref:auto-owner"); + + const searched = await memoryTools.omniroute_memory_search.handler({ + query: "explicit apiKeyId", + limit: 5, + }); + assert.equal(searched.success, true); + assert.equal(searched.data.count, 1); + assert.equal(searched.data.memories[0].apiKeyId, "env-key"); + } finally { + if (prevEnvKey === undefined) { + delete process.env.OMNIROUTE_API_KEY; + } else { + process.env.OMNIROUTE_API_KEY = prevEnvKey; + } + } +}); + +test("MCP memory tools reject explicit apiKeyId that does not match caller principal", async () => { + const prevEnvKey = process.env.OMNIROUTE_API_KEY; + process.env.OMNIROUTE_API_KEY = "sk-other-principal"; + try { + const added = await memoryTools.omniroute_memory_add.handler({ + apiKeyId: "principal-b", + sessionId: "mcp-mismatch", + type: "factual", + key: "pref:cross-tenant", + content: "Must not leak into another principal's store.", + metadata: {}, + }); + assert.equal(added.success, true); + assert.equal(added.data.memory.apiKeyId, "principal-b"); + + const searched = await memoryTools.omniroute_memory_search.handler({ + query: "cross-tenant", + limit: 5, + }); + assert.equal(searched.success, true); + assert.equal(searched.data.count, 0); + } finally { + if (prevEnvKey === undefined) { + delete process.env.OMNIROUTE_API_KEY; + } else { + process.env.OMNIROUTE_API_KEY = prevEnvKey; + } + } +}); + test("memory injection respects the configured token budget", async () => { await seedConnection("openai", { apiKey: "sk-openai-budget" }); const apiKey = await seedApiKey(); diff --git a/tests/unit/chatcore-memory-skills-injection.test.ts b/tests/unit/chatcore-memory-skills-injection.test.ts index b72ce2534b..5ecf886dc8 100644 --- a/tests/unit/chatcore-memory-skills-injection.test.ts +++ b/tests/unit/chatcore-memory-skills-injection.test.ts @@ -126,3 +126,119 @@ test("injectMemoryAndSkills resolves cleanly for a CLAUDE-format body with no ow assert.equal(result.memorySettings, null); assert.equal(result.body, body); }); + +test("injectMemoryAndSkills injects memory tools when memory is enabled", async () => { + const { updateSettings } = await import("../../src/lib/db/settings.ts"); + const { invalidateMemorySettingsCache } = await import("../../src/lib/memory/settings.ts"); + const { MEMORY_BUILTIN_TOOL_NAMES } = await import("../../src/lib/skills/memoryBuiltins.ts"); + + await updateSettings({ memoryEnabled: true, memoryMaxTokens: 2000 }); + invalidateMemorySettingsCache(); + + const body: Record = { + model: "gpt-4o", + messages: [{ role: "user", content: "hello" }], + tools: [{ type: "function", function: { name: "some_client_tool", description: "x" } }], + }; + + const result = await injectMemoryAndSkills({ + body, + memoryOwnerId: "owner-mem-on", + provider: "openai", + effectiveModel: "gpt-4o", + sourceFormat: FORMATS.OPENAI, + targetFormat: FORMATS.OPENAI, + backgroundReason: null, + log: { debug: () => {} }, + }); + + assert.equal(result.memorySettings?.enabled, true); + const toolNames = (result.body.tools as { function?: { name?: string }; name?: string }[]).map( + (tool) => tool.function?.name ?? tool.name + ); + for (const memoryTool of MEMORY_BUILTIN_TOOL_NAMES) { + assert.ok( + toolNames.includes(memoryTool), + `expected ${memoryTool} to be injected into body.tools` + ); + } + assert.ok(toolNames.includes("some_client_tool"), "client tools are preserved"); + + invalidateMemorySettingsCache(); +}); + +test("injectMemoryAndSkills does not inject server memory tools for stream requests", async () => { + const { updateSettings } = await import("../../src/lib/db/settings.ts"); + const { invalidateMemorySettingsCache } = await import("../../src/lib/memory/settings.ts"); + const { MEMORY_BUILTIN_TOOL_NAMES } = await import("../../src/lib/skills/memoryBuiltins.ts"); + + await updateSettings({ memoryEnabled: true, memoryMaxTokens: 2000 }); + invalidateMemorySettingsCache(); + + const body: Record = { + model: "gpt-4o", + stream: true, + messages: [{ role: "user", content: "hello" }], + }; + + const result = await injectMemoryAndSkills({ + body, + memoryOwnerId: "owner-stream", + provider: "openai", + effectiveModel: "gpt-4o", + sourceFormat: FORMATS.OPENAI, + targetFormat: FORMATS.OPENAI, + backgroundReason: null, + log: { debug: () => {} }, + }); + + assert.equal(result.memorySettings?.enabled, true); + const tools = (result.body.tools as { function?: { name?: string }; name?: string }[] | undefined) ?? []; + const toolNames = tools.map((tool) => tool.function?.name ?? tool.name); + for (const memoryTool of MEMORY_BUILTIN_TOOL_NAMES) { + assert.equal( + toolNames.includes(memoryTool), + false, + `expected ${memoryTool} to be absent for stream requests (client-side MCP path)` + ); + } + + invalidateMemorySettingsCache(); +}); + +test("injectMemoryAndSkills does not inject memory tools when memory is disabled", async () => { + const { updateSettings } = await import("../../src/lib/db/settings.ts"); + const { invalidateMemorySettingsCache } = await import("../../src/lib/memory/settings.ts"); + const { MEMORY_BUILTIN_TOOL_NAMES } = await import("../../src/lib/skills/memoryBuiltins.ts"); + + await updateSettings({ memoryEnabled: false }); + invalidateMemorySettingsCache(); + + const body: Record = { + model: "gpt-4o", + messages: [{ role: "user", content: "hello" }], + }; + + const result = await injectMemoryAndSkills({ + body, + memoryOwnerId: "owner-mem-off", + provider: "openai", + effectiveModel: "gpt-4o", + sourceFormat: FORMATS.OPENAI, + targetFormat: FORMATS.OPENAI, + backgroundReason: null, + log: { debug: () => {} }, + }); + + const tools = (result.body.tools as { function?: { name?: string }; name?: string }[] | undefined) ?? []; + const toolNames = tools.map((tool) => tool.function?.name ?? tool.name); + for (const memoryTool of MEMORY_BUILTIN_TOOL_NAMES) { + assert.equal( + toolNames.includes(memoryTool), + false, + `expected ${memoryTool} to be absent when memory is disabled` + ); + } + + invalidateMemorySettingsCache(); +}); diff --git a/tests/unit/skills-memory-builtins.test.ts b/tests/unit/skills-memory-builtins.test.ts new file mode 100644 index 0000000000..aaabcdf0a0 --- /dev/null +++ b/tests/unit/skills-memory-builtins.test.ts @@ -0,0 +1,227 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-memory-builtins-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const coreDb = await import("../../src/lib/db/core.ts"); +const { + memoryBuiltinHandlers, + buildMemoryToolsForProvider, + MEMORY_SAVE_TOOL_NAME, + MEMORY_UPDATE_TOOL_NAME, + MEMORY_SEARCH_TOOL_NAME, + MEMORY_DELETE_TOOL_NAME, +} = await import("../../src/lib/skills/memoryBuiltins.ts"); +const { interceptToolCalls } = await import("../../src/lib/skills/interception.ts"); +const { listMemories } = await import("../../src/lib/memory/store.ts"); + +function getMemoryMap() { + return { apiKeyId: "key-mem", sessionId: "session-mem" }; +} + +test.beforeEach(() => { + coreDb.resetDbInstance(); + fs.rmSync(path.join(TEST_DATA_DIR, "storage.sqlite"), { force: true }); + fs.rmSync(path.join(TEST_DATA_DIR, "storage.sqlite-wal"), { force: true }); + fs.rmSync(path.join(TEST_DATA_DIR, "storage.sqlite-shm"), { force: true }); +}); + +test.after(() => { + coreDb.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("memory_save creates a new memory entry", async () => { + const result = await memoryBuiltinHandlers[MEMORY_SAVE_TOOL_NAME]( + { key: "preference:coffee", content: "prefers dark roast", type: "factual" }, + getMemoryMap() + ); + assert.equal(result.success, true); + assert.equal(result.memory.key, "preference:coffee"); + assert.equal(result.memory.content, "prefers dark roast"); + assert.equal(result.memory.type, "factual"); + + const stored = await listMemories({ apiKeyId: "key-mem" }); + assert.equal(stored.data.length, 1); +}); + +test("memory_save upserts when the key already exists", async () => { + await memoryBuiltinHandlers[MEMORY_SAVE_TOOL_NAME]( + { key: "fact:city", content: "lives in Berlin" }, + getMemoryMap() + ); + const second = await memoryBuiltinHandlers[MEMORY_SAVE_TOOL_NAME]( + { key: "fact:city", content: "lives in Madrid" }, + getMemoryMap() + ); + + assert.equal(second.success, true); + assert.equal(second.memory.content, "lives in Madrid"); + const stored = await listMemories({ apiKeyId: "key-mem" }); + assert.equal(stored.data.length, 1, "same key must upsert, not duplicate"); +}); + +test("memory_save rejects missing key or content", async () => { + await assert.rejects( + () => memoryBuiltinHandlers[MEMORY_SAVE_TOOL_NAME]({ content: "no key" }, getMemoryMap()), + /Missing required field: key/ + ); + await assert.rejects( + () => memoryBuiltinHandlers[MEMORY_SAVE_TOOL_NAME]({ key: "k" }, getMemoryMap()), + /Missing required field: content/ + ); +}); + +test("memory_search returns saved memories by query and type", async () => { + await memoryBuiltinHandlers[MEMORY_SAVE_TOOL_NAME]( + { key: "pref:color", content: "likes green", type: "factual" }, + getMemoryMap() + ); + const found = await memoryBuiltinHandlers[MEMORY_SEARCH_TOOL_NAME]( + { query: "green", type: "factual" }, + getMemoryMap() + ); + assert.equal(found.success, true); + assert.equal(found.data.count, 1); + assert.equal(found.data.memories[0].content, "likes green"); + + const none = await memoryBuiltinHandlers[MEMORY_SEARCH_TOOL_NAME]( + { type: "episodic" }, + getMemoryMap() + ); + assert.equal(none.data.count, 0); +}); + +test("memory_update changes content and re-saves", async () => { + const saved = await memoryBuiltinHandlers[MEMORY_SAVE_TOOL_NAME]( + { key: "fact:job", content: "works as engineer" }, + getMemoryMap() + ); + const updated = await memoryBuiltinHandlers[MEMORY_UPDATE_TOOL_NAME]( + { id: saved.memory.id, content: "works as architect" }, + getMemoryMap() + ); + assert.equal(updated.success, true); + + const found = await memoryBuiltinHandlers[MEMORY_SEARCH_TOOL_NAME]( + { query: "architect" }, + getMemoryMap() + ); + assert.equal(found.data.count, 1); + assert.equal(found.data.memories[0].content, "works as architect"); +}); + +test("memory_update rejects memory owned by another API key", async () => { + const saved = await memoryBuiltinHandlers[MEMORY_SAVE_TOOL_NAME]( + { key: "fact:secret", content: "mine" }, + getMemoryMap() + ); + await assert.rejects( + () => + memoryBuiltinHandlers[MEMORY_UPDATE_TOOL_NAME]( + { id: saved.memory.id, content: "theirs" }, + { apiKeyId: "key-other", sessionId: "s" } + ), + /does not belong/ + ); +}); + +test("memory_delete removes the entry and rejects foreign keys", async () => { + const saved = await memoryBuiltinHandlers[MEMORY_SAVE_TOOL_NAME]( + { key: "fact:temp", content: "to be deleted" }, + getMemoryMap() + ); + + await assert.rejects( + () => + memoryBuiltinHandlers[MEMORY_DELETE_TOOL_NAME]( + { id: saved.memory.id }, + { apiKeyId: "key-other", sessionId: "s" } + ), + /does not belong/ + ); + + const deleted = await memoryBuiltinHandlers[MEMORY_DELETE_TOOL_NAME]( + { id: saved.memory.id }, + getMemoryMap() + ); + assert.equal(deleted.success, true); + + const stored = await listMemories({ apiKeyId: "key-mem" }); + assert.equal(stored.data.length, 0); +}); + +test("buildMemoryToolsForProvider emits provider-shaped tool definitions", () => { + const openai = buildMemoryToolsForProvider("openai") as { + type: string; + function: { name: string; description: string; parameters: { required: string[] } }; + }[]; + assert.equal(openai.length, 4); + assert.equal(openai[0].type, "function"); + assert.equal(openai[0].function.name, MEMORY_SAVE_TOOL_NAME); + assert.deepEqual(openai[0].function.parameters.required, ["key", "content"]); + + const claude = buildMemoryToolsForProvider("anthropic") as { + name: string; + input_schema: { required: string[] }; + }[]; + assert.equal(claude.length, 4); + assert.equal(claude[0].name, MEMORY_SAVE_TOOL_NAME); + assert.ok(claude[0].input_schema, "anthropic tools use input_schema"); + + const gemini = buildMemoryToolsForProvider("google") as { + name: string; + parameters: { required: string[] }; + }[]; + assert.equal(gemini.length, 4); + assert.equal(gemini[2].name, MEMORY_SEARCH_TOOL_NAME); + assert.ok(gemini[0].parameters, "gemini tools use parameters"); +}); + +test("interceptToolCalls executes memory tools when allowed via builtinToolNames", async () => { + const results = await interceptToolCalls( + [ + { + id: "call-save", + name: MEMORY_SAVE_TOOL_NAME, + arguments: { key: "pref:tea", content: "likes oolong" }, + }, + ], + { + apiKeyId: "key-mem", + sessionId: "session-mem", + requestId: "request-mem", + builtinToolNames: [MEMORY_SAVE_TOOL_NAME], + } + ); + + assert.equal(results.length, 1); + assert.equal(results[0].id, "call-save"); + assert.equal(results[0].result.success, true); + + const stored = await listMemories({ apiKeyId: "key-mem" }); + assert.equal(stored.data.length, 1); + assert.equal(stored.data[0].content, "likes oolong"); +}); + +test("interceptToolCalls skips memory tools not allowed by builtinToolNames", async () => { + const results = await interceptToolCalls( + [ + { id: "call-x", name: MEMORY_DELETE_TOOL_NAME, arguments: { id: "anything" } }, + ], + { + apiKeyId: "key-mem", + sessionId: "session-mem", + requestId: "request-mem", + builtinToolNames: [], + } + ); + // The tool is not in the allowed builtin list, so it falls through to the + // custom-skill resolver, which has no such skill registered. + assert.equal(results.length, 1); + assert.match(String(results[0].result.error), /Skill not found/); +});