mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-21 14:42:20 +03:00
Compare commits
2 Commits
fix/8450-r
...
fix/9115-m
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
21c75528a8 | ||
|
|
9136d84007 |
1
changelog.d/fixes/10887-memory-mcp-tools.md
Normal file
1
changelog.d/fixes/10887-memory-mcp-tools.md
Normal file
@@ -0,0 +1 @@
|
||||
- **fix(memory):** enable agent memory save/update via MCP tools (`memory_save`/`update`/`search`/`delete` builtins with per-provider schemas, `apiKeyId` optional with caller-principal fallback) and gate server-side memory builtin injection to non-stream requests only ([#10887](https://github.com/diegosouzapw/OmniRoute/pull/10887)) — thanks @Egorich-print
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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<string, unknown> | null;
|
||||
if (!record || typeof record !== "object") return [];
|
||||
const fn = record.function as Record<string, unknown> | 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<string, unknown>;
|
||||
const name =
|
||||
(record.function as Record<string, unknown> | 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
|
||||
|
||||
@@ -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<string> {
|
||||
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<typeof MemorySearchSchema>) => {
|
||||
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<typeof MemoryAddSchema>) => {
|
||||
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<typeof MemoryClearSchema>) => {
|
||||
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)
|
||||
|
||||
@@ -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<string, string> = {
|
||||
[OMNIROUTE_WEB_FETCH_FALLBACK_TOOL_NAME]: "web_fetch",
|
||||
};
|
||||
|
||||
const MEMORY_TOOL_NAMES = new Set<string>(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<string, unknown> | 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,
|
||||
|
||||
294
src/lib/skills/memoryBuiltins.ts
Normal file
294
src/lib/skills/memoryBuiltins.ts
Normal file
@@ -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<void> {
|
||||
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<ReturnType<typeof createMemory>>) {
|
||||
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<string, unknown>, context: { apiKeyId: string; sessionId: string }) {
|
||||
const { type, key, content, metadata } = input as {
|
||||
type?: string;
|
||||
key: string;
|
||||
content: string;
|
||||
metadata?: Record<string, unknown>;
|
||||
};
|
||||
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<string, unknown>, context: { apiKeyId: string }) {
|
||||
const { id, type, key, content, metadata } = input as {
|
||||
id: string;
|
||||
type?: string;
|
||||
key?: string;
|
||||
content?: string;
|
||||
metadata?: Record<string, unknown>;
|
||||
};
|
||||
if (!id || typeof id !== "string") throw new Error("Missing required field: id");
|
||||
|
||||
await assertOwner(id, context.apiKeyId);
|
||||
|
||||
const updates: Record<string, unknown> = {};
|
||||
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<string, unknown>, 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<string, unknown>, 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<string, unknown>) => ({
|
||||
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<string, unknown>) => ({
|
||||
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<string, unknown>) => ({
|
||||
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();
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
|
||||
@@ -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<string, unknown> = {
|
||||
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<string, unknown> = {
|
||||
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<string, unknown> = {
|
||||
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();
|
||||
});
|
||||
|
||||
227
tests/unit/skills-memory-builtins.test.ts
Normal file
227
tests/unit/skills-memory-builtins.test.ts
Normal file
@@ -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/);
|
||||
});
|
||||
Reference in New Issue
Block a user