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

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -112,6 +112,7 @@ const tools = createMemoryAiTools({
namespace: "support",
profile: "user-42",
sessionId: "chat-7",
locked: true,
});

const { text } = await generateText({
Expand All @@ -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

Expand Down
46 changes: 33 additions & 13 deletions extensions/shared/memory-tool-definitions.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import {
} from "@agntn/tools";
import type { AbortOptions } from "@agntn/memory";
import {
createLockedMemoryToolSchemas,
createOperationalMemoryToolSchemas,
namespaceListParameters,
namespaceParameters,
Expand All @@ -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. */
Expand Down Expand Up @@ -64,14 +82,7 @@ function loadExecutors(): Promise<typeof MemoryTools> {
/** 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<MemoryToolName> = new Set([
"memory_namespace_list",
"memory_namespace_get",
"memory_namespace_create",
"memory_namespace_delete",
"memory_profile_list",
"memory_profile_delete",
]);
const administrationTools: ReadonlySet<MemoryToolName> = new Set(administrationToolNames);

/**
* Arguments as a call line shows them, with the default scope filled in where the call left it out.
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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,
Expand Down
49 changes: 39 additions & 10 deletions extensions/shared/memory-tool-schemas.ts
Original file line number Diff line number Diff line change
Expand Up @@ -359,6 +359,17 @@ const hostSessionSchemas: ReadonlySet<OperationalSchemaName> = 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`.
*
Expand All @@ -367,13 +378,35 @@ const hostSessionSchemas: ReadonlySet<OperationalSchemaName> = 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<OperationalSchemaName, TObject>} The rebuilt schemas.
*/
function rebuildOperationalSchemas(
rebuild: (name: OperationalSchemaName) => TObject,
): Record<OperationalSchemaName, TObject> {
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<OperationalSchemaName, TObject>} The schemas the model sees.
*/
export function createLockedMemoryToolSchemas(): Record<OperationalSchemaName, TObject> {
return rebuildOperationalSchemas(scopelessSchema);
}

/**
Expand All @@ -387,9 +420,5 @@ export function createOperationalMemoryToolSchemas(
): Record<OperationalSchemaName, TObject> {
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));
}
98 changes: 85 additions & 13 deletions src/ai.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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<MemoryAiToolOptions, "locked"> {
readonly namespace: string;
readonly profile: string;
/** No call can name another scope, and the namespace and profile tools are gone. */
readonly locked: true;
}

/**
Expand All @@ -38,35 +53,92 @@ 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<Pick<MemoryAiToolOptions, "namespace" | "profile">>,
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<ProfileMemoryToolName, Tool>} 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<ProfileMemoryToolName, Tool>;
/**
* Every memory tool for the AI SDK, keyed by the name MCP, Pi and OMP use for it.
*
* @param options - Namespace, profile and write session the app binds the tools to.
* @returns {Record<MemoryToolName, Tool>} 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<MemoryToolName, Tool>;
/**
* Build the AI SDK tools for either form.
*
* @param options - The app's binding, locked or not.
* @returns {Partial<Record<MemoryToolName, Tool>>} 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<MemoryToolName, Tool> {
return toAiTools(createMemoryTools(memoryToolOptions(options))) as Record<MemoryToolName, Tool>;
options: MemoryAiToolOptions | LockedMemoryAiToolOptions = {},
): Partial<Record<MemoryToolName, Tool>> {
return toAiTools(createMemoryTools(memoryToolOptions(options))) as Partial<
Record<MemoryToolName, Tool>
>;
}
Loading
Loading