feat(ai-sdk): createAgentMemory, an AI SDK memory provider [stack 2/3]

.changeset/ai-sdk-memory-provider.md

@@ -0,0 +1,5 @@
+---
+"@upstash/agentkit-ai-sdk": minor
+---
+
+Add `createAgentMemory`, an AI SDK memory provider backed by Upstash Redis Search. Create it once, then per request use `memory.tools({ userId })` (the `recall_memory` / `save_memory` tools), `memory.prepareStep({ userId })` to inject the memories relevant to the latest user message into the instructions automatically (works with `generateText`, `streamText` and `ToolLoopAgent`, with `userId` optionally derived from `runtimeContext`), `memory.recallContext()` to build instructions yourself, and `memory.instructions` for tool-usage guidance. The model is never wrapped, and a failed recall never fails the call.

CLAUDE.md

@@ -62,6 +62,18 @@ implements TanStack AI's own backend contracts (see its section below) — keep
   wrapper**), `cachedTools`, `createMemoryTools`, `createSearchTools`,
   `createChatHistory` (→ `ChatHistory<UIMessage>`; also re-exports `ChatHistory`/`ChatRecord`/etc.).
   (**No model cache** — removed.)
+- **`createAgentMemory`** (`src/agent-memory.ts`) — the AI SDK *memory provider* shape (cf.
+  `vercel/ai` `content/docs/03-agents/06-memory.mdx`: Mem0/Letta/Supermemory/Hindsight/`@redis/vercel-ai-sdk`).
+  One factory over a core `AgentMemory` (`prefix`/`indexName`/`topK`/`minScore`/tool names), then per
+  request: `.tools({ userId })` (the same tools as `createMemoryTools`, via the shared internal
+  `buildMemoryTools`), `.prepareStep({ userId, query?, onError? })` (automatic recall: on step 0 it
+  appends a `<memories>` "treat as data" block to the instructions — string, message or array — and
+  the override carries forward), `.recallContext()`, `.instructions` (names the configured tools),
+  `.store`. **Deliberately NOT a model wrapper:** the other providers' "automatic memory" is a
+  `wrapLanguageModel` middleware (`memory.model()`, `withSupermemory(model)`, `mem0(modelId)`), which
+  this repo forbids; `prepareStep` is the allowed seam and works for `generateText`/`streamText`/
+  `ToolLoopAgent` (userId may come from `runtimeContext`). Recall failures are swallowed (`onError`).
+  No automatic extraction/recording — saving stays explicit; transcripts go through `createChatHistory`.
 - **Only `cachedTools`** (no singular `cachedTool` — removed): `cachedTools(map, { userId, redis?, ttlSeconds? })`
   takes a map of `tool()`-built tools (so each keeps inference) and caches each under its **map key as
   the toolName**, scoped to `userId` — so the user never passes a tool name.

examples/ai-sdk-demo/app/api/chat/route.ts

@@ -12,7 +12,7 @@ import { z } from "zod";
 import {
   Ratelimit,
   cachedTools,
-  createMemoryTools,
+  createAgentMemory,
   createRateLimit,
   createSearchTools,
 } from "@upstash/agentkit-ai-sdk";
@@ -46,8 +46,10 @@ export async function POST(req: Request) {
     return new Response("Rate limited", { status: 429 });
   }
 
-  // Agent memory tools — recall_memory / save_memory, scoped to this user.
-  const memoryTools = createMemoryTools({ redis, userId });
+  // Agent memory — recall_memory / save_memory tools scoped to this user, plus automatic recall: the
+  // memories relevant to the latest message are appended to the instructions via `prepareStep`.
+  const memory = createAgentMemory({ redis });
+  const memoryTools = memory.tools({ userId });
 
   // Schema-driven search tools — search / aggregate / count over the books index (created reactively).
   // The books index is intentionally SHARED across users, so no per-user scoping here.
@@ -71,9 +73,11 @@ export async function POST(req: Request) {
 
   const result = streamText({
     model,
+    instructions: `You are a helpful assistant.\n\n${memory.instructions}`,
     messages: await convertToModelMessages(messages),
     tools,
     stopWhen: stepCountIs(5),
+    prepareStep: memory.prepareStep({ userId }),
   });
 
   const history = getHistory();

packages/ai-sdk/README.md

@@ -7,6 +7,7 @@ from this package.
 | Import | Feature |
 | --- | --- |
 | `createChatHistory` | Durable chat history on Redis Search — save, list, and `$smart`-search a user's transcripts. |
+| `createAgentMemory` | Memory provider: memory tools **and** automatic recall into the instructions via `prepareStep`. |
 | `createMemoryTools` | `recall_memory` + `save_memory` tools so the model reads and writes long-term memory. |
 | `createSearchTools` | `search` / `aggregate` / `count` tools over a Redis Search index (this is how you do RAG). |
 | `createRateLimit` | A configured Upstash Ratelimit to call before the model. |
@@ -114,6 +115,52 @@ value. Memories are stored at `agentkit:memory:<userId>:<id>`.
 
 </details>
 
+### Memory provider
+
+`createAgentMemory` is the AI SDK [memory provider](https://ai-sdk.dev/docs/agents/memory#memory-providers)
+shape: create it once, then per request give the model the memory tools, inject the memories
+relevant to the latest user message automatically, or both. The model is **not** wrapped — recall
+runs in `prepareStep` and appends a `<memories>` block to your instructions on the first step (the
+override carries forward to later steps). Works with `generateText`, `streamText` and `ToolLoopAgent`.
+
+```ts
+import { createAgentMemory } from "@upstash/agentkit-ai-sdk";
+import { ToolLoopAgent } from "ai";
+
+const memory = createAgentMemory(); // once, at module level
+
+// per request, with `userId` from your auth
+const agent = new ToolLoopAgent({
+  model,
+  instructions: `You are a helpful assistant.\n\n${memory.instructions}`,
+  tools: memory.tools({ userId }), // recall_memory + save_memory
+  prepareStep: memory.prepareStep({ userId }), // relevant memories injected up front
+});
+
+await agent.generate({ prompt: "What editor do I use?" });
+```
+
+<details>
+<summary>Options</summary>
+
+- `createAgentMemory({ redis?, prefix?, indexName?, topK?, minScore?, recallToolName?, saveToolName?, enableTelemetry? })`
+- `memory.tools({ userId, topK?, minScore? })` — the same tools as `createMemoryTools`.
+- `memory.prepareStep({ userId, query?, topK?, minScore?, onError? })` — `userId` may be a function of
+  the `prepareStep` options (e.g. `({ runtimeContext }) => runtimeContext.userId` on a shared agent);
+  `query` defaults to the last user message's text. A failed recall never fails the call: the step
+  runs with its original instructions and `onError` gets the error. Already have a `prepareStep`?
+  Call `memory.prepareStep(...)(step)` inside yours and merge its `instructions`.
+- `memory.recallContext({ userId, query?, topK?, minScore? })` — the formatted block (or `undefined`),
+  for building instructions yourself.
+- `memory.instructions` — tool-usage guidance naming your configured tool names.
+- `memory.store` — the underlying `AgentMemory` for direct `add` / `list` / `get` / `forget`.
+
+Recall is Upstash Redis Search `$smart` (fuzzy, BM25) — no embeddings, no vector database. Saving is
+explicit (the `save_memory` tool or `memory.store.add`); nothing is extracted from transcripts
+behind your back. Persist the transcript itself with `createChatHistory`.
+
+</details>
+
 ## Search tools
 
 `search` / `aggregate` / `count` over an Upstash Redis Search index; the model-facing descriptions are

packages/ai-sdk/src/agent-memory.test.ts

@@ -0,0 +1,181 @@
+import { generateText, stepCountIs, tool, ToolLoopAgent, type ModelMessage } from "ai";
+import { z } from "zod";
+import { MockLanguageModelV4 } from "ai/test";
+import type { LanguageModelV4GenerateResult, LanguageModelV4Prompt } from "@ai-sdk/provider";
+import { Redis } from "@upstash/redis";
+import { afterAll, beforeAll, describe, expect, it } from "vitest";
+import { createAgentMemory } from "./agent-memory.js";
+import { cleanupKeys, hasRedisCreds, testRedis, uniqueUserId } from "./test-support.js";
+
+/** Poll a read until it reflects a just-written doc — insurance for residual indexing lag. */
+async function pollUntil<T>(read: () => Promise<T>, ready: (value: T) => boolean): Promise<T> {
+  const deadline = Date.now() + 8_000;
+  let value = await read();
+  while (!ready(value) && Date.now() < deadline) {
+    await new Promise((resolve) => setTimeout(resolve, 250));
+    value = await read();
+  }
+  return value;
+}
+
+const TOOL_OPTS = { toolCallId: "t", messages: [] } as never;
+function call<R>(execute: unknown, input: unknown): Promise<R> {
+  return (execute as (i: unknown, o: unknown) => Promise<R>)(input, TOOL_OPTS);
+}
+
+const RESULT: LanguageModelV4GenerateResult = {
+  content: [{ type: "text", text: "ok" }],
+  finishReason: { unified: "stop", raw: undefined },
+  usage: {
+    inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined },
+    outputTokens: { total: 1, text: 1, reasoning: undefined },
+  },
+  warnings: [],
+};
+
+/** The system text the mocked model received on its first call. */
+function systemText(prompt: LanguageModelV4Prompt): string {
+  return prompt
+    .filter((m) => m.role === "system")
+    .map((m) => m.content)
+    .join("\n");
+}
+
+describe("createAgentMemory (no Redis)", () => {
+  // A client for an unreachable host: none of these cases ever sends a request.
+  const memory = createAgentMemory({
+    redis: new Redis({ url: "https://unreachable.invalid", token: "unused" }),
+    recallToolName: "remember_what",
+    saveToolName: "note_down",
+  });
+
+  it("names the configured tools in its instructions and tool map", () => {
+    expect(memory.instructions).toContain("`remember_what`");
+    expect(memory.instructions).toContain("`note_down`");
+    expect(Object.keys(memory.tools({ userId: "u" })).sort()).toEqual([
+      "note_down",
+      "remember_what",
+    ]);
+  });
+
+  it("swallows a failed recall and reports it to onError", async () => {
+    const errors: unknown[] = [];
+    const model = new MockLanguageModelV4({ doGenerate: RESULT });
+    const result = await generateText({
+      model,
+      instructions: "base",
+      prompt: "hi",
+      // ':' is rejected by the core userId validation before any Redis call.
+      prepareStep: memory.prepareStep({ userId: "bad:id", onError: (e) => errors.push(e) }),
+    });
+    expect(result.text).toBe("ok");
+    expect(errors).toHaveLength(1);
+    expect(systemText(model.doGenerateCalls[0]!.prompt)).toBe("base");
+  });
+});
+
+describe.skipIf(!hasRedisCreds)("createAgentMemory (live Redis)", () => {
+  const redis = testRedis();
+  const userId = uniqueUserId("aisdk-provider");
+  const emptyUser = uniqueUserId("aisdk-provider-empty");
+  const memory = createAgentMemory({ redis });
+
+  // Provision the index before seeding: waitIndexing() on a missing index is a silent no-op.
+  beforeAll(async () => {
+    await memory.store.recall({ userId, query: "provisioning probe" });
+    await memory.store.add({ userId, text: "The user's favorite editor is Neovim" });
+    await memory.store.add({ userId, text: "The user lives in Lisbon" });
+    await memory.store.searchIndex.waitIndexing();
+    await pollUntil(
+      () => memory.store.list({ userId }),
+      (all) => all.length >= 2,
+    );
+  });
+
+  afterAll(async () => {
+    await cleanupKeys(redis, `agentkit:memory:${userId}`);
+  });
+
+  it("tools() saves and recalls through the shared store", async () => {
+    const tools = memory.tools({ userId });
+    const recalled = await pollUntil(
+      () => call<{ text: string }[]>(tools.recall_memory!.execute, { query: "editor" }),
+      (hits) => hits.some((h) => h.text.includes("Neovim")),
+    );
+    expect(recalled.some((h) => h.text.includes("Neovim"))).toBe(true);
+  });
+
+  it("recallContext formats hits and returns undefined for a user with none", async () => {
+    const block = await pollUntil(
+      () => memory.recallContext({ userId, query: "editor" }),
+      (b) => Boolean(b?.includes("Neovim")),
+    );
+    expect(block).toMatch(/^<memories>/);
+    expect(block).toContain("- The user's favorite editor is Neovim");
+    expect(await memory.recallContext({ userId: emptyUser })).toBeUndefined();
+  });
+
+  it("prepareStep appends relevant memories to string instructions on step 0 only", async () => {
+    await pollUntil(
+      () => memory.recallContext({ userId, query: "Which editor do I use?" }),
+      (b) => Boolean(b?.includes("Neovim")),
+    );
+    // Step 0 calls a tool, step 1 answers: the block is injected once and carries forward.
+    const toolCall: LanguageModelV4GenerateResult = {
+      ...RESULT,
+      content: [{ type: "tool-call", toolCallId: "c1", toolName: "ping", input: "{}" }],
+      finishReason: { unified: "tool-calls", raw: undefined },
+    };
+    const twoStep = new MockLanguageModelV4({ doGenerate: [toolCall, RESULT] });
+    await generateText({
+      model: twoStep,
+      instructions: "You are helpful.",
+      prompt: "Which editor do I use?",
+      tools: { ping: tool({ inputSchema: z.object({}), execute: async () => "pong" }) },
+      stopWhen: stepCountIs(3),
+      prepareStep: memory.prepareStep({ userId }),
+    });
+    expect(twoStep.doGenerateCalls).toHaveLength(2);
+    const system = systemText(twoStep.doGenerateCalls[0]!.prompt);
+    expect(system.startsWith("You are helpful.\n\n<memories>")).toBe(true);
+    expect(system).toContain("Neovim");
+    expect(systemText(twoStep.doGenerateCalls[1]!.prompt)).toBe(system);
+    expect(system.match(/<memories>/g)).toHaveLength(1);
+  });
+
+  it("works on a shared ToolLoopAgent with userId from runtimeContext and a custom query", async () => {
+    const model = new MockLanguageModelV4({ doGenerate: RESULT });
+    const queries: string[] = [];
+    const agent = new ToolLoopAgent({
+      model,
+      runtimeContext: { userId },
+      prepareStep: memory.prepareStep({
+        userId: ({ runtimeContext }) => runtimeContext.userId,
+        query: (messages: ModelMessage[]) => {
+          queries.push(String(messages.at(-1)?.content));
+          return "Lisbon";
+        },
+      }),
+    });
+    await pollUntil(
+      () => memory.recallContext({ userId, query: "Lisbon" }),
+      (b) => Boolean(b?.includes("Lisbon")),
+    );
+    await agent.generate({ prompt: "Where do I live?" });
+    expect(queries).toEqual(["Where do I live?"]);
+    const system = systemText(model.doGenerateCalls[0]!.prompt);
+    expect(system).toMatch(/^<memories>/);
+    expect(system).toContain("Lisbon");
+  });
+
+  it("leaves instructions untouched when nothing is recalled", async () => {
+    const model = new MockLanguageModelV4({ doGenerate: RESULT });
+    await generateText({
+      model,
+      instructions: "base",
+      prompt: "hello",
+      prepareStep: memory.prepareStep({ userId: emptyUser }),
+    });
+    expect(systemText(model.doGenerateCalls[0]!.prompt)).toBe("base");
+  });
+});