diff --git a/README.md b/README.md index de8f5ca..fa00fc5 100644 --- a/README.md +++ b/README.md @@ -98,7 +98,7 @@ Don't have an API token? The tools can also use `cf` CLI authentication. Run `cf ### AI SDK -You get the same 17 tools as MCP, Pi, and OMP. No more writing your own wrappers. Install the optional AI SDK peer, version 7 or newer. +For a per-user app, lock memory to that user. No more relying on the model to respect the neighbours. Install the optional AI SDK peer, version 7 or newer. ```bash pnpm add @agntn/memory ai @@ -112,6 +112,7 @@ const tools = createMemoryAiTools({ namespace: "support", profile: "user-42", sessionId: "chat-7", + locked: true, }); const { text } = await generateText({ @@ -122,7 +123,13 @@ const { text } = await generateText({ }); ``` -Set namespace and profile together, or tool creation throws. Calls can name both to reach another profile, but never just one. With a profile set, the model can't override your write session. No Claude Code session hides in your chat app: scope and session come from your code, not environment variables. +Locked gives you 11 tools that work inside one profile. The six `memory_namespace_*` and `memory_profile_*` tools aren't returned, and TypeScript knows they're missing. + +No tool schema has a `namespace` or `profile` field. Calls that name either are refused before reaching Cloudflare. The model in user-42's chat can't read or delete user-43's memories. Creation throws unless both namespace and profile are set. + +Without `locked`, you get the same 17 tools as MCP, Pi, and OMP. That's for an agent managing its own memory. Namespace and profile must still go together, or creation throws. Calls can name both to reach another profile, never just one. + +With a profile set, writes carry your `sessionId`, and the model can't pass its own. Scope and session come from your code, never environment variables. No Claude Code session hides in your chat app. ## Behavior worth knowing diff --git a/extensions/shared/memory-tool-definitions.ts b/extensions/shared/memory-tool-definitions.ts index 0f415ef..adf38cf 100644 --- a/extensions/shared/memory-tool-definitions.ts +++ b/extensions/shared/memory-tool-definitions.ts @@ -8,6 +8,7 @@ import { } from "@agntn/tools"; import type { AbortOptions } from "@agntn/memory"; import { + createLockedMemoryToolSchemas, createOperationalMemoryToolSchemas, namespaceListParameters, namespaceParameters, @@ -27,12 +28,29 @@ import { toolErrorMessage } from "./tool-errors.ts"; export type MemoryToolName = keyof typeof memoryToolMeta; +const administrationToolNames = [ + "memory_namespace_list", + "memory_namespace_get", + "memory_namespace_create", + "memory_namespace_delete", + "memory_profile_list", + "memory_profile_delete", +] as const satisfies readonly MemoryToolName[]; + +/** The tools that work inside one profile, all a locked surface keeps. */ +export type ProfileMemoryToolName = Exclude< + MemoryToolName, + (typeof administrationToolNames)[number] +>; + /** How each surface binds the calls it serves. */ export interface MemoryToolOptions { /** Namespace and profile a call gets when it names neither of them. */ readonly scope?: ProfileParams; /** The host's conversation session for a write, read from the call's `host`. */ readonly writeSession?: (host: unknown) => string | undefined; + /** Keep every call in `scope`: no scope fields for the model, and no administration tools. */ + readonly locked?: boolean; } /** How far the configured scope reaches into a call: administration, a profile read, or a write. */ @@ -64,14 +82,7 @@ function loadExecutors(): Promise { /** MCP sends a still-waiting note this often while a call can run up to its timeout. */ const PROGRESS_INTERVAL_MS = 30_000; -const administrationTools: ReadonlySet = new Set([ - "memory_namespace_list", - "memory_namespace_get", - "memory_namespace_create", - "memory_namespace_delete", - "memory_profile_list", - "memory_profile_delete", -]); +const administrationTools: ReadonlySet = new Set(administrationToolNames); /** * Arguments as a call line shows them, with the default scope filled in where the call left it out. @@ -192,11 +203,14 @@ async function runMemoryTool( /** * The tool table, in the order every surface lists it. * - * @param scope - Namespace and profile the host defaults to, if any. + * @param options - Scope binding of the surface, which decides the schemas the model sees. * @returns {readonly MemoryToolSpec[]} One entry per tool. */ -function memoryToolSpecs(scope: ProfileParams | undefined): readonly MemoryToolSpec[] { - const schemas = createOperationalMemoryToolSchemas(scope); +function memoryToolSpecs(options: MemoryToolOptions): readonly MemoryToolSpec[] { + const schemas = + options.locked === true + ? createLockedMemoryToolSchemas() + : createOperationalMemoryToolSchemas(options.scope); return [ { name: "memory_namespace_list", @@ -378,10 +392,16 @@ function memoryToolSpecs(scope: ProfileParams | undefined): readonly MemoryToolS * Every memory tool as one `@agntn/tools` definition, for every surface from MCP to the AI SDK. * * @param options - Scope and session binding of the surface serving them. - * @returns {readonly ToolDefinition[]} The definitions, in listing order. + * @returns {readonly ToolDefinition[]} The definitions in listing order, locked ones without admin. + * @throws {Error} When the surface is locked without a scope to lock it to. */ export function createMemoryTools(options: MemoryToolOptions = {}): readonly ToolDefinition[] { - return memoryToolSpecs(options.scope).map((spec) => + const locked = options.locked === true; + if (locked && options.scope === undefined) { + throw new Error("Locked memory tools need a namespace and profile to lock to"); + } + const specs = memoryToolSpecs(options); + return (locked ? specs.filter((spec) => spec.reach !== "none") : specs).map((spec) => defineTool({ name: spec.name, title: memoryToolMeta[spec.name].label, diff --git a/extensions/shared/memory-tool-schemas.ts b/extensions/shared/memory-tool-schemas.ts index 73abb87..f67e9a2 100644 --- a/extensions/shared/memory-tool-schemas.ts +++ b/extensions/shared/memory-tool-schemas.ts @@ -359,6 +359,17 @@ const hostSessionSchemas: ReadonlySet = new Set([ "ingestParameters", ]); +/** + * One operational schema with no namespace or profile, and no `sessionId` on a write. + * + * @param name - Which raw schema to rebuild. + * @returns {TObject} The closed schema a locked tool takes. + */ +function scopelessSchema(name: OperationalSchemaName): TObject { + const hostOwned = hostSessionSchemas.has(name) ? ["sessionId"] : []; + return Type.Omit(rawOperationalSchemas[name], ["namespace", "profile", ...hostOwned], closed); +} + /** * One operational schema under a bound scope: the pair turns optional, a write loses `sessionId`. * @@ -367,13 +378,35 @@ const hostSessionSchemas: ReadonlySet = new Set([ * @returns {TObject} The closed schema the model sees. */ function boundSchema(name: OperationalSchemaName, target: string): TObject { - const hostOwned = hostSessionSchemas.has(name) ? ["sessionId"] : []; - const rest = Type.Omit( - rawOperationalSchemas[name], - ["namespace", "profile", ...hostOwned], + return Type.Object( + { ...boundProfileParameters(target), ...scopelessSchema(name).properties }, closed, ); - return Type.Object({ ...boundProfileParameters(target), ...rest.properties }, closed); +} + +/** + * Rebuild every operational schema the same way. + * + * @param rebuild - Turns one raw schema name into the schema the model sees. + * @returns {Record} The rebuilt schemas. + */ +function rebuildOperationalSchemas( + rebuild: (name: OperationalSchemaName) => TObject, +): Record { + const names = Object.keys(rawOperationalSchemas) as OperationalSchemaName[]; + return Object.fromEntries(names.map((name) => [name, rebuild(name)])) as Record< + OperationalSchemaName, + TObject + >; +} + +/** + * The schemas of a surface locked to one profile: no scope fields, and no `sessionId` on a write. + * + * @returns {Record} The schemas the model sees. + */ +export function createLockedMemoryToolSchemas(): Record { + return rebuildOperationalSchemas(scopelessSchema); } /** @@ -387,9 +420,5 @@ export function createOperationalMemoryToolSchemas( ): Record { if (scope === undefined) return rawOperationalSchemas; const target = `${scope.namespace}/${scope.profile}`; - const names = Object.keys(rawOperationalSchemas) as OperationalSchemaName[]; - return Object.fromEntries(names.map((name) => [name, boundSchema(name, target)])) as Record< - OperationalSchemaName, - TObject - >; + return rebuildOperationalSchemas((name) => boundSchema(name, target)); } diff --git a/src/ai.ts b/src/ai.ts index 245f438..285c090 100644 --- a/src/ai.ts +++ b/src/ai.ts @@ -6,11 +6,16 @@ import { createMemoryTools, type MemoryToolName, type MemoryToolOptions, + type ProfileMemoryToolName, } from "../extensions/shared/memory-tool-definitions.ts"; +import type { ProfileParams } from "../extensions/shared/memory-tools.types.ts"; import { validateMemoryScope } from "../extensions/shared/memory-scope.ts"; import { MAX_TOOL_SESSION_ID_CHARACTERS } from "../extensions/shared/memory-tool-contract.ts"; -export type { MemoryToolName } from "../extensions/shared/memory-tool-definitions.ts"; +export type { + MemoryToolName, + ProfileMemoryToolName, +} from "../extensions/shared/memory-tool-definitions.ts"; /** Where the tools land. The app names it, and no scope variable in the env gets a vote. */ export interface MemoryAiToolOptions { @@ -20,6 +25,16 @@ export interface MemoryAiToolOptions { readonly profile?: string; /** Session for a write that names none. Bound to a profile, the model can't name one. */ readonly sessionId?: string; + /** Leave it out, or use {@link LockedMemoryAiToolOptions} to lock the tools to the profile. */ + readonly locked?: false; +} + +/** One profile and nothing else, for an app that serves many users from one account. */ +export interface LockedMemoryAiToolOptions extends Omit { + readonly namespace: string; + readonly profile: string; + /** No call can name another scope, and the namespace and profile tools are gone. */ + readonly locked: true; } /** @@ -38,26 +53,73 @@ function validateSessionId(sessionId: string): string { return sessionId; } +/** + * Read `locked` strictly: a plain JS caller passing `"true"` must not get unlocked tools. + * + * @param locked - The option as the app passed it. + * @returns {boolean} Whether the tools are locked. + * @throws {TypeError} When the option is neither boolean nor left out. + */ +function lockedOption(locked: unknown): boolean { + if (locked === undefined) return false; + if (typeof locked !== "boolean") { + throw new TypeError("Memory AI tool option locked must be a boolean"); + } + return locked; +} + +/** + * The scope to bind. A locked set needs one, so it can't fall back to letting the model choose. + * + * @param options - Namespace and profile from the app. + * @param locked - Whether the tools are locked. + * @returns {ProfileParams | undefined} The checked pair, or nothing for unbound tools. + * @throws {Error} When the pair is half set, missing under a lock, or a name is invalid. + */ +function memoryToolScope( + options: Readonly>, + locked: boolean, +): ProfileParams | undefined { + const { namespace, profile } = options; + if (!locked && namespace === undefined && profile === undefined) return undefined; + return validateMemoryScope( + namespace, + profile, + locked ? "Locked memory AI tool options" : "Memory AI tool options", + ); +} + /** * Turn the app's options into the binding the shared definitions take, checked up front. * - * @param options - Namespace, profile and write session from the app. - * @returns {MemoryToolOptions} The scope and write session for every tool. - * @throws {Error} When the pair is half set, a name is invalid, or the session id is unusable. + * @param options - Namespace, profile, write session and lock from the app. + * @returns {MemoryToolOptions} The scope, write session and lock for every tool. + * @throws {Error} When the pair is half set or missing under a lock, or a name or id is invalid. */ -function memoryToolOptions(options: MemoryAiToolOptions): MemoryToolOptions { - const { namespace, profile, sessionId } = options; - const scope = - namespace === undefined && profile === undefined - ? undefined - : validateMemoryScope(namespace, profile, "Memory AI tool options"); +function memoryToolOptions( + options: MemoryAiToolOptions | LockedMemoryAiToolOptions, +): MemoryToolOptions { + const locked = lockedOption(options.locked); + const scope = memoryToolScope(options, locked); + const { sessionId } = options; const session = sessionId === undefined ? undefined : validateSessionId(sessionId); return { ...(scope === undefined ? {} : { scope }), ...(session === undefined ? {} : { writeSession: () => session }), + ...(locked ? { locked } : {}), }; } +/** + * The profile tools for the AI SDK, locked to the app's namespace and profile. + * + * @param options - Namespace, profile and write session the tools never leave. + * @returns {Record} The AI SDK tools, in listing order. + * @throws {Error} When the pair is incomplete, a name is invalid, or the session id is unusable. + */ +export function createMemoryAiTools( + options: LockedMemoryAiToolOptions, +): Record; /** * Every memory tool for the AI SDK, keyed by the name MCP, Pi and OMP use for it. * @@ -65,8 +127,18 @@ function memoryToolOptions(options: MemoryAiToolOptions): MemoryToolOptions { * @returns {Record} The AI SDK tools, in listing order. * @throws {Error} When the pair is half set, a name is invalid, or the session id is unusable. */ +export function createMemoryAiTools(options?: MemoryAiToolOptions): Record; +/** + * Build the AI SDK tools for either form. + * + * @param options - The app's binding, locked or not. + * @returns {Partial>} The AI SDK tools, in listing order. + * @throws {Error} When the options can't bind the tools as written. + */ export function createMemoryAiTools( - options: MemoryAiToolOptions = {}, -): Record { - return toAiTools(createMemoryTools(memoryToolOptions(options))) as Record; + options: MemoryAiToolOptions | LockedMemoryAiToolOptions = {}, +): Partial> { + return toAiTools(createMemoryTools(memoryToolOptions(options))) as Partial< + Record + >; } diff --git a/test/ai.test.ts b/test/ai.test.ts index 1538564..4f21fec 100644 --- a/test/ai.test.ts +++ b/test/ai.test.ts @@ -1,6 +1,6 @@ import { generateText, isStepCount, type Tool } from "ai"; import { MockLanguageModelV4 } from "ai/test"; -import { afterEach, describe, expect, it, vi } from "vite-plus/test"; +import { afterEach, describe, expect, expectTypeOf, it, vi } from "vite-plus/test"; import { createMemoryAiTools } from "../src/ai.ts"; import { createMemoryTools } from "../extensions/shared/memory-tool-definitions.ts"; @@ -220,3 +220,118 @@ describe("@agntn/memory/ai", () => { expect(() => createMemoryAiTools(options)).toThrow(message); }); }); + +describe("@agntn/memory/ai locked to one profile", () => { + const locked = { namespace: "support", profile: "user-42", locked: true } as const; + + it("drops every namespace and profile tool and keeps the rest in listing order", () => { + const administration = [ + "memory_namespace_list", + "memory_namespace_get", + "memory_namespace_create", + "memory_namespace_delete", + "memory_profile_list", + "memory_profile_delete", + ]; + const tools = createMemoryAiTools(locked); + + expect(Object.keys(tools)).toEqual( + createMemoryTools() + .map((tool) => tool.name) + .filter((name) => !administration.includes(name)), + ); + expect(Object.keys(tools)).toHaveLength(11); + expectTypeOf(tools).not.toHaveProperty("memory_profile_delete"); + expectTypeOf(tools).toHaveProperty("memory_recall"); + }); + + it("offers the model no scope or session field in any schema", () => { + const tools = createMemoryTools({ + scope: { namespace: "support", profile: "user-42" }, + locked: true, + }); + + for (const tool of tools) { + expect(Object.keys(tool.input.properties), tool.name).not.toContain("namespace"); + expect(Object.keys(tool.input.properties), tool.name).not.toContain("profile"); + } + const writes = tools.filter((tool) => + ["memory_remember", "memory_supersede", "memory_ingest"].includes(tool.name), + ); + expect(writes).toHaveLength(3); + for (const tool of writes) { + expect(Object.keys(tool.input.properties), tool.name).not.toContain("sessionId"); + } + }); + + it("refuses a call that names another profile before any request", async () => { + const requests = stubService(); + const tools = createMemoryAiTools(locked); + + await expect( + run(tools.memory_recall, { namespace: "support", profile: "user-43", query: "who?" }), + ).rejects.toThrow('unknown property "namespace"'); + await expect( + run(tools.memory_delete, { namespace: "support", profile: "user-43", memoryId: "mem-1" }), + ).rejects.toThrow('unknown property "profile"'); + expect(requests).toEqual([]); + }); + + it("keeps a model aiming at the neighbour's profile at home through generateText", async () => { + const requests = stubService(); + const model = scriptedModel("memory_recall", [ + { namespace: "support", profile: "user-43", query: "What does user-43 prefer?" }, + { query: "What does this user prefer?" }, + ]); + + const run = await generateText({ + model, + prompt: "What does this user prefer?", + tools: createMemoryAiTools(locked), + stopWhen: isStepCount(3), + }); + + expect(run.steps[0]?.content).toContainEqual( + expect.objectContaining({ type: "tool-error", toolName: "memory_recall" }), + ); + expect(requests.map(({ path }) => path)).toEqual([ + "/namespaces/support/profiles/user-42/recall", + ]); + }); + + it("tags a locked write with the app's session", async () => { + const requests = stubService(); + const tools = createMemoryAiTools({ ...locked, sessionId: "chat-7" }); + + await run(tools.memory_remember, { content: "Prefers short answers." }); + await expect( + run(tools.memory_remember, { content: "hi", sessionId: "the-model-picked-this" }), + ).rejects.toThrow(/sessionId/); + + expect(requests).toEqual([ + { + path: "/namespaces/support/profiles/user-42/remember", + method: "POST", + body: { content: "Prefers short answers.", sessionId: "chat-7" }, + }, + ]); + }); + + it.each([ + [{ locked: true }, "Locked memory AI tool options must define both namespace and profile"], + [{ locked: true, namespace: "support" }, "must define both namespace and profile"], + [{ locked: true, profile: "user-42" }, "must define both namespace and profile"], + [{ ...locked, profile: "User-42" }, "Memory profile must match"], + [{ ...locked, locked: "true" }, "locked must be a boolean"], + ])("refuses %j before any tool exists", (options, message) => { + expect(() => { + Reflect.apply(createMemoryAiTools, undefined, [options]); + }).toThrow(message); + }); + + it("refuses a locked definition set with nothing to lock to", () => { + expect(() => createMemoryTools({ locked: true })).toThrow( + "Locked memory tools need a namespace and profile to lock to", + ); + }); +}); diff --git a/test/eval-packed.mjs b/test/eval-packed.mjs index 6cdcab3..0e44ed8 100644 --- a/test/eval-packed.mjs +++ b/test/eval-packed.mjs @@ -268,6 +268,23 @@ try { assert.match(requests.at(-1).url, /\/namespaces\/packed\/profiles\/app\/remember$/); assert.equal(requests.at(-1).body.sessionId, "chat-1"); + /* The locked form README leads with: no administration tools, and no way out of the profile. */ + const lockedTools = createMemoryAiTools({ namespace: "packed", profile: "app", locked: true }); + assert.deepEqual( + Object.keys(lockedTools), + expectedToolNames.filter((name) => !/^memory_(namespace|profile)_/.test(name)), + "wrong locked tool set from packed dist/ai.mjs", + ); + const sentBeforeEscape = requests.length; + await assert.rejects( + lockedTools.memory_recall.execute( + { namespace: "packed", profile: "other", query: "who?" }, + { toolCallId: "packed-ai-locked", messages: [], context: {} }, + ), + /unknown property "namespace"/, + ); + assert.equal(requests.length, sentBeforeEscape, "a locked call left its profile"); + /* An app that imports MemoryError has to recognize the one a failed tool call carries. */ failNextRequest = true; const aiFailure = await aiTools.memory_remember