diff --git a/.changeset/lua-key-locking.md b/.changeset/lua-key-locking.md new file mode 100644 index 0000000..6533d13 --- /dev/null +++ b/.changeset/lua-key-locking.md @@ -0,0 +1,5 @@ +--- +"@upstash/agentkit-eve": patch +--- + +`redisDocuments()` now runs its compare-and-swap Lua script with the `allow-key-locking` flag, so Upstash locks only the document's key instead of the whole database while the script runs. diff --git a/.changeset/tanstack-ai-initial.md b/.changeset/tanstack-ai-initial.md new file mode 100644 index 0000000..71d883a --- /dev/null +++ b/.changeset/tanstack-ai-initial.md @@ -0,0 +1,5 @@ +--- +"@upstash/agentkit-tanstack-ai": minor +--- + +New package: production backends for TanStack AI on Upstash Redis. `upstashPersistence()` covers every persistence store (messages, runs, interrupts, metadata, generation runs, artifacts, and with an Upstash Blob bucket, blobs) and passes TanStack's conformance suite with nothing skipped. Also `upstashStream()` (resumable `StreamDurability` on Redis Streams), `upstashLocks()` (distributed `LockStore`), `upstashMemory()` (a `MemoryAdapter` ranked in Redis Search, passing TanStack's memory contract), `toolCache()` and `rateLimit()` chat middlewares, and `createSearchTools()`. The Redis primitives underneath are exported too: `RedisLock` (a lease lock with fencing tokens) and `EventLog` (a resumable append-only log on Redis Streams). diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 31347bb..1027515 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -56,10 +56,28 @@ jobs: env: UPSTASH_REDIS_REST_URL: ${{ secrets.UPSTASH_REDIS_REST_URL }} UPSTASH_REDIS_REST_TOKEN: ${{ secrets.UPSTASH_REDIS_REST_TOKEN }} + UPSTASH_BLOB_TOKEN: ${{ secrets.UPSTASH_BLOB_TOKEN }} - name: Build example apps run: pnpm -r --filter "./examples/*" build + - name: E2E (TanStack AI demo, two instances, mocked model) + # Starts two production servers of examples/tanstack-ai-demo against the same Upstash Redis, + # with a scripted model (no provider key), and checks what only shared backends can do: + # resume a dropped stream on the other instance, rebuild a thread mid-run there, recall + # memory across threads and instances, and produce a doubly-POSTed run exactly once. + # Uses the app built by the step above; skipped without Redis secrets (e.g. fork PRs). + working-directory: examples/tanstack-ai-demo + env: + UPSTASH_REDIS_REST_URL: ${{ secrets.UPSTASH_REDIS_REST_URL }} + UPSTASH_REDIS_REST_TOKEN: ${{ secrets.UPSTASH_REDIS_REST_TOKEN }} + run: | + if [ -z "$UPSTASH_REDIS_REST_URL" ]; then + echo "No Redis secrets available — skipping the TanStack AI e2e suite." + exit 0 + fi + pnpm e2e + - name: E2E eval (eve memory slots, mocked model) # Boots the eve demo's agent with a scripted mockModel (no model provider, no # OPENAI_API_KEY) and asserts both Upstash Redis memory integrations run at eve's real diff --git a/CLAUDE.md b/CLAUDE.md index 78913f7..e42cd95 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -17,11 +17,14 @@ embeddings — keep that in mind when naming/among scoring. | `@upstash/agentkit-sdk` (`packages/sdk`) | Core, framework-agnostic primitives. **No `ai` dependency** (redis-only). | | `@upstash/agentkit-ai-sdk` (`packages/ai-sdk`) | Vercel AI SDK adapter. | | `@upstash/agentkit-eve` (`packages/eve`) | Eve framework adapter. Depends on the ai-sdk package. | +| `@upstash/agentkit-tanstack-ai` (`packages/tanstack-ai`) | TanStack AI backends: persistence stores, `StreamDurability`, `LockStore`, `MemoryAdapter`, middlewares, search tools. | | `@upstash/agentkit-eve-extension` (`packages/eve-extension`) | AgentKit as a mountable **eve extension** (eve ≥0.24): one `agent/extensions/.ts` file composes memory tools, search tools, a chat-history hook, and an instructions fragment under `__*`. | Examples (`examples/`): `ai-sdk-demo` (hand-written Next.js), `eve-demo` (a real `eve` CLI scaffold), and `eve-extension-demo` (a minimal eve scaffold that mounts the extension). -`langchain` and `tanstack-ai` packages were **removed** — don't reintroduce them. +`langchain` was **removed** — don't reintroduce it. A first `tanstack-ai` adapter was removed in June 2026 +(it only wrapped memory + a model cache); the current `packages/tanstack-ai` is a different package that +implements TanStack AI's own backend contracts (see its section below) — keep it. ### Core SDK exports (`@upstash/agentkit-sdk`) - `AgentMemory`, `ToolCache`, `ChatHistory`, `createSearchToolDefs` (the framework-agnostic search-tool @@ -416,6 +419,20 @@ and `eve-extension-demo` (a minimal eve scaffold that mounts the extension). - **Removed entirely:** the model cache (`ModelCache`/`SemanticCache`, `cachedModel`, `modelCacheMiddleware`), Telemetry, the generic Sandbox (sandbox is eve-only), and dead core exports `ChatMessage`/`Logger`/`noopLogger`. +## Lua scripts: always `allow-key-locking` + +- **Every `EVAL` script starts with `#!lua flags=allow-key-locking` on its very first line** (same as + `@upstash/ratelimit` since #154). Upstash then locks only the keys the script declares instead of the + whole database. The price: **every key a script touches must be in `KEYS`** — a key built or read + inside Lua fails with `ERR Dynamic keys are not allowed in Lua scripts when 'allow-key-locking' flag is + set` (verified live 2026-09-27). Don't pass placeholder keys for optional indexes either; make `KEYS` + variable-length and test `if KEYS[n]` in the script. +- Scripts that need a key they can only learn by reading (the artifact store's old run/thread index) + read it first, declare it, and compare-and-swap inside the script, retrying on a mismatch. +- Current scripts: `RedisLock` acquire/release/extend (tanstack-ai), eve `redisDocuments()` CAS, and the + tanstack-ai persistence/generation stores. Upstash tolerates a leading newline before the shebang, + standard Redis does not — keep it on line 1. + ## API conventions - **Naming of the knobs (consistent across all features):** - `prefix` — the base `agentkit:X` key prefix (config level). `ToolCache`/`AgentMemory`/`ChatHistory` @@ -1135,3 +1152,116 @@ find node_modules -path "*/zod/package.json" | while read f; do echo "$f $(node - ~~`gpt-5.4-mini` (demo model) may not exist~~ — verified live (2026-08): it exists and responds in both demos. No swap needed. - The `19.2.17` `@types/react` may linger as an unpruned orphan in `.pnpm`; harmless (nothing links it). + +## tanstack-ai (`packages/tanstack-ai`) + +- **Entry points** (tsup + `exports`): `.` = everything that needs only `@tanstack/ai` (stream, locks, + middleware, search, `RedisLock`/`EventLog`); `./persistence` (`src/persistence/index.ts`, needs + `@tanstack/ai-persistence`) and `./memory` (`src/memory/index.ts`, needs `@tanstack/ai-memory`). + Rule: **the root entry must never import types from an optional peer** — check `dist/index.d.ts` + has no `@tanstack/ai-persistence` / `@tanstack/ai-memory` / `@upstash/blob` import after a build. + (Copilot review on #48; verified 2026-09-28 with a consumer that has only `@tanstack/ai` and + `skipLibCheck: false` — the only errors are two pre-existing ones inside `@tanstack/ai`'s own + `spawn.d.ts`.) +- **Demo + E2E** (`examples/tanstack-ai-demo`): Next.js port of TanStack's `ts-react-chat` persistent-chat + route onto our backends (detached run, `resumeServerSentEventsResponse`, `reconstructChat`, memory), + with `RedisLock` replacing its process-local "one producer per run" set: `withLock(runId, fn, + { acquireTimeoutMs: 0 })` (timeout = someone else produces it, just tail), lease renewed for the run, + and on lease loss the run is aborted and the producer neither appends nor closes the log (the new + owner may be writing). The producer promise goes to Next.js `after()` (+ `maxDuration`) so it + survives the response on serverless; truly durable background runs (QStash) are a follow-up. `AGENTKIT_MOCK_MODEL=1` swaps in + a scripted model (`lib/model.ts`: word-by-word stream, `remember:` → `save_memory`, echoes recalled + memories). `pnpm e2e` (`e2e/`, files `*.e2e.ts` so root `pnpm test` skips them) starts **two** `next start` + servers on free ports, one Redis, a per-run `DEMO_PREFIX`, and checks cross-instance resume, mid-run + reconstruct, memory across threads, single production of a doubly-POSTed run, single production of + a run that outlives its lease (`PRODUCER_LEASE_MS=1000` + a `long:` ~6s mock answer), and a 400 on a + bad body. Count `TEXT_MESSAGE_START` to detect a second producer: two producers **interleave** their + words in one log, so text regexes miss it (a no-renewal build passed a regex check and failed this). + Gotchas hit building it: AG-UI bodies need a message `id`; Next.js route handlers don't return a thrown + `Response` (catch `chatParamsFromRequest`'s 400 yourself); spawn `node_modules/.bin/next` detached and + kill the process group (killing `npx` orphans `next`); backends are created lazily so `next build` + needs no credentials. CI runs it after "Build example apps"; skipped without Redis secrets. +- **`src/package-exports.test.ts`** pins `package.json#exports` to the tsup entries — added after a + pushed commit shipped `dist/persistence.js` without exporting it (only the demo's build caught it). +- **Layout** (`src/`): one folder per feature — `persistence/` (stores, `records.ts` plumbing, + `blob-store.ts`), `stream/` (`upstashStream` + `EventLog`), `locks/` (`upstashLocks` + `RedisLock`), + `memory/`, `middleware/`, `search/`, and `testing/` (scripted adapter, test bucket, env helpers — + never imported by `index.ts`). Root: `index.ts`, `telemetry.ts`, `version.ts`, `integration.test.ts`. +- **Redis primitives** (`src/locks/redis-lock.ts`, `src/stream/event-log.ts`; moved here from core 2026-09-28, exported from this package): `RedisLock` (+ `LockLease`, `LockAcquireTimeoutError`, + `LockLostError`) — `SET NX PX` lease + a never-expiring `INCR` fencing counter, release/extend are + ownership-checked Lua; `withLock` renews every `leaseMs/3` and aborts the section's signal on loss. + `EventLog` (+ `LogEntry`, `EventLogClosedError`) — Redis Streams: `append` is one `EVAL` script of `XADD`s (atomic, + ordered ids) that refuses a closed log and refreshes the stream TTL, + `read` resumes with an **exclusive** `XRANGE (id` and tails by polling (REST keeps no blocking + reads), `close` sets a separate `:closed` key and readers do a final drain after seeing it, so + nothing appended before close is lost. Values carry an `agentkit-event-v1:` marker because + `@upstash/redis` auto-deserializes stream fields (`"123"` would come back as `123`). Both verified + live over REST (2026-09-27): `MULTI`+`XADD`, `XRANGE (`, `EVAL`, `SET NX PX` all work. +- **Why it exists:** TanStack AI (≥0.61) defines backend contracts for production state and ships + only in-memory implementations (`memoryPersistence`, `memoryStream`, `InMemoryLockStore`, + `InMemoryRunStore`, the ai-memory `inMemory()`/ioredis-shaped `redis()` adapters). As of 2026-09 + there were no Redis/Postgres persistence backends on npm. This package fills the seams; it does + **not** wrap models or reinvent TanStack features. +- **Already upstream, don't duplicate:** `@tanstack/ai-sandbox-upstash-box` (Box sandbox provider) + is TanStack's own package. +- Exports: `upstashPersistence` (messages/runs/interrupts/metadata), `upstashStream` + (`StreamDurability`, offsets `upstash:v1::`), `upstashLocks` + (`LockStore` over `RedisLock`, `src/locks/`), `upstashMemory` + `memoryScopeKey` (`MemoryAdapter` over core + `AgentMemory`), `toolCache`/`rateLimit` middlewares + `RateLimitExceededError`, `createSearchTools` + (returns a `Tool[]` built with `toolDefinition().server()`), re-exported `createRateLimit`/`Ratelimit`. +- **Peers:** `@tanstack/ai` required; `@tanstack/ai-persistence` and `@tanstack/ai-memory` optional + (only their *types* are imported, so the runtime never needs them). `@tanstack/ai` is pinned exactly + as a devDep (0.61.0) — TanStack AI is pre-1.0 and moves fast; re-run the conformance suite on bumps. +- **Persistence layout** (`agentkit:tanstack:*`): every record (run, interrupt, generation run, + artifact, blob record, metadata value, transcript) is a **RedisJSON document**, so values keep their + JSON types and need no codec. Shared plumbing is in `src/records.ts`: + - `CREATE`: `JSON.SET … NX` + `ZADD` into each index, returns the stored doc (idempotent + `createOrResume`). `PATCH`: `EXISTS` guard (a bare `JSON.MERGE` on a missing key *creates* it, + verified live 2026-09-28) then `JSON.MERGE`. + - `toMergePatches` turns `{...existing, ...patch}` semantics into two merge patches: a key present + with `undefined` becomes `null` (merge-patch delete); object-valued fields are nulled in the first + patch and set in the second, because merge patches *merge* nested objects instead of replacing + them. A literal `null` in a patch therefore deletes the field. + - `checkRun` validates run status at read time (TanStack's readers act destructively on it). +- Indexes: `threadRuns:`/`parentRuns:` (zset by `startedAt`), `detachedRuns` (by + `detachedSince`, a **candidate** list — `listReclaimable` re-checks each record), + `threadInterrupts`/`runInterrupts` (by `requestedAt`), `threadGenerationRuns`, + `runArtifacts`/`threadArtifacts` (by `createdAt`, ties byte-ordered like the reference store). + `commitBatch` checks every interrupt is present and pending (`JSON.GET $.status`) and applies all + patches in one script. +- Artifacts: which indexes a doc is in lives in a side hash `artifactIndexes:` (raw index-key + strings) so the doc stays exactly the record. Re-saving under another run/thread: the caller reads + the side hash, declares the old index keys (required by `allow-key-locking`), and the script + compare-and-swaps, retrying on a race. +- Metadata: one doc per pair, key `meta::`, value + wrapped as `{ v }` so a bare string round-trips. +- **Blobs** (`blob-store.ts`, only when `upstashPersistence({ bucket })`): bytes in Upstash Blob at an + immutable **versioned** path `/`; `blobVersion:` points + at the current one and is swapped in the same script as the record (`PUT_RECORD` returns the replaced + path, deleted after commit; `DELETE_RECORD` likewise), and `get` reads record + pointer in one + `MULTI`, so concurrent puts can never leave a record describing another writer's bytes; the record is a JSON doc plus a lexical zset of keys. + `PUT_RECORD` replaces the doc but restores the old `createdAt` (read via `JSON.GET $.createdAt`). + Blob listings omit content type + metadata, hence records in Redis. Ranges use `signedReadUrl` + + `Range`; a 200 is sliced locally. `@upstash/blob` is an optional peer (structural + `BlobBucketLike`). Live-Blob conformance runs with `UPSTASH_BLOB_TOKEN` (get it via the MCP's + `blob_bucket get` + `include_credentials`): 26/26 on 2026-09-28. +- **Config types reuse the primitives and core:** `UpstashLocksConfig = Omit & {redis?}`, + `UpstashStreamConfig` from `EventLogConfig`, `ToolCacheMiddlewareConfig` from `ToolCacheConfig`, + `UpstashMemoryConfig` picks from `AgentMemoryConfig`, `CreateSearchToolsConfig` from + `SearchToolDefsConfig` — spread straight through, no per-field copying. +- **Memory:** own keyspace `agentkit:tanstackMemory` (its schema adds an indexed `source` field, so + it must not share `agentkit:memory` — see the eve memory-slot notes on why). Scope → `userId` via + `memoryScopeKey`: per user across threads by default, parts escaped so `.`/`_`/`:` can't forge a + collision. Recalled lines are labelled by source like eve's `redisMemory()`. +- **Testing:** `persistence.test.ts` runs TanStack's `runPersistenceConformance` (from + `@tanstack/ai-persistence/testkit`; it declares a vitest ^4 peer but runs fine on the repo's vitest + 2) with a fresh prefix per case — all seven stores, nothing skipped, 26/26 on 2026-09-27. `memory.test.ts` also runs + `runMemoryAdapterContract` (`@tanstack/ai-memory/testkit`), each scope pinned under a unique + tenant — it needs `waitForIndexing` (save provisions the index once, writes, then waits), exactly + like eve's `redisMemory()` capture. Middleware/memory/search tests drive a real + `chat()` agent loop through `src/test-adapter.ts` (a scripted `TextAdapter` emitting AG-UI + `TOOL_CALL_*`/`TEXT_MESSAGE_*`/`RUN_FINISHED` chunks) — no model provider needed. +- **Not built yet (proposed):** Code Mode isolate driver on Box (port of `@tanstack/ai-isolate-daytona`'s + need_tools/replay loop), Redis `SandboxInstanceStore`/`SandboxCheckpointStore` (+ Upstash Blob + `BlobStore`), QStash-backed background runs writing to `upstashStream`, and a QStash schedule for + `reapDetachedRuns`. diff --git a/README.md b/README.md index 7a94a74..901ae3d 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ are powered by [Upstash Redis Search](https://upstash.com/docs/redis/search/intr | [`@upstash/agentkit-sdk`](./packages/sdk) | Core, framework-agnostic primitives. | | [`@upstash/agentkit-ai-sdk`](./packages/ai-sdk) | Adapter for the [Vercel AI SDK](https://ai-sdk.dev). | | [`@upstash/agentkit-eve`](./packages/eve) | Adapter for the Vercel Eve framework. | +| [`@upstash/agentkit-tanstack-ai`](./packages/tanstack-ai) | Production backends for [TanStack AI](https://tanstack.com/ai) — chat persistence, resumable streams, distributed locks, memory, tool caching, rate limiting and search tools. | | [`@upstash/agentkit-eve-extension`](./packages/eve-extension) | The same capabilities as a mountable [Eve extension](https://eve.dev/docs/extensions) — one file in `agent/extensions/` adds memory tools, search tools, and durable chat history the agent can search. | ## Core features @@ -36,9 +37,11 @@ are powered by [Upstash Redis Search](https://upstash.com/docs/redis/search/intr ## Examples Runnable demos (real Upstash Redis + a mock/real model) live in [`examples/`](./examples): -[`ai-sdk-demo`](./examples/ai-sdk-demo), [`eve-demo`](./examples/eve-demo), and +[`ai-sdk-demo`](./examples/ai-sdk-demo), [`eve-demo`](./examples/eve-demo), [`eve-extension-demo`](./examples/eve-extension-demo) (an eve agent that mounts -`@upstash/agentkit-eve-extension`). +`@upstash/agentkit-eve-extension`), and [`tanstack-ai-demo`](./examples/tanstack-ai-demo) (a +TanStack AI chat whose persistence, resumable stream and memory work across server instances, with +a two-instance E2E suite). ## Development diff --git a/examples/tanstack-ai-demo/.env.example b/examples/tanstack-ai-demo/.env.example new file mode 100644 index 0000000..afb7212 --- /dev/null +++ b/examples/tanstack-ai-demo/.env.example @@ -0,0 +1,7 @@ +# Upstash Redis (https://console.upstash.com) +UPSTASH_REDIS_REST_URL= +UPSTASH_REDIS_REST_TOKEN= + +# OpenAI, or set AGENTKIT_MOCK_MODEL=1 to run a scripted model instead +OPENAI_API_KEY= +# OPENAI_MODEL=gpt-5.4-mini diff --git a/examples/tanstack-ai-demo/.gitignore b/examples/tanstack-ai-demo/.gitignore new file mode 100644 index 0000000..fdd6fde --- /dev/null +++ b/examples/tanstack-ai-demo/.gitignore @@ -0,0 +1,5 @@ +node_modules +.next +.env* +!.env.example +*.tsbuildinfo diff --git a/examples/tanstack-ai-demo/README.md b/examples/tanstack-ai-demo/README.md new file mode 100644 index 0000000..3e5e47a --- /dev/null +++ b/examples/tanstack-ai-demo/README.md @@ -0,0 +1,51 @@ +# TanStack AI on Upstash — demo + +A chat app built on [TanStack AI](https://tanstack.com/ai) with every piece of its state in Upstash +Redis, via [`@upstash/agentkit-tanstack-ai`](../../packages/tanstack-ai): + +- **Persistence**: transcripts, runs and approvals (`upstashPersistence`), so a reload rebuilds the + thread. +- **Resumable streams**: every chunk goes to a Redis Stream first (`upstashStream`), so a reload — + or the same URL on another device — continues the answer, on any server instance. +- **Detached runs**: closing the tab doesn't stop the model; the run finishes and is persisted. + Next.js `after()` keeps the invocation alive for it on serverless hosts (up to `maxDuration`). + `RedisLock` makes sure a duplicate request never starts a second run: its lease is renewed for the + whole run, and a producer that ever loses it stops. +- **Memory**: facts saved in one thread are recalled in the next (`upstashMemory`). + +It is TanStack AI's own persistent-chat pattern (from their `ts-react-chat` example) with the +in-memory and SQLite pieces swapped for Upstash, which is what makes it work across instances. + +## Run it + +```bash +cp .env.example .env # then fill in: +# UPSTASH_REDIS_REST_URL, UPSTASH_REDIS_REST_TOKEN, and OPENAI_API_KEY +pnpm dev +``` + +No OpenAI key? `AGENTKIT_MOCK_MODEL=1 pnpm dev` runs a scripted model instead. + +Try it: send a message, reload mid-answer, and watch it continue. Send `remember: I like green tea`, +then open a new thread and ask what it knows about you. + +## End-to-end tests + +```bash +pnpm build && pnpm e2e +``` + +Starts **two** production servers of this app against the same Upstash Redis, with the scripted +model, and checks what only shared backends can do: + +| Scenario | What it proves | +| --- | --- | +| A client drops mid-answer on A, resumes on B | The stream is resumable across instances, with nothing missed or repeated | +| A reload on B while A is still generating | `reconstructChat` sees the in-flight run, then the finished transcript | +| `remember:` in one thread on A, ask in a new thread on B | Memory is shared across threads and instances, and isolated per user | +| The same run POSTed to both instances at once | Exactly one model run (`RedisLock`), one answer, one transcript entry | +| A run re-POSTed after its producer lease would have expired | The lease is renewed for the whole run, so it is still produced once | +| A malformed request | A 400, not a 500 | + +Keys are written under a per-run prefix and deleted afterwards. Needs `UPSTASH_REDIS_REST_URL` and +`UPSTASH_REDIS_REST_TOKEN` (the suite skips without them). diff --git a/examples/tanstack-ai-demo/app/api/chat/route.ts b/examples/tanstack-ai-demo/app/api/chat/route.ts new file mode 100644 index 0000000..ad527e7 --- /dev/null +++ b/examples/tanstack-ai-demo/app/api/chat/route.ts @@ -0,0 +1,128 @@ +/** + * The chat endpoint — TanStack AI's "persistent chat" pattern, on Upstash instead of in-memory + * stores, so it works across server instances: + * + * - POST starts the model run **detached from the request** and answers by tailing the run's + * durable log. Closing the tab doesn't stop the run; it finishes, and is persisted. + * - GET with `Last-Event-ID` / `?offset` resumes a run from the log — on any instance. + * - GET with `?threadId` rebuilds the thread (`reconstructChat`): the transcript plus a pointer to + * the run still in flight, which the client then tails. + */ +import { + EventType, + chat, + chatParamsFromRequest, + chatParamsFromRequestBody, + maxIterations, + resumeServerSentEventsResponse, +} from "@tanstack/ai"; +import type { StreamChunk } from "@tanstack/ai"; +import { memoryMiddleware } from "@tanstack/ai-memory"; +import { LockAcquireTimeoutError } from "@upstash/agentkit-tanstack-ai"; +import { after } from "next/server"; +import { reconstructChat, withPersistence } from "@tanstack/ai-persistence"; +import { backends, runLog } from "@/lib/backends"; +import { chatModel } from "@/lib/model"; + +export const runtime = "nodejs"; +export const dynamic = "force-dynamic"; +/** The longest a detached run may take on a serverless host (seconds). */ +export const maxDuration = 300; + +type ChatParams = Awaited>; + +/** + * The user a request belongs to. A demo stand-in: a real app takes this from its auth session, + * never from something the client can set — it is what keeps one user's memories from another's. + */ +function userIdOf(request: Request): string { + return request.headers.get("x-demo-user") || "demo-user"; +} + +/** + * Produce a run into its durable log — exactly once across every instance. Resolves when the run is + * done, or right away when another instance already owns it (a client retry that landed on another + * server just tails the log the first one is writing). + * + * The lock is a lease renewed for as long as the run lasts, so long runs stay single-producer. If + * renewal ever reports the lease lost, the run is aborted and this producer stops touching the log, + * since another instance may now own it. + */ +function produceRun(params: ChatParams, userId: string): Promise { + const { persistence, memory, producerLock } = backends(); + return producerLock + .withLock( + params.runId, + async (leaseLost) => { + const abortController = new AbortController(); + leaseLost.addEventListener("abort", () => abortController.abort(leaseLost.reason), { + once: true, + }); + const log = runLog({ runId: params.runId }); + const stream = chat({ + adapter: chatModel(), + middleware: [ + withPersistence(persistence, { snapshotStreaming: true }), + memoryMiddleware({ adapter: memory, scope: { threadId: params.threadId, userId } }), + ], + agentLoopStrategy: maxIterations(5), + systemPrompts: ["You are a concise, friendly assistant."], + messages: params.messages, + threadId: params.threadId, + runId: params.runId, + abortController, + ...(params.parentRunId ? { parentRunId: params.parentRunId } : {}), + ...(params.resume ? { resume: params.resume } : {}), + }); + try { + for await (const chunk of stream) { + if (leaseLost.aborted) return; + await log.append([chunk]); + } + } catch (error) { + if (leaseLost.aborted) return; + await log.append([ + { + type: EventType.RUN_ERROR, + message: error instanceof Error ? error.message : String(error), + timestamp: Date.now(), + } as StreamChunk, + ]); + } finally { + if (!leaseLost.aborted) await log.close(); + } + }, + // Don't wait: a held lock means the run is already being produced elsewhere. + { acquireTimeoutMs: 0 }, + ) + .catch((error) => { + if (error instanceof LockAcquireTimeoutError) return; + throw error; + }); +} + +export async function POST(request: Request): Promise { + let params: ChatParams; + try { + params = await chatParamsFromRequest(request); + } catch (error) { + // A malformed body throws a 400 Response; Next.js route handlers don't return thrown ones. + if (error instanceof Response) return error; + throw error; + } + // Detached from this request: the run keeps going if the client leaves. `after` keeps the + // invocation alive until it finishes (up to the platform's max duration), which is what makes this + // safe on serverless; on a long-running server it is simply a background task. + after(produceRun(params, userIdOf(request))); + // Answer by reading the run's log from the start; the run itself is not tied to this response. + return resumeServerSentEventsResponse({ adapter: runLog({ runId: params.runId, offset: "-1" }) }); +} + +export async function GET(request: Request): Promise { + const log = runLog(request); + if (log.resumeFrom() !== null) return resumeServerSentEventsResponse({ adapter: log }); + return reconstructChat(backends().persistence, request, { + // A real app checks that the session's user may read this thread. + authorize: (threadId) => threadId.length > 0, + }); +} diff --git a/examples/tanstack-ai-demo/app/globals.css b/examples/tanstack-ai-demo/app/globals.css new file mode 100644 index 0000000..42325f9 --- /dev/null +++ b/examples/tanstack-ai-demo/app/globals.css @@ -0,0 +1,60 @@ +:root { + color-scheme: light dark; + font-family: system-ui, sans-serif; +} +body { + margin: 0; +} +main { + max-width: 720px; + margin: 0 auto; + padding: 24px 16px; +} +header { + display: flex; + align-items: center; + justify-content: space-between; +} +.status { + font-size: 12px; + opacity: 0.7; +} +.hint { + font-size: 14px; + opacity: 0.8; +} +.thread { + list-style: none; + padding: 0; + display: grid; + gap: 12px; +} +.thread li { + padding: 10px 12px; + border-radius: 10px; + background: color-mix(in srgb, currentColor 6%, transparent); +} +.thread li.user { + background: color-mix(in srgb, royalblue 18%, transparent); +} +.thread p { + margin: 0; + white-space: pre-wrap; +} +.tool { + font-size: 13px; + opacity: 0.7; +} +form { + display: flex; + gap: 8px; +} +input { + flex: 1; + padding: 10px; + font: inherit; +} +button { + padding: 10px 16px; + font: inherit; +} diff --git a/examples/tanstack-ai-demo/app/layout.tsx b/examples/tanstack-ai-demo/app/layout.tsx new file mode 100644 index 0000000..ea0aa4f --- /dev/null +++ b/examples/tanstack-ai-demo/app/layout.tsx @@ -0,0 +1,12 @@ +import type { ReactNode } from "react"; +import "./globals.css"; + +export const metadata = { title: "TanStack AI on Upstash" }; + +export default function RootLayout({ children }: { children: ReactNode }) { + return ( + + {children} + + ); +} diff --git a/examples/tanstack-ai-demo/app/page.tsx b/examples/tanstack-ai-demo/app/page.tsx new file mode 100644 index 0000000..90a19be --- /dev/null +++ b/examples/tanstack-ai-demo/app/page.tsx @@ -0,0 +1,90 @@ +"use client"; + +/** + * A chat whose state lives on the server, in Upstash. The page keeps nothing of its own: on mount + * `useChat` asks the route for the thread (transcript + any run still in flight) and tails that run. + * Reload mid-answer, or open the same URL on another device, and it picks up where the stream is. + */ +import { useEffect, useState } from "react"; +import { fetchServerSentEvents } from "@tanstack/ai-client"; +import { useChat } from "@tanstack/ai-react"; + +const connection = fetchServerSentEvents("/api/chat"); + +/** The thread lives in the URL (`?thread=`), so a reload or a shared link lands on the same one. */ +function useThreadId(): string | null { + const [threadId, setThreadId] = useState(null); + useEffect(() => { + const url = new URL(window.location.href); + let id = url.searchParams.get("thread"); + if (!id) { + id = `thread-${crypto.randomUUID()}`; + url.searchParams.set("thread", id); + window.history.replaceState(null, "", url); + } + setThreadId(id); + }, []); + return threadId; +} + +function Chat({ threadId }: { threadId: string }) { + const { messages, sendMessage, isLoading, connectionStatus } = useChat({ + threadId, + connection, + persistence: true, + }); + const [input, setInput] = useState(""); + + return ( +
+
+

TanStack AI on Upstash

+ {connectionStatus} +
+

+ Transcripts, runs and the stream live in Upstash Redis. Reload mid-answer, or open this URL + in another tab: the answer continues. Try remember: I like green tea, then ask + something in a new thread. +

+
    + {messages.map((message) => ( +
  1. + {message.parts.map((part, i) => + part.type === "text" ? ( +

    {part.content}

    + ) : part.type === "tool-call" ? ( +

    + ⚙ {part.name} +

    + ) : null, + )} +
  2. + ))} +
+
{ + event.preventDefault(); + const text = input.trim(); + if (!text || isLoading) return; + setInput(""); + void sendMessage(text); + }} + > + setInput(event.target.value)} + placeholder="Say something" + aria-label="Message" + /> + +
+
+ ); +} + +export default function Page() { + const threadId = useThreadId(); + return threadId ? : null; +} diff --git a/examples/tanstack-ai-demo/e2e/chat.e2e.ts b/examples/tanstack-ai-demo/e2e/chat.e2e.ts new file mode 100644 index 0000000..65fb3f5 --- /dev/null +++ b/examples/tanstack-ai-demo/e2e/chat.e2e.ts @@ -0,0 +1,159 @@ +/** + * End to end, across two server instances of the demo app sharing one Upstash Redis: things TanStack + * AI's in-memory backends cannot do, because each instance only sees its own memory. + */ +import { randomUUID } from "node:crypto"; +import { afterAll, describe, expect, inject, it } from "vitest"; +import { Redis } from "@upstash/redis"; +import { memoryScopeKey } from "@upstash/agentkit-tanstack-ai/memory"; +import { + assistantText, + newTurn, + readEvents, + reconstruct, + resume, + send, + textOf, + until, +} from "./client.js"; + +const instances = inject("instances"); +const demoPrefix = inject("demoPrefix"); +const users: string[] = []; + +/** How many assistant messages a stream carries: two producers show up as two starts, interleaved. */ +const messageStarts = (events: { chunk: { type: string } }[]) => + events.filter((e) => e.chunk.type === "TEXT_MESSAGE_START").length; + +const deltaCount = (n: number) => (events: { chunk: { type: string } }[]) => + events.filter((e) => e.chunk.type === "TEXT_MESSAGE_CONTENT").length >= n; + +/** The whole answer of a run, replayed from the start of its durable log. */ +async function fullAnswer(base: string, runId: string): Promise { + const res = await fetch(`${base}/api/chat?offset=-1&runId=${encodeURIComponent(runId)}`); + return textOf(await readEvents(res)); +} + +describe.skipIf(!instances)("demo app, two instances, one Upstash Redis", () => { + const { a, b } = instances ?? { a: "", b: "" }; + + afterAll(async () => { + const redis = Redis.fromEnv(); + const patterns = [ + `${demoPrefix}:*`, + ...users.map((u) => `agentkit:tanstackMemory:${memoryScopeKey({ threadId: "_", userId: u })}:*`), + ]; + for (const match of patterns) { + let cursor = "0"; + do { + const [next, keys] = await redis.scan(cursor, { match, count: 500 }); + if (keys.length) await redis.del(...keys); + cursor = String(next); + } while (cursor !== "0"); + } + }); + + it("a client that drops mid-answer on A resumes the same run on B, missing nothing", async () => { + const turn = newTurn("hello there"); + const controller = new AbortController(); + const before = await readEvents(await send(a, turn, controller.signal), deltaCount(5)); + controller.abort(); + const lastOffset = before.findLast((e) => e.id)?.id; + expect(lastOffset).toBeTruthy(); + + // The "reloaded" browser reconnects to the OTHER instance with the last offset it saw. + const after = await readEvents(await resume(b, lastOffset!)); + expect(after.at(-1)?.chunk.type).toBe("RUN_FINISHED"); + expect(after.some((e) => e.id === lastOffset)).toBe(false); // strictly after, no duplicates + + const stitched = textOf(before) + textOf(after); + expect(stitched).toBe(await fullAnswer(b, turn.runId)); + expect(stitched.startsWith("You said: hello there.")).toBe(true); + expect(stitched.endsWith("the very last word arrives.")).toBe(true); + }); + + it("a reload on B mid-answer finds the run still in flight on A, then the finished transcript", async () => { + const turn = newTurn("how are you"); + const controller = new AbortController(); + await readEvents(await send(a, turn, controller.signal), deltaCount(3)); + controller.abort(); // closed tab: the run keeps going on A + + const midway = await reconstruct(b, turn.threadId); + expect(midway.activeRun?.runId).toBe(turn.runId); + + const done = await until( + () => reconstruct(b, turn.threadId), + (t) => t.activeRun === null && assistantText(t).includes("arrives."), + ); + expect(done.activeRun).toBeNull(); + expect(assistantText(done)).toBe(await fullAnswer(b, turn.runId)); + expect(done.messages.filter((m) => m.role === "user")).toHaveLength(1); + }); + + it("a fact saved in one thread on A is recalled in a new thread on B", async () => { + const user = `e2e-user-${randomUUID()}`; + users.push(user); + const fact = `I like green tea ${randomUUID().slice(0, 8)}`; + + const saving = await readEvents(await send(a, newTurn(`remember: ${fact}`, { user }))); + expect(saving.some((e) => e.chunk.type === "TOOL_CALL_START")).toBe(true); + expect(textOf(saving)).toContain("Saved."); + + const asking = await readEvents(await send(b, newTurn("what do you know about me?", { user }))); + // Recalled into the prompt on B: the saved fact (and the captured message it came from). + expect(textOf(asking)).toMatch(/I remember: .*I like green tea/); + expect(textOf(asking)).toContain(fact); + + // Another user sees none of it. + const stranger = `e2e-user-${randomUUID()}`; + users.push(stranger); + const other = await readEvents(await send(b, newTurn("what do you know about me?", { user: stranger }))); + expect(textOf(other)).not.toContain(fact); + }); + + it("a run that outlives its producer lease is still produced once when re-POSTed elsewhere", async () => { + const turn = newTurn("long: a long one"); + const first = send(a, turn).then((res) => readEvents(res)); + // Mid-run (the answer takes ~6s) and past the 1s lease: without renewal, B would take the lock + // and start a second producer writing into the same log. + await new Promise((r) => setTimeout(r, 2_000)); + const second = await readEvents(await send(b, turn)); + const fromA = await first; + expect(textOf(second)).toBe(textOf(fromA)); + // A second producer would interleave a second answer into the same log. + expect(messageStarts(fromA)).toBe(1); + expect(messageStarts(second)).toBe(1); + const thread = await until( + () => reconstruct(a, turn.threadId), + (t) => t.activeRun === null, + ); + expect(thread.messages.filter((m) => m.role === "assistant")).toHaveLength(1); + }); + + it("rejects a malformed request with a 400", async () => { + const res = await fetch(`${a}/api/chat`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ messages: "nope" }), + }); + expect(res.status).toBe(400); + }); + + it("the same run POSTed to both instances at once is produced exactly once", async () => { + const turn = newTurn("only once please"); + const [fromA, fromB] = await Promise.all([ + send(a, turn).then((res) => readEvents(res)), + send(b, turn).then((res) => readEvents(res)), + ]); + // Both clients see the same single answer, tailed from one log. + expect(textOf(fromA)).toBe(textOf(fromB)); + expect(messageStarts(fromA)).toBe(1); + expect(messageStarts(fromB)).toBe(1); + + const thread = await until( + () => reconstruct(a, turn.threadId), + (t) => t.activeRun === null, + ); + expect(thread.messages.filter((m) => m.role === "assistant")).toHaveLength(1); + }); +}); diff --git a/examples/tanstack-ai-demo/e2e/client.ts b/examples/tanstack-ai-demo/e2e/client.ts new file mode 100644 index 0000000..2e0a492 --- /dev/null +++ b/examples/tanstack-ai-demo/e2e/client.ts @@ -0,0 +1,118 @@ +/** A small AG-UI-over-SSE client for the E2E suite: what `useChat` does, without a browser. */ +import { randomUUID } from "node:crypto"; + +export interface Event { + /** The SSE `id:` — the resume offset of this chunk in the run's durable log. */ + id?: string; + chunk: { type: string; delta?: string; [key: string]: unknown }; +} + +export interface Turn { + threadId: string; + runId: string; + text: string; + user?: string; +} + +export function newTurn(text: string, init: Partial = {}): Turn { + return { threadId: `thread-${randomUUID()}`, runId: `run-${randomUUID()}`, text, ...init }; +} + +/** POST a user message; returns the SSE response (the run keeps going if it is dropped). */ +export function send(base: string, turn: Turn, signal?: AbortSignal): Promise { + return fetch(`${base}/api/chat`, { + method: "POST", + headers: { + "content-type": "application/json", + ...(turn.user ? { "x-demo-user": turn.user } : {}), + }, + body: JSON.stringify({ + threadId: turn.threadId, + runId: turn.runId, + messages: [{ id: `msg-${randomUUID()}`, role: "user", content: turn.text }], + tools: [], + context: [], + }), + ...(signal ? { signal } : {}), + }); +} + +/** Resume a run from an offset, the way a reloaded browser does (`Last-Event-ID`). */ +export function resume(base: string, offset: string, signal?: AbortSignal): Promise { + return fetch(`${base}/api/chat`, { + headers: { "last-event-id": offset }, + ...(signal ? { signal } : {}), + }); +} + +/** What `useChat` fetches on mount: the transcript plus the run still in flight, if any. */ +export async function reconstruct( + base: string, + threadId: string, +): Promise<{ + messages: Array<{ role: string; parts: Array<{ type: string; content?: string }> }>; + activeRun: { runId: string } | null; +}> { + const res = await fetch(`${base}/api/chat?threadId=${encodeURIComponent(threadId)}`); + if (!res.ok) throw new Error(`reconstruct failed: ${res.status} ${await res.text()}`); + return res.json(); +} + +/** + * Read SSE events from a response. Stops at the end of the stream, or after `stopAfter` returns + * true (then the caller aborts the request, like a closed tab). + */ +export async function readEvents( + res: Response, + stopAfter?: (events: Event[]) => boolean, +): Promise { + if (!res.ok || !res.body) throw new Error(`stream failed: ${res.status} ${await res.text()}`); + const events: Event[] = []; + const reader = res.body.pipeThrough(new TextDecoderStream()).getReader(); + let buffer = ""; + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + buffer += value; + let sep: number; + while ((sep = buffer.indexOf("\n\n")) !== -1) { + const raw = buffer.slice(0, sep); + buffer = buffer.slice(sep + 2); + let id: string | undefined; + let data = ""; + for (const line of raw.split("\n")) { + if (line.startsWith("id:")) id = line.slice(3).trim(); + else if (line.startsWith("data:")) data += line.slice(5).trim(); + } + if (!data) continue; + events.push({ ...(id ? { id } : {}), chunk: JSON.parse(data) }); + if (stopAfter?.(events)) { + await reader.cancel().catch(() => undefined); + return events; + } + } + } + return events; +} + +export const textOf = (events: Event[]) => + events + .filter((e) => e.chunk.type === "TEXT_MESSAGE_CONTENT") + .map((e) => e.chunk.delta ?? "") + .join(""); + +export const assistantText = (thread: Awaited>) => + thread.messages + .filter((m) => m.role === "assistant") + .map((m) => m.parts.filter((p) => p.type === "text").map((p) => p.content ?? "").join("")) + .join("\n"); + +export async function until(read: () => Promise, ok: (v: T) => boolean, ms = 20_000): Promise { + const deadline = Date.now() + ms; + let value = await read(); + while (!ok(value) && Date.now() < deadline) { + await new Promise((r) => setTimeout(r, 200)); + value = await read(); + } + return value; +} diff --git a/examples/tanstack-ai-demo/e2e/servers.ts b/examples/tanstack-ai-demo/e2e/servers.ts new file mode 100644 index 0000000..29fca59 --- /dev/null +++ b/examples/tanstack-ai-demo/e2e/servers.ts @@ -0,0 +1,106 @@ +/** + * Global setup: build the app if needed, then start TWO production servers (`next start`) on free + * ports against the same Upstash Redis, with the scripted model. Tests talk to both, so every + * scenario can start something on one instance and finish it on the other. + * + * Skipped (no servers, tests skip too) without UPSTASH_REDIS_REST_URL / _TOKEN. + */ +import { spawn, spawnSync, type ChildProcess } from "node:child_process"; +import { existsSync } from "node:fs"; +import { createServer } from "node:net"; +import { randomUUID } from "node:crypto"; +import { fileURLToPath } from "node:url"; +import { config } from "dotenv"; +import type { GlobalSetupContext } from "vitest/node"; + +const appDir = fileURLToPath(new URL("..", import.meta.url)); +const nextBin = `${appDir}node_modules/.bin/next`; +config({ path: fileURLToPath(new URL("../../../.env", import.meta.url)) }); + +declare module "vitest" { + export interface ProvidedContext { + instances: { a: string; b: string } | null; + demoPrefix: string; + } +} + +function freePort(): Promise { + return new Promise((resolve, reject) => { + const server = createServer(); + server.listen(0, "127.0.0.1", () => { + const { port } = server.address() as { port: number }; + server.close(() => resolve(port)); + }); + server.on("error", reject); + }); +} + +async function waitUntilUp(url: string, child: ChildProcess, log: () => string): Promise { + const deadline = Date.now() + 60_000; + while (Date.now() < deadline) { + if (child.exitCode !== null) throw new Error(`server exited early:\n${log()}`); + try { + if ((await fetch(url)).ok) return; + } catch { + /* not listening yet */ + } + await new Promise((r) => setTimeout(r, 250)); + } + throw new Error(`server at ${url} did not come up:\n${log()}`); +} + +export default async function setup({ provide }: GlobalSetupContext) { + const demoPrefix = `tanstack-ai-demo-e2e:${randomUUID()}`; + provide("demoPrefix", demoPrefix); + if (!process.env.UPSTASH_REDIS_REST_URL || !process.env.UPSTASH_REDIS_REST_TOKEN) { + provide("instances", null); + return; + } + + if (!existsSync(`${appDir}.next/BUILD_ID`)) { + const build = spawnSync(nextBin, ["build"], { cwd: appDir, stdio: "inherit" }); + if (build.status !== 0) throw new Error("next build failed"); + } + + const env = { + ...process.env, + AGENTKIT_MOCK_MODEL: "1", + MOCK_WORD_DELAY_MS: process.env.MOCK_WORD_DELAY_MS ?? "60", + DEMO_PREFIX: demoPrefix, + // Shorter than one mock answer (~2s), so a run outlives its first lease and must renew it. + PRODUCER_LEASE_MS: "1000", + NODE_ENV: "production" as const, + }; + const children: ChildProcess[] = []; + const urls: string[] = []; + for (let i = 0; i < 2; i++) { + const port = await freePort(); + // `next` itself (not `npx`, which would leave it orphaned), in its own process group so + // teardown can stop it and anything it spawned. + const child = spawn(nextBin, ["start", "-p", String(port), "-H", "127.0.0.1"], { + cwd: appDir, + env, + stdio: ["ignore", "pipe", "pipe"], + detached: true, + }); + let output = ""; + child.stdout?.on("data", (d) => (output += d)); + child.stderr?.on("data", (d) => (output += d)); + children.push(child); + const url = `http://127.0.0.1:${port}`; + await waitUntilUp(`${url}/`, child, () => output); + urls.push(url); + } + provide("instances", { a: urls[0]!, b: urls[1]! }); + + return async () => { + for (const child of children) { + if (child.pid === undefined || child.exitCode !== null) continue; + try { + process.kill(-child.pid, "SIGTERM"); + } catch { + /* already gone */ + } + } + }; +} diff --git a/examples/tanstack-ai-demo/e2e/vitest.config.ts b/examples/tanstack-ai-demo/e2e/vitest.config.ts new file mode 100644 index 0000000..3baa330 --- /dev/null +++ b/examples/tanstack-ai-demo/e2e/vitest.config.ts @@ -0,0 +1,13 @@ +import { defineConfig } from "vitest/config"; + +// The E2E suite: two production servers of this app, one Upstash Redis, a scripted model. +// Files are `*.e2e.ts` so the repo's unit-test run never picks them up. +export default defineConfig({ + test: { + include: ["e2e/**/*.e2e.ts"], + globalSetup: ["e2e/servers.ts"], + testTimeout: 60_000, + hookTimeout: 180_000, + fileParallelism: false, + }, +}); diff --git a/examples/tanstack-ai-demo/lib/backends.ts b/examples/tanstack-ai-demo/lib/backends.ts new file mode 100644 index 0000000..3f9c570 --- /dev/null +++ b/examples/tanstack-ai-demo/lib/backends.ts @@ -0,0 +1,49 @@ +/** + * The Upstash backends this app runs on. Every one of them is shared state in Upstash Redis, which + * is the point: any server instance can serve any request for any thread. + * + * Created on first use, not at import, so `next build` (which loads route modules) needs no + * credentials. + */ +import { Redis } from "@upstash/redis"; +import { RedisLock, upstashStream } from "@upstash/agentkit-tanstack-ai"; +import { upstashPersistence } from "@upstash/agentkit-tanstack-ai/persistence"; +import { upstashMemory } from "@upstash/agentkit-tanstack-ai/memory"; + +/** `DEMO_PREFIX` isolates a run's keys (the E2E suite sets one per run). */ +const prefix = () => process.env.DEMO_PREFIX ?? "tanstack-ai-demo"; + +function create() { + const redis = Redis.fromEnv(); + return { + redis, + /** Transcripts, run records and interrupts — what `reconstructChat` rebuilds a reload from. */ + persistence: upstashPersistence({ redis, prefix: `${prefix()}:state` }), + /** Long-term memory, per user across threads. */ + memory: upstashMemory({ redis }), + /** + * "Only one producer per run", across instances: a duplicate POST for the same run (a client + * retry landing on another server) must not start a second model run into the same log. + */ + // A short lease, renewed while the run lasts: a crashed instance frees the run within seconds. + // (`PRODUCER_LEASE_MS` lets the E2E suite make it shorter than a run, to prove the renewal.) + producerLock: new RedisLock({ + redis, + prefix: `${prefix()}:producer`, + leaseMs: Number(process.env.PRODUCER_LEASE_MS ?? 15_000), + }), + }; +} + +let instance: ReturnType | undefined; +export const backends = () => (instance ??= create()); + +/** The durable delivery log of one run, on Redis Streams. */ +export function runLog(source: Request | { runId: string; offset?: string | null }) { + return upstashStream(source, { + redis: backends().redis, + prefix: `${prefix()}:stream`, + // A reader that joins a run the moment it is started may briefly beat the first chunk. + firstChunkDeadlineMs: 10_000, + }); +} diff --git a/examples/tanstack-ai-demo/lib/model.ts b/examples/tanstack-ai-demo/lib/model.ts new file mode 100644 index 0000000..a054c66 --- /dev/null +++ b/examples/tanstack-ai-demo/lib/model.ts @@ -0,0 +1,102 @@ +/** + * The chat model. With `OPENAI_API_KEY` set it is OpenAI; with `AGENTKIT_MOCK_MODEL=1` (CI, the E2E + * suite) it is a scripted stand-in, so the whole app runs deterministically with no provider: + * + * - A message starting with `remember:` makes the model call the `save_memory` tool with the rest. + * - Otherwise it answers `You said: .`, then `I remember: …` with any memories that were + * recalled into the system prompt, then a short filler (three times over for `long:` messages) — streamed one word at a time + * (`MOCK_WORD_DELAY_MS`, default 40) so a client can disconnect mid-answer. + */ +import { openaiText } from "@tanstack/ai-openai"; +import type { StreamChunk } from "@tanstack/ai"; + +type Adapter = ReturnType; + +const FILLER = + "This answer is streamed word by word so a reload in the middle of it can be resumed from the " + + "durable log on any server instance until the very last word arrives."; + +interface CallOptions { + messages: Array<{ role: string; content: unknown }>; + systemPrompts?: unknown; +} + +const textOf = (content: unknown): string => + typeof content === "string" + ? content + : Array.isArray(content) + ? content.map((p) => (p && typeof p === "object" && "content" in p ? String(p.content) : "")).join("") + : ""; + +const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)); + +function mockModel(): Adapter { + const delay = Number(process.env.MOCK_WORD_DELAY_MS ?? 40); + let calls = 0; + const adapter = { + kind: "text", + name: "mock", + model: "mock-1", + async *chatStream(options: CallOptions): AsyncIterable { + calls++; + const timestamp = Date.now(); + const last = options.messages.at(-1); + const lastUser = [...options.messages].reverse().find((m) => m.role === "user"); + const said = textOf(lastUser?.content).trim(); + + // First step of a "remember:" turn: call save_memory. The tool result comes back as a `tool` + // message, and the next step (below) answers in text. + if (said.toLowerCase().startsWith("remember:") && last?.role !== "tool") { + const id = `call-${calls}-${timestamp}`; + const args = { text: said.slice("remember:".length).trim() }; + yield { type: "TOOL_CALL_START", toolCallId: id, toolCallName: "save_memory", toolName: "save_memory", timestamp } as unknown as StreamChunk; + yield { type: "TOOL_CALL_ARGS", toolCallId: id, delta: JSON.stringify(args), timestamp } as unknown as StreamChunk; + yield { type: "TOOL_CALL_END", toolCallId: id, toolCallName: "save_memory", toolName: "save_memory", input: args, timestamp } as unknown as StreamChunk; + yield { type: "RUN_FINISHED", finishReason: "tool_calls", timestamp } as unknown as StreamChunk; + return; + } + + // Memories recalled into the system prompt, one "- ()" line each. + const prompts = Array.isArray(options.systemPrompts) ? options.systemPrompts : []; + const recalled = prompts + .map((p) => (typeof p === "string" ? p : textOf((p as { content?: unknown })?.content))) + .join("\n") + .split("\n") + .filter((line) => line.startsWith("- ")) + .map((line) => line.slice(2).replace(/ \((you saved this|the user said this)\)$/, "")) + .sort(); + const answer = + last?.role === "tool" + ? "Saved." + : [ + `You said: ${said}.`, + recalled.length ? `I remember: ${recalled.join("; ")}.` : "", + // "long:" triples the filler, for a run that must outlast a short producer lease. + ...(said.toLowerCase().startsWith("long:") ? [FILLER, FILLER, FILLER] : [FILLER]), + ] + .filter(Boolean) + .join(" "); + + const messageId = `msg-${calls}-${timestamp}`; + yield { type: "TEXT_MESSAGE_START", messageId, role: "assistant", timestamp } as unknown as StreamChunk; + const words = answer.split(" "); + for (const [i, word] of words.entries()) { + if (delay > 0) await sleep(delay); + yield { type: "TEXT_MESSAGE_CONTENT", messageId, delta: i === 0 ? word : ` ${word}`, timestamp: Date.now() } as unknown as StreamChunk; + } + yield { type: "TEXT_MESSAGE_END", messageId, timestamp: Date.now() } as unknown as StreamChunk; + yield { type: "RUN_FINISHED", finishReason: "stop", timestamp: Date.now() } as unknown as StreamChunk; + }, + structuredOutput: async () => { + throw new Error("The mock model does not do structured output."); + }, + }; + // The mock implements the same streaming surface chat() drives; it is typed as the real adapter + // so the route code is identical in both modes. + return adapter as unknown as Adapter; +} + +export function chatModel(): Adapter { + if (process.env.AGENTKIT_MOCK_MODEL === "1") return mockModel(); + return openaiText((process.env.OPENAI_MODEL ?? "gpt-5.4-mini") as Parameters[0]); +} diff --git a/examples/tanstack-ai-demo/next-env.d.ts b/examples/tanstack-ai-demo/next-env.d.ts new file mode 100644 index 0000000..9edff1c --- /dev/null +++ b/examples/tanstack-ai-demo/next-env.d.ts @@ -0,0 +1,6 @@ +/// +/// +import "./.next/types/routes.d.ts"; + +// NOTE: This file should not be edited +// see https://nextjs.org/docs/app/api-reference/config/typescript for more information. diff --git a/examples/tanstack-ai-demo/next.config.ts b/examples/tanstack-ai-demo/next.config.ts new file mode 100644 index 0000000..cb651cd --- /dev/null +++ b/examples/tanstack-ai-demo/next.config.ts @@ -0,0 +1,5 @@ +import type { NextConfig } from "next"; + +const nextConfig: NextConfig = {}; + +export default nextConfig; diff --git a/examples/tanstack-ai-demo/package.json b/examples/tanstack-ai-demo/package.json new file mode 100644 index 0000000..e3b1db1 --- /dev/null +++ b/examples/tanstack-ai-demo/package.json @@ -0,0 +1,34 @@ +{ + "name": "tanstack-ai-demo", + "version": "0.0.0", + "private": true, + "type": "module", + "scripts": { + "dev": "next dev", + "build": "next build", + "start": "next start", + "e2e": "vitest run --config e2e/vitest.config.ts" + }, + "dependencies": { + "@tanstack/ai": "0.63.0", + "@tanstack/ai-client": "0.36.0", + "@tanstack/ai-memory": "0.2.8", + "@tanstack/ai-openai": "0.25.1", + "@tanstack/ai-persistence": "0.7.1", + "@tanstack/ai-react": "0.29.3", + "@upstash/agentkit-tanstack-ai": "workspace:*", + "@upstash/redis": "^1.38.4", + "next": "16.2.9", + "react": "19.2.6", + "react-dom": "19.2.6", + "zod": "^4" + }, + "devDependencies": { + "@types/node": "^20", + "@types/react": "19.2.15", + "@types/react-dom": "19.2.3", + "dotenv": "^16.4.5", + "typescript": "^5", + "vitest": "^2.0.0" + } +} diff --git a/examples/tanstack-ai-demo/tsconfig.json b/examples/tanstack-ai-demo/tsconfig.json new file mode 100644 index 0000000..803331c --- /dev/null +++ b/examples/tanstack-ai-demo/tsconfig.json @@ -0,0 +1,41 @@ +{ + "compilerOptions": { + "target": "ES2022", + "lib": [ + "dom", + "dom.iterable", + "esnext" + ], + "allowJs": true, + "skipLibCheck": true, + "strict": true, + "noEmit": true, + "esModuleInterop": true, + "module": "esnext", + "moduleResolution": "bundler", + "resolveJsonModule": true, + "isolatedModules": true, + "jsx": "react-jsx", + "incremental": true, + "plugins": [ + { + "name": "next" + } + ], + "paths": { + "@/*": [ + "./*" + ] + } + }, + "include": [ + "next-env.d.ts", + "**/*.ts", + "**/*.tsx", + ".next/types/**/*.ts", + ".next/dev/types/**/*.ts" + ], + "exclude": [ + "node_modules" + ] +} diff --git a/packages/eve/src/memory/documents.ts b/packages/eve/src/memory/documents.ts index f6bfd5b..c562d30 100644 --- a/packages/eve/src/memory/documents.ts +++ b/packages/eve/src/memory/documents.ts @@ -135,7 +135,7 @@ const CONTENT_MARKER = "eve-memory-document-v1:"; * caller turns the second case into eve's `MemoryDocumentConflictError`. Returning the *current* * version rather than a bare `0` keeps the failure debuggable. */ -const CAS_SCRIPT = ` +const CAS_SCRIPT = `#!lua flags=allow-key-locking local current = redis.call('HGET', KEYS[1], 'version') if current == false then current = '' end if current ~= ARGV[2] then return {0, current} end diff --git a/packages/tanstack-ai/LICENSE b/packages/tanstack-ai/LICENSE new file mode 100644 index 0000000..e748b2f --- /dev/null +++ b/packages/tanstack-ai/LICENSE @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2026 Upstash, Inc. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/packages/tanstack-ai/README.md b/packages/tanstack-ai/README.md new file mode 100644 index 0000000..a410367 --- /dev/null +++ b/packages/tanstack-ai/README.md @@ -0,0 +1,175 @@ +# @upstash/agentkit-tanstack-ai + +Production backends for [TanStack AI](https://tanstack.com/ai), all on one +[Upstash Redis](https://upstash.com/). TanStack AI defines the contracts for chat persistence, +resumable streaming, distributed locks and memory, and ships in-memory implementations that only +work inside one process. This package implements them on Redis, so they hold across serverless +instances, reloads and devices. + +| TanStack AI seam | This package | What you get | +| --- | --- | --- | +| `withPersistence()` / `withGenerationPersistence()` stores | `upstashPersistence()` | Transcripts, runs, human-in-the-loop interrupts, metadata, generation jobs, artifacts and (with Upstash Blob) their bytes | +| `StreamDurability` | `upstashStream()` | Reload mid-answer or open the thread on another device and pick up where the stream is | +| `withLocks()` `LockStore` | `upstashLocks()` | Distributed locks for TanStack AI middleware, e.g. so concurrent requests never create duplicate sandboxes | +| `memoryMiddleware()` adapter | `upstashMemory()` | Long-term memory ranked in Redis Search (BM25, typo-tolerant), no per-turn full scan | +| chat middleware | `toolCache()`, `rateLimit()` | Skip repeated deterministic tool calls; throttle users before the model runs | +| tools | `createSearchTools()` | `search` / `aggregate` / `count` over your own documents for RAG | + +`upstashPersistence()` and `upstashMemory()` are checked against TanStack AI's own conformance +suites (`runPersistenceConformance`, `runMemoryAdapterContract`). + +## Install + +```bash +npm install @upstash/agentkit-tanstack-ai @tanstack/ai +# plus, for the features you use: +npm install @tanstack/ai-persistence @tanstack/ai-memory +``` + +Persistence and memory have their own entry points, `@upstash/agentkit-tanstack-ai/persistence` and +`@upstash/agentkit-tanstack-ai/memory`, so the root entry never needs those optional packages. + +Set `UPSTASH_REDIS_REST_URL` and `UPSTASH_REDIS_REST_TOKEN`; every factory defaults to +`Redis.fromEnv()` and accepts `redis` to pass a client explicitly. + +## Chat persistence + +```ts +import { chat } from "@tanstack/ai"; +import { withPersistence } from "@tanstack/ai-persistence"; +import { upstashPersistence } from "@upstash/agentkit-tanstack-ai/persistence"; + +const persistence = upstashPersistence({ + messagesTtlSeconds: 60 * 60 * 24 * 30, // optional: expire idle transcripts +}); + +chat({ adapter, messages, threadId, middleware: [withPersistence(persistence)] }); +``` + +Stores `messages`, `runs`, `interrupts` and `metadata` (the full `ChatPersistence` shape). Runs are +indexed by thread, parent run and detach time, so reconnect (`findActiveRun`), subagent cards +(`listByParentRun`) and the sandbox reaper (`listReclaimable`) are index reads. Every mutation is a +single command or one Lua script, so concurrent instances cannot interleave a write. + +It also stores `generationRuns` and `artifacts`, for one-shot generation jobs (image, video, speech) +saved by `withGenerationPersistence()`. Pass an [Upstash Blob](https://upstash.com/docs/blob) bucket to +add a `blobs` store for the generated bytes: + +```ts +import { Bucket } from "@upstash/blob"; + +const persistence = upstashPersistence({ bucket: Bucket.fromEnv() }); +``` + +The bytes go to the bucket and each blob's record (size, content type, custom metadata) to Redis, so +`head` and `list` are Redis reads and byte ranges are served with an HTTP `Range` request. + +## Resumable streams + +```ts +import { chat, toServerSentEventsResponse } from "@tanstack/ai"; +import { upstashStream } from "@upstash/agentkit-tanstack-ai"; + +export async function POST(request: Request) { + const stream = chat({ adapter, messages, threadId }); + return toServerSentEventsResponse(stream, { durability: { adapter: upstashStream(request) } }); +} +``` + +Each chunk is appended to a Redis Stream before it is delivered. A client that reconnects with +`Last-Event-ID` (or `?offset`) replays everything after it and keeps tailing the live run, on any +instance. Without a `Request`, use `upstashStream({ runId, offset })`. + +| Option | Default | | +| --- | --- | --- | +| `ttlSeconds` | `86400` | How long a run stays resumable after its last chunk | +| `pollIntervalMs` | `150` | Tail poll interval (the REST API keeps no blocking reads) | +| `firstChunkDeadlineMs` | `2000` | How long a from-start join waits for a run that has not produced yet | + +## Distributed locks + +```ts +import { withLocks } from "@tanstack/ai/locks"; +import { upstashLocks } from "@upstash/agentkit-tanstack-ai"; + +chat({ adapter, messages, middleware: [withLocks(upstashLocks()), withSandbox(sandbox)] }); +``` + +`withLocks` doesn't lock anything by itself: it hands the lock store to later middleware, which +take a lock around the one step they must not run twice. `withSandbox` does this so two concurrent +requests for a thread don't both create a sandbox, and your own middleware can do the same via +`getLocks(ctx)`. It does not serialize whole chat turns. TanStack's built-in `InMemoryLockStore` only +coordinates inside one process; `upstashLocks()` works across instances. + +Each lock is a lease (`leaseMs`, default 30 s) renewed while the section runs; the section's `signal` +aborts if the lease is lost. Built on `RedisLock`, exported from this package too, which also exposes a +fencing token (`EventLog`, the Redis Streams log under `upstashStream`, is exported as well). + +## Long-term memory + +```ts +import { memoryMiddleware } from "@tanstack/ai-memory"; +import { upstashMemory } from "@upstash/agentkit-tanstack-ai/memory"; + +chat({ + adapter, + messages, + middleware: [ + memoryMiddleware({ + adapter: upstashMemory(), + // Derive these server-side from the session, never from the request body. + scope: (ctx) => ({ threadId: ctx.threadId, userId: session.userId }), + }), + ], +}); +``` + +- Recall runs before the model: the top matches for the user's message are injected as a system + prompt block, each labelled with its source (`you saved this` / `the user said this`). +- The model gets a `save_memory` tool for durable facts (`saveTool: false` to turn it off). +- Each save waits for the index (`waitForIndexing`, default on), so a memory is recallable on the very + next turn. +- Each turn's user message is captured (`captureUserMessages: false` for model-curated only). +- Memory is per user across threads by default (`scopeBy: "thread"` for per-conversation); + `tenantId` and `namespace` always partition. + +## Tool caching and rate limiting + +```ts +import { rateLimit, Ratelimit, toolCache } from "@upstash/agentkit-tanstack-ai"; + +chat({ + adapter, + messages, + tools: [getWeather, sendEmail], + middleware: [ + rateLimit({ limiter: Ratelimit.slidingWindow(10, "60 s"), identifier: userId }), + toolCache({ tools: ["get_weather"], userId, ttlSeconds: 600 }), // allowlist deterministic tools only + ], +}); +``` + +`rateLimit` fails the run with `RateLimitExceededError` before the model is called. To answer with +an HTTP 429 instead, call `createRateLimit({ limiter }).limit(userId)` in the route before `chat()`. + +## Search tools (RAG) + +```ts +import { s } from "@upstash/redis"; +import { createSearchTools } from "@upstash/agentkit-tanstack-ai"; + +const tools = createSearchTools({ + indexName: "products", + schema: s.object({ name: s.string(), price: s.number(), category: s.string().noTokenize() }), +}); +chat({ adapter, messages, tools }); +``` + +## Telemetry + +Reports the package name and version as a header on your redis client's requests. Opt out with +`enableTelemetry: false` on any factory, on the redis client, or with `UPSTASH_DISABLE_TELEMETRY`. + +## License + +MIT diff --git a/packages/tanstack-ai/package.json b/packages/tanstack-ai/package.json new file mode 100644 index 0000000..d5baf00 --- /dev/null +++ b/packages/tanstack-ai/package.json @@ -0,0 +1,84 @@ +{ + "name": "@upstash/agentkit-tanstack-ai", + "version": "0.0.0", + "description": "Upstash AgentKit backends for TanStack AI: chat persistence, resumable streams, distributed locks, long-term memory, tool caching, rate limiting, and Redis-Search tools — all on one Upstash Redis.", + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/upstash/agentkit.git", + "directory": "packages/tanstack-ai" + }, + "homepage": "https://github.com/upstash/agentkit/tree/main/packages/tanstack-ai", + "bugs": { + "url": "https://github.com/upstash/agentkit/issues" + }, + "type": "module", + "main": "./dist/index.js", + "module": "./dist/index.js", + "types": "./dist/index.d.ts", + "exports": { + ".": { + "types": "./dist/index.d.ts", + "import": "./dist/index.js" + }, + "./persistence": { + "types": "./dist/persistence.d.ts", + "import": "./dist/persistence.js" + }, + "./memory": { + "types": "./dist/memory.d.ts", + "import": "./dist/memory.js" + } + }, + "files": [ + "dist", + "README.md" + ], + "scripts": { + "build": "tsup", + "dev": "tsup --watch", + "clean": "rm -rf dist", + "typecheck": "tsc --noEmit" + }, + "keywords": [ + "upstash", + "tanstack", + "tanstack-ai", + "ai", + "agent", + "persistence", + "resumable-stream", + "memory", + "lock", + "rag" + ], + "dependencies": { + "@upstash/agentkit-sdk": "workspace:*", + "@upstash/redis": "^1.38.4", + "zod": "^3.23.8 || ^4" + }, + "devDependencies": { + "@tanstack/ai": "0.63.0", + "@tanstack/ai-memory": "0.2.8", + "@tanstack/ai-persistence": "0.7.1", + "dotenv": "^16.4.5", + "@upstash/blob": "0.0.7" + }, + "peerDependencies": { + "@tanstack/ai": ">=0.63.0", + "@tanstack/ai-memory": ">=0.2.8", + "@tanstack/ai-persistence": ">=0.7.1", + "@upstash/blob": ">=0.0.7" + }, + "peerDependenciesMeta": { + "@tanstack/ai-memory": { + "optional": true + }, + "@tanstack/ai-persistence": { + "optional": true + }, + "@upstash/blob": { + "optional": true + } + } +} diff --git a/packages/tanstack-ai/src/index.ts b/packages/tanstack-ai/src/index.ts new file mode 100644 index 0000000..231b2d8 --- /dev/null +++ b/packages/tanstack-ai/src/index.ts @@ -0,0 +1,36 @@ +/** + * Root entry: everything that needs only `@tanstack/ai`. The features built on TanStack's optional + * packages have their own entry points, so their types never reach this one: + * `@upstash/agentkit-tanstack-ai/persistence` (needs `@tanstack/ai-persistence`) and + * `@upstash/agentkit-tanstack-ai/memory` (needs `@tanstack/ai-memory`). + */ + +// Resumable delivery (`StreamDurability`) on Redis Streams — reload/reconnect/second device. +export { upstashStream } from "./stream/stream.js"; +export type { UpstashStreamConfig, UpstashStreamInit } from "./stream/stream.js"; + +// The Redis primitives under `upstashLocks()` and `upstashStream()`, usable directly: a lease lock with +// fencing tokens, and a resumable append-only event log on Redis Streams. +export { RedisLock, LockAcquireTimeoutError, LockLostError } from "./locks/redis-lock.js"; +export type { RedisLockConfig, LockLease } from "./locks/redis-lock.js"; +export { EventLog, EventLogClosedError } from "./stream/event-log.js"; +export type { EventLogConfig, LogEntry } from "./stream/event-log.js"; + +// Distributed `LockStore` for `withLocks()`. +export { upstashLocks } from "./locks/locks.js"; +export type { UpstashLocksConfig } from "./locks/locks.js"; + +// Chat middlewares: tool-result caching and per-run rate limiting. +export { toolCache, rateLimit, RateLimitExceededError } from "./middleware/middleware.js"; +export type { + ToolCacheMiddlewareConfig, + RateLimitMiddlewareConfig, +} from "./middleware/middleware.js"; + +// Schema-driven Redis Search tools (search / aggregate / count) for RAG. +export { createSearchTools } from "./search/search-tools.js"; +export type { CreateSearchToolsConfig } from "./search/search-tools.js"; + +// Rate limiting primitives, re-exported so users never import `@upstash/ratelimit` directly. +export { createRateLimit, Ratelimit } from "@upstash/agentkit-sdk"; +export type { RateLimitConfig, Duration } from "@upstash/agentkit-sdk"; diff --git a/packages/tanstack-ai/src/integration.test.ts b/packages/tanstack-ai/src/integration.test.ts new file mode 100644 index 0000000..473a0ae --- /dev/null +++ b/packages/tanstack-ai/src/integration.test.ts @@ -0,0 +1,57 @@ +import { randomUUID } from "node:crypto"; +import { afterAll, describe, expect, it } from "vitest"; +import { chat, replayRunStream, toServerSentEventsResponse } from "@tanstack/ai"; +import type { StreamChunk } from "@tanstack/ai"; +import { withPersistence } from "@tanstack/ai-persistence"; +import { upstashPersistence } from "./persistence/persistence.js"; +import { upstashStream } from "./stream/stream.js"; +import { scriptedAdapter } from "./testing/test-adapter.js"; +import { cleanupKeys, hasRedisCreds, testRedis, uniquePrefix } from "./testing/test-support.js"; + +// The pieces together, the way an app route wires them: persistence middleware on `chat()`, the +// durable stream on the SSE response, and a second "instance" reading both back. +describe.skipIf(!hasRedisCreds)("persistence + resumable stream in a real chat() route", () => { + const redis = testRedis(); + const prefix = uniquePrefix("tsint"); + afterAll(async () => { + await cleanupKeys(redis, prefix); + }); + + it("a finished run leaves a reloadable transcript, a completed run record, and a replayable stream", async () => { + const threadId = `thread-${randomUUID()}`; + const runId = randomUUID(); + const persistence = upstashPersistence({ redis, prefix: `${prefix}:state` }); + const adapter = scriptedAdapter([{ text: "Hi Arda, how can I help?" }]); + + const stream = chat({ + adapter: adapter as never, + threadId, + runId, + messages: [{ role: "user", content: "hello" }], + middleware: [withPersistence(persistence)], + } as never); + const response = toServerSentEventsResponse(stream as AsyncIterable, { + durability: { adapter: upstashStream({ runId }, { redis, prefix: `${prefix}:stream` }) }, + }); + expect(await response.text()).toContain("how can I help"); + + // Another instance: fresh store objects over the same Redis. + const reloaded = upstashPersistence({ redis: testRedis(), prefix: `${prefix}:state` }); + const transcript = await reloaded.stores.messages.loadThread(threadId); + expect(JSON.stringify(transcript)).toContain("hello"); + expect(JSON.stringify(transcript)).toContain("how can I help"); + + const runs = await reloaded.stores.runs.listByThread!(threadId); + expect(runs.length).toBeGreaterThan(0); + expect(runs.every((r) => r.status === "completed")).toBe(true); + expect(await reloaded.stores.runs.findActiveRun(threadId)).toBeNull(); + + const replayed: string[] = []; + const joiner = upstashStream({ runId }, { redis: testRedis(), prefix: `${prefix}:stream` }); + for await (const c of replayRunStream(joiner)) { + const delta = (c as { delta?: string }).delta; + if (delta) replayed.push(delta); + } + expect(replayed.join("")).toBe("Hi Arda, how can I help?"); + }); +}); diff --git a/packages/tanstack-ai/src/locks/locks.test.ts b/packages/tanstack-ai/src/locks/locks.test.ts new file mode 100644 index 0000000..fcf8205 --- /dev/null +++ b/packages/tanstack-ai/src/locks/locks.test.ts @@ -0,0 +1,56 @@ +import { describe, expect, it } from "vitest"; +import { afterAll } from "vitest"; +import { upstashLocks } from "./locks.js"; +import { cleanupKeys, hasRedisCreds, testRedis, uniquePrefix } from "../testing/test-support.js"; + +const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)); + +describe.skipIf(!hasRedisCreds)("upstashLocks (live Redis)", () => { + const redis = testRedis(); + const prefix = uniquePrefix("tslock"); + afterAll(async () => { + await cleanupKeys(redis, prefix); + }); + + it("two LockStores (two instances) never run the same key's section concurrently", async () => { + const a = upstashLocks({ redis, prefix, retryDelayMs: 20 }); + const b = upstashLocks({ redis: testRedis(), prefix, retryDelayMs: 20 }); + let inside = 0; + let max = 0; + const section = async () => { + inside++; + max = Math.max(max, inside); + await sleep(80); + inside--; + return "done"; + }; + const results = await Promise.all([ + a.withLock("thread:1", section), + b.withLock("thread:1", section), + a.withLock("thread:1", section), + ]); + expect(results).toEqual(["done", "done", "done"]); + expect(max).toBe(1); + }); + + it("different keys do not block each other", async () => { + const locks = upstashLocks({ redis, prefix }); + const started = Date.now(); + await Promise.all([ + locks.withLock("k:a", () => sleep(300)), + locks.withLock("k:b", () => sleep(300)), + ]); + expect(Date.now() - started).toBeLessThan(550); + }); + + it("passes a live signal and releases on failure", async () => { + const locks = upstashLocks({ redis, prefix }); + await expect( + locks.withLock("k:c", async (signal) => { + expect(signal.aborted).toBe(false); + throw new Error("x"); + }), + ).rejects.toThrow("x"); + expect(await locks.withLock("k:c", async () => 1)).toBe(1); + }); +}); diff --git a/packages/tanstack-ai/src/locks/locks.ts b/packages/tanstack-ai/src/locks/locks.ts new file mode 100644 index 0000000..7d995e8 --- /dev/null +++ b/packages/tanstack-ai/src/locks/locks.ts @@ -0,0 +1,41 @@ +import type { Redis } from "@upstash/redis"; +import { Redis as RedisClient } from "@upstash/redis"; +import { RedisLock } from "./redis-lock.js"; +import type { RedisLockConfig } from "./redis-lock.js"; +import type { LockStore } from "@tanstack/ai/locks"; +import { addTelemetry } from "../telemetry.js"; + +/** + * The core {@link RedisLockConfig} (lease, acquire timeout, retry delay), with `redis` optional. + * `prefix` defaults to `agentkit:tanstackLock`. + */ +export type UpstashLocksConfig = Omit & { + /** Upstash Redis client. Defaults to `Redis.fromEnv()`. */ + redis?: Redis; +}; + +/** + * A distributed TanStack AI `LockStore` on Upstash Redis — the multi-instance replacement for + * `InMemoryLockStore`. Pass it to `withLocks()`; anything that reads the locks capability (notably + * `withSandbox`'s ensure step and the run driver's claim) then coordinates across server instances. + * + * Lease-backed, as the contract asks: the lease is renewed while the section runs, and the section's + * `signal` aborts as soon as ownership can no longer be guaranteed. + * + * ```ts + * import { withLocks } from "@tanstack/ai/locks"; + * import { upstashLocks } from "@upstash/agentkit-tanstack-ai"; + * + * chat({ adapter, messages, middleware: [withLocks(upstashLocks()), withSandbox(sandbox)] }); + * ``` + */ +export function upstashLocks(config: UpstashLocksConfig = {}): LockStore { + const redis = config.redis ?? RedisClient.fromEnv(); + addTelemetry(redis, config.enableTelemetry); + const lock = new RedisLock({ + ...config, + redis, + prefix: config.prefix ?? "agentkit:tanstackLock", + }); + return { withLock: (key, fn) => lock.withLock(key, (signal) => fn(signal)) }; +} diff --git a/packages/tanstack-ai/src/locks/redis-lock.test.ts b/packages/tanstack-ai/src/locks/redis-lock.test.ts new file mode 100644 index 0000000..b764aa8 --- /dev/null +++ b/packages/tanstack-ai/src/locks/redis-lock.test.ts @@ -0,0 +1,109 @@ +import { afterAll, describe, expect, it } from "vitest"; +import { LockAcquireTimeoutError, LockLostError, RedisLock } from "./redis-lock.js"; +import { cleanupKeys, hasRedisCreds, testRedis, uniquePrefix } from "../testing/test-support.js"; + +const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)); + +describe.skipIf(!hasRedisCreds)("RedisLock (live Redis)", () => { + const redis = testRedis(); + const prefix = uniquePrefix("lock"); + + afterAll(async () => { + await cleanupKeys(redis, prefix); + }); + + it("grants the key to one holder at a time, with increasing fencing tokens", async () => { + const lock = new RedisLock({ redis, prefix }); + const a = await lock.tryAcquire("k1"); + expect(a).not.toBeNull(); + expect(await lock.tryAcquire("k1")).toBeNull(); + expect(await a!.release()).toBe(true); + const b = await lock.tryAcquire("k1"); + expect(b!.fencingToken).toBeGreaterThan(a!.fencingToken); + await b!.release(); + }); + + it("a stale holder cannot release or extend the next holder's lease", async () => { + const lock = new RedisLock({ redis, prefix }); + const stale = await lock.tryAcquire("k2", { leaseMs: 300 }); + await sleep(500); + const fresh = await lock.tryAcquire("k2"); + expect(fresh).not.toBeNull(); + expect(await stale!.release()).toBe(false); + expect(await stale!.extend()).toBe(false); + expect(await lock.tryAcquire("k2")).toBeNull(); + await fresh!.release(); + }); + + it("withLock serializes concurrent critical sections", async () => { + const lock = new RedisLock({ redis, prefix, retryDelayMs: 20 }); + let inside = 0; + let maxInside = 0; + const order: number[] = []; + await Promise.all( + [1, 2, 3].map((n) => + lock.withLock("k3", async () => { + inside += 1; + maxInside = Math.max(maxInside, inside); + await sleep(100); + order.push(n); + inside -= 1; + }), + ), + ); + expect(maxInside).toBe(1); + expect(order.sort()).toEqual([1, 2, 3]); + expect(await redis.exists(`${prefix}:lease:k3`)).toBe(0); + }); + + it("releases the lease when the critical section throws", async () => { + const lock = new RedisLock({ redis, prefix }); + await expect( + lock.withLock("k4", async () => { + throw new Error("boom"); + }), + ).rejects.toThrow("boom"); + expect(await lock.tryAcquire("k4")).not.toBeNull(); + }); + + it("renews the lease while the section runs past its lifetime", async () => { + const lock = new RedisLock({ redis, prefix, leaseMs: 600 }); + await lock.withLock("k5", async (signal) => { + await sleep(1_500); + expect(signal.aborted).toBe(false); + expect(await lock.tryAcquire("k5")).toBeNull(); + }); + }); + + it("aborts the signal when the lease is taken away", async () => { + const lock = new RedisLock({ redis, prefix, leaseMs: 600 }); + const reason = await lock.withLock("k6", async (signal) => { + await redis.del(`${prefix}:lease:k6`); // simulate expiry + takeover + await sleep(800); + return signal.reason; + }); + expect(reason).toBeInstanceOf(LockLostError); + }); + + it("keeps leases and fencing counters apart for colliding-looking keys", async () => { + const lock = new RedisLock({ redis, prefix }); + const x = await lock.tryAcquire("x"); + expect(x).not.toBeNull(); + const fenceX = await lock.tryAcquire("fence:x"); + expect(fenceX).not.toBeNull(); // would be refused if its lease were x's fencing counter + expect(await x!.release()).toBe(true); + expect(await fenceX!.release()).toBe(true); + const again = await lock.tryAcquire("x"); + expect(again!.fencingToken).toBeGreaterThan(x!.fencingToken); + await again!.release(); + }); + + it("times out acquiring a held key", async () => { + const lock = new RedisLock({ redis, prefix, acquireTimeoutMs: 300, retryDelayMs: 50 }); + const held = await lock.tryAcquire("k7"); + await expect(lock.withLock("k7", async () => 1)).rejects.toBeInstanceOf( + LockAcquireTimeoutError, + ); + await held!.release(); + }); +}); diff --git a/packages/tanstack-ai/src/locks/redis-lock.ts b/packages/tanstack-ai/src/locks/redis-lock.ts new file mode 100644 index 0000000..f3a451b --- /dev/null +++ b/packages/tanstack-ai/src/locks/redis-lock.ts @@ -0,0 +1,230 @@ +import { randomUUID } from "node:crypto"; +import type { Redis } from "@upstash/redis"; +import { addTelemetry } from "../telemetry.js"; + +/** + * Every script runs with `allow-key-locking`, so Upstash locks only the keys it declares rather than + * the whole database. That makes declaring every touched key in `KEYS` mandatory (an undeclared key + * is an error), which these scripts do. + * + * Acquire: take the lease key only if it is free, and on success bump a per-key counter that never + * expires. The counter is the **fencing token** — strictly increasing across every holder the key + * ever had, so a holder that stalled past its lease can compare its token against the store and + * learn it was superseded (a lease alone tells the winner it won, but gives a loser nothing to read). + */ +const ACQUIRE = `#!lua flags=allow-key-locking +if redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2]) then + return redis.call("INCR", KEYS[2]) +end +return 0`; + +/** Release only our own lease: a holder whose lease expired must not delete the next holder's. */ +const RELEASE = `#!lua flags=allow-key-locking +if redis.call("GET", KEYS[1]) == ARGV[1] then + return redis.call("DEL", KEYS[1]) +end +return 0`; + +/** Extend only our own lease — the same ownership check as release. */ +const EXTEND = `#!lua flags=allow-key-locking +if redis.call("GET", KEYS[1]) == ARGV[1] then + return redis.call("PEXPIRE", KEYS[1], ARGV[2]) +end +return 0`; + +export interface RedisLockConfig { + /** Upstash Redis client. */ + redis: Redis; + /** Base key prefix. Defaults to `agentkit:lock`. */ + prefix?: string; + /** + * How long a lease lives without renewal, in ms. `withLock` renews it in the background, so this + * only bounds how long a crashed holder blocks the key. + * @default 30000 + */ + leaseMs?: number; + /** + * How long `withLock` keeps retrying a held key before throwing {@link LockAcquireTimeoutError}. + * @default 30000 + */ + acquireTimeoutMs?: number; + /** + * Delay between acquire attempts while the key is held, in ms. + * @default 100 + */ + retryDelayMs?: number; + /** + * Report the sdk name + version to Upstash as a header on the requests made by your redis client. + * Can also be disabled with the `UPSTASH_DISABLE_TELEMETRY` env var. Defaults to `true`. + */ + enableTelemetry?: boolean; +} + +/** A held lease. */ +export interface LockLease { + /** The lock key (without the prefix). */ + key: string; + /** Opaque per-acquisition owner token. */ + token: string; + /** + * Strictly increasing across every acquisition of this key, ever. Store it next to anything the + * holder writes and reject writes carrying an older one to make stale holders harmless. + */ + fencingToken: number; + /** Push the expiry out by `leaseMs` (or the configured lease). `false` = the lease was lost. */ + extend(leaseMs?: number): Promise; + /** Release the lease if we still hold it. `false` = it had already expired or been taken. */ + release(): Promise; +} + +/** Thrown by `withLock` when the key stays held past `acquireTimeoutMs`. */ +export class LockAcquireTimeoutError extends Error { + constructor( + readonly key: string, + readonly timeoutMs: number, + ) { + super(`RedisLock: could not acquire "${key}" within ${timeoutMs}ms.`); + this.name = "LockAcquireTimeoutError"; + } +} + +/** The reason `withLock`'s signal aborts with when renewal finds the lease gone. */ +export class LockLostError extends Error { + constructor(readonly key: string) { + super(`RedisLock: lease on "${key}" was lost before the critical section finished.`); + this.name = "LockLostError"; + } +} + +function assertKey(key: string): void { + if (typeof key !== "string" || key === "") { + throw new Error("RedisLock: `key` is required and must be a non-empty string."); + } +} + +const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); + +/** + * A distributed mutex on Upstash Redis: a lease key (`SET NX PX`) plus a fencing token, with + * ownership-checked release/extend in Lua. Works over the REST API (every operation is a single + * command or one `EVAL`, so it needs no connection-scoped state like `WATCH`). + * + * Keys: `:lease:` (the lease) and `:fence:` (the never-expiring counter). + * + * ```ts + * const lock = new RedisLock({ redis }); + * await lock.withLock("thread:42", async (signal) => { + * // only one process at a time gets here; stop writing if `signal` aborts (lease lost) + * }); + * ``` + */ +export class RedisLock { + private redis: Redis; + private prefix: string; + private leaseMs: number; + private acquireTimeoutMs: number; + private retryDelayMs: number; + + constructor(config: RedisLockConfig) { + this.redis = config.redis; + addTelemetry(config.redis, config.enableTelemetry); + this.prefix = config.prefix ?? "agentkit:lock"; + this.leaseMs = config.leaseMs ?? 30_000; + this.acquireTimeoutMs = config.acquireTimeoutMs ?? 30_000; + this.retryDelayMs = config.retryDelayMs ?? 100; + } + + // Leases and fencing counters live in separate namespaces so no user key can name the other's + // Redis key (the lease for "fence:x" must not be the counter for "x"). + private leaseKey(key: string): string { + return `${this.prefix}:lease:${key}`; + } + + private fenceKey(key: string): string { + return `${this.prefix}:fence:${key}`; + } + + /** One acquire attempt. Returns the lease, or `null` when the key is held. */ + async tryAcquire(key: string, opts: { leaseMs?: number } = {}): Promise { + assertKey(key); + const leaseMs = opts.leaseMs ?? this.leaseMs; + const token = randomUUID(); + const leaseKey = this.leaseKey(key); + const fencingToken = Number( + await this.redis.eval(ACQUIRE, [leaseKey, this.fenceKey(key)], [token, String(leaseMs)]), + ); + if (!fencingToken) return null; + const redis = this.redis; + return { + key, + token, + fencingToken, + extend: async (ms?: number) => + Number(await redis.eval(EXTEND, [leaseKey], [token, String(ms ?? leaseMs)])) === 1, + release: async () => Number(await redis.eval(RELEASE, [leaseKey], [token])) === 1, + }; + } + + /** + * Acquire, retrying every `retryDelayMs` until `acquireTimeoutMs` (then + * {@link LockAcquireTimeoutError}). Pass `signal` to give up early. + */ + async acquire( + key: string, + opts: { leaseMs?: number; acquireTimeoutMs?: number; signal?: AbortSignal } = {}, + ): Promise { + const timeoutMs = opts.acquireTimeoutMs ?? this.acquireTimeoutMs; + const deadline = Date.now() + timeoutMs; + for (;;) { + opts.signal?.throwIfAborted(); + const lease = await this.tryAcquire( + key, + opts.leaseMs !== undefined ? { leaseMs: opts.leaseMs } : {}, + ); + if (lease) return lease; + if (Date.now() >= deadline) throw new LockAcquireTimeoutError(key, timeoutMs); + await sleep(Math.min(this.retryDelayMs, Math.max(0, deadline - Date.now()))); + } + } + + /** + * Run `fn` while holding `key`. The lease is renewed every third of its lifetime; if a renewal finds + * it gone (expired after a stall, or taken), `signal` aborts with {@link LockLostError} — stop + * externally visible writes when it does. The lease is released when `fn` settles. + */ + async withLock( + key: string, + fn: (signal: AbortSignal, lease: LockLease) => Promise, + opts: { leaseMs?: number; acquireTimeoutMs?: number; signal?: AbortSignal } = {}, + ): Promise { + const lease = await this.acquire(key, opts); + const leaseMs = opts.leaseMs ?? this.leaseMs; + const controller = new AbortController(); + let renewing = false; + const timer = setInterval( + () => { + if (renewing || controller.signal.aborted) return; + renewing = true; + lease + .extend(leaseMs) + .then((held) => { + if (!held) controller.abort(new LockLostError(key)); + }) + // A failed renewal is not proof of loss, but ownership can no longer be guaranteed. + .catch(() => controller.abort(new LockLostError(key))) + .finally(() => { + renewing = false; + }); + }, + Math.max(10, Math.floor(leaseMs / 3)), + ); + // Never keep the process alive just to renew a lease. + (timer as { unref?: () => void }).unref?.(); + try { + return await fn(controller.signal, lease); + } finally { + clearInterval(timer); + await lease.release().catch(() => false); + } + } +} diff --git a/packages/tanstack-ai/src/memory/index.ts b/packages/tanstack-ai/src/memory/index.ts new file mode 100644 index 0000000..c0b4237 --- /dev/null +++ b/packages/tanstack-ai/src/memory/index.ts @@ -0,0 +1,4 @@ +// `@upstash/agentkit-tanstack-ai/memory` — a long-term memory `MemoryAdapter` for +// `memoryMiddleware()`, ranked in Redis Search. Needs `@tanstack/ai-memory`. +export { upstashMemory, memoryScopeKey } from "./memory.js"; +export type { UpstashMemoryConfig } from "./memory.js"; diff --git a/packages/tanstack-ai/src/memory/memory.test.ts b/packages/tanstack-ai/src/memory/memory.test.ts new file mode 100644 index 0000000..2bc9ac9 --- /dev/null +++ b/packages/tanstack-ai/src/memory/memory.test.ts @@ -0,0 +1,189 @@ +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { chat } from "@tanstack/ai"; +import type { StreamChunk } from "@tanstack/ai"; +import { memoryMiddleware } from "@tanstack/ai-memory"; +import { runMemoryAdapterContract } from "@tanstack/ai-memory/testkit"; +import { s } from "@upstash/redis"; +import { AgentMemory } from "@upstash/agentkit-sdk"; +import { memoryScopeKey, upstashMemory } from "./memory.js"; +import { scriptedAdapter } from "../testing/test-adapter.js"; +import { cleanupKeys, hasRedisCreds, testRedis, uniqueUserId } from "../testing/test-support.js"; + +async function drain(stream: unknown): Promise { + const out: StreamChunk[] = []; + for await (const c of stream as AsyncIterable) out.push(c); + return out; +} + +async function pollUntil(read: () => Promise, ready: (v: T) => boolean): Promise { + const deadline = Date.now() + 8_000; + let value = await read(); + while (!ready(value) && Date.now() < deadline) { + await new Promise((r) => setTimeout(r, 250)); + value = await read(); + } + return value; +} + +const promptText = (call: { systemPrompts?: unknown }) => JSON.stringify(call.systemPrompts ?? ""); + +describe("memoryScopeKey", () => { + it("partitions by user across threads, and cannot be forged by separators", () => { + expect(memoryScopeKey({ threadId: "t1", userId: "u" })).toBe( + memoryScopeKey({ threadId: "t2", userId: "u" }), + ); + expect(memoryScopeKey({ threadId: "t1", userId: "u" }, "thread")).not.toBe( + memoryScopeKey({ threadId: "t2", userId: "u" }, "thread"), + ); + expect(memoryScopeKey({ threadId: "t", userId: "a.u.b" })).not.toBe( + memoryScopeKey({ threadId: "t", userId: "a", tenantId: "u.b" }), + ); + expect(memoryScopeKey({ threadId: "t", userId: "_" })).not.toBe( + memoryScopeKey({ threadId: "t" }), + ); + expect(memoryScopeKey({ threadId: "x:y", userId: "a:b" })).not.toContain(":"); + }); +}); + +describe.skipIf(!hasRedisCreds)("upstashMemory (live Redis, real chat loop)", () => { + const redis = testRedis(); + const userId = uniqueUserId("tsmem"); + const other = uniqueUserId("tsmem-other"); + const adapter = upstashMemory({ redis }); + // A handle on the same index, only to wait for indexing between a write and the next recall. + const index = new AgentMemory({ + redis, + prefix: "agentkit:tanstackMemory", + metadataSchema: { source: s.string().noTokenize() }, + }).searchIndex; + + beforeAll(async () => { + // Provision the index before anything is written into its keyspace (see CLAUDE.md, Testing). + await adapter.recall({ threadId: "probe", userId }, "provisioning probe"); + }); + + const manyUsers: string[] = []; + afterAll(async () => { + for (const u of [userId, other, ...manyUsers]) { + await cleanupKeys( + redis, + `agentkit:tanstackMemory:${memoryScopeKey({ threadId: "x", userId: u })}:`, + ); + } + }); + + it("a fact the model saves in one thread is recalled into the prompt of another thread", async () => { + const mw = () => + memoryMiddleware({ adapter, scope: (ctx) => ({ threadId: ctx.threadId, userId }) }); + + // Thread 1: the model calls save_memory, which the adapter offered as a tool this turn. + const first = scriptedAdapter([ + { + toolCalls: [ + { id: "s1", name: "save_memory", args: { text: "User is allergic to hazelnuts" } }, + ], + }, + { text: "Noted." }, + ]); + await drain( + chat({ + adapter: first as never, + threadId: "thread-1", + messages: [{ role: "user", content: "remember: no hazelnuts for me" }], + middleware: [mw()], + }), + ); + const offered = (first.calls[0]!.tools ?? []) as { name: string }[]; + expect(offered.map((t) => t.name)).toContain("save_memory"); + await index.waitIndexing(); + + // Thread 2: recall runs before the model and injects the fact into the system prompt. + const facts = await pollUntil( + () => adapter.listFacts!({ threadId: "thread-2", userId }), + (f) => f.some((x) => x.text.includes("hazelnuts") && x.source === "agent"), + ); + expect(facts.some((x) => x.text.includes("hazelnuts"))).toBe(true); + + const second = scriptedAdapter([{ text: "Try the almond cake." }]); + await drain( + chat({ + adapter: second as never, + threadId: "thread-2", + messages: [{ role: "user", content: "suggest a dessert without hazelnuts" }], + middleware: [mw()], + }), + ); + const prompt = promptText(second.calls[0]!); + expect(prompt).toContain("User is allergic to hazelnuts"); + expect(prompt).toContain("you saved this"); + }); + + it("captures user messages on save, labelled apart from saved facts", async () => { + await adapter.save({ threadId: "t", userId }, { user: "I live in Izmir", assistant: "Nice!" }); + await index.waitIndexing(); + const recalled = await pollUntil( + () => adapter.recall({ threadId: "t9", userId }, "where do I live Izmir"), + (r) => r.systemPrompt.includes("Izmir"), + ); + expect(recalled.systemPrompt).toContain("I live in Izmir (the user said this)"); + }); + + it("listFacts returns every fact in a scope, past the default page size", async () => { + const many = uniqueUserId("tsmem-many"); + manyUsers.push(many); + const store = new AgentMemory({ + redis, + prefix: "agentkit:tanstackMemory", + metadataSchema: { source: s.string().noTokenize() }, + }); + const scopeUser = memoryScopeKey({ threadId: "x", userId: many }); + await store.count({ userId: scopeUser }); // provision before writing + await Promise.all( + Array.from({ length: 120 }, (_, i) => + store.add({ + userId: scopeUser, + text: `fact number ${i}`, + id: `f${i}`, + metadata: { source: "agent" }, + }), + ), + ); + await store.searchIndex.waitIndexing(); + const facts = await pollUntil( + () => adapter.listFacts!({ threadId: "x", userId: many }), + (f) => f.length >= 120, + ); + expect(facts).toHaveLength(120); + }); + + it("is isolated per user", async () => { + const r = await adapter.recall({ threadId: "t", userId: other }, "hazelnuts Izmir"); + expect(r.systemPrompt).toBe(""); + expect(r.fragments).toEqual([]); + }); +}); + +// TanStack AI's own MemoryAdapter contract suite (save receipts, round-trip recall, scope isolation +// by thread / user / tenant, listFacts), run against a real Upstash Redis. +describe.skipIf(!hasRedisCreds)("upstashMemory contract (live Redis)", () => { + const redis = testRedis(); + // Share the default index; each contract case uses its own scopes, isolated by a unique tenant. + const tenant = uniqueUserId("tscontract"); + afterAll(async () => { + await cleanupKeys(redis, `agentkit:tanstackMemory:${encodeURIComponent(tenant)}`); + }); + runMemoryAdapterContract("upstashMemory", () => { + const adapter = upstashMemory({ redis }); + // Pin every scope under this run's tenant so the suite cannot see earlier runs' data. + const pin = (scope: T): T => ({ + ...scope, + tenantId: `${tenant}:${scope.tenantId ?? ""}`, + }); + return { + ...adapter, + recall: (scope, query) => adapter.recall(pin(scope), query), + save: (scope, turn) => adapter.save(pin(scope), turn), + listFacts: (scope) => adapter.listFacts!(pin(scope)), + }; + }); +}); diff --git a/packages/tanstack-ai/src/memory/memory.ts b/packages/tanstack-ai/src/memory/memory.ts new file mode 100644 index 0000000..c582475 --- /dev/null +++ b/packages/tanstack-ai/src/memory/memory.ts @@ -0,0 +1,226 @@ +import type { Redis } from "@upstash/redis"; +import { Redis as RedisClient, s } from "@upstash/redis"; +import { z } from "zod"; +import { AgentMemory, stableHash } from "@upstash/agentkit-sdk"; +import type { AgentMemoryConfig } from "@upstash/agentkit-sdk"; +import { toolDefinition } from "@tanstack/ai"; +import type { Tool } from "@tanstack/ai"; +import type { + MemoryAdapter, + MemoryFact, + MemoryScope, + RecallResult, + SaveReceipt, +} from "@tanstack/ai-memory"; +import { addTelemetry } from "../telemetry.js"; + +/** Where a memory came from — shown next to it in the recalled block, since they differ in weight. */ +type Source = "agent" | "userMessage"; + +const METADATA = { source: s.string().noTokenize() }; + +/** + * The core `AgentMemoryConfig` options this adapter passes through (`prefix`, `indexName`, + * `minScore`, `enableTelemetry`), with `redis` optional, plus how the adapter scopes, recalls and + * captures. `prefix` defaults to `agentkit:tanstackMemory` — its own keyspace (and search index), + * because the store indexes a `source` field the plain `agentkit:memory` store does not have. + */ +export type UpstashMemoryConfig = Pick< + AgentMemoryConfig, + "prefix" | "indexName" | "minScore" | "enableTelemetry" +> & { + /** Upstash Redis client. Defaults to `Redis.fromEnv()`. */ + redis?: Redis; + /** + * What memory is partitioned by. `"user"` keeps one long-term memory per user across every + * thread (falling back to the thread when the scope has no `userId`); `"thread"` keeps it per + * conversation. `tenantId` and `namespace` always partition as well. + * @default "user" + */ + scopeBy?: "user" | "thread"; + /** + * Max memories injected per turn. + * @default 5 + */ + topK?: number; + /** + * Store each turn's user message automatically. `false` = only what the model saves with the + * `save_memory` tool. + * @default true + */ + captureUserMessages?: boolean; + /** + * Expose a `save_memory` tool so the model can store durable facts deliberately. Such facts are + * labelled apart from captured messages in the recalled block. + * @default true + */ + saveTool?: boolean; + /** + * Name of the save tool. + * @default "save_memory" + */ + saveToolName?: string; + /** + * Longest message captured (characters); longer ones are truncated. + * @default 2000 + */ + maxMemoryCharacters?: number; + /** + * Wait for the search index to catch up after each write, so a memory saved in one turn is + * recallable on the very next one. Upstash Search indexing otherwise lags by minutes. Costs one + * extra round trip per save, which `memoryMiddleware` runs after the response is delivered. + * @default true + */ + waitForIndexing?: boolean; +}; + +/** Escape one scope part so no value can forge a separator (`.`) or the key separator (`:`). */ +const part = (v: string | undefined) => + v === undefined || v === "" + ? "_" + : encodeURIComponent(v).replace(/\./g, "%2E").replace(/_/g, "%5F"); + +/** The `AgentMemory` userId a scope maps to — `:`-free, so it is safe as a key part. */ +export function memoryScopeKey(scope: MemoryScope, scopeBy: "user" | "thread" = "user"): string { + const subject = + scopeBy === "user" && scope.userId ? `u.${part(scope.userId)}` : `t.${part(scope.threadId)}`; + return [part(scope.tenantId), part(scope.namespace), subject].join("."); +} + +const LABEL: Record = { + agent: "you saved this", + userMessage: "the user said this", +}; + +/** + * A TanStack AI `MemoryAdapter` on Upstash Redis Search — plug it into `memoryMiddleware()`. + * + * Unlike a plain key/value memory store, ranking happens **in the database**: recall is one BM25 + * `$smart` (typo-tolerant, fuzzy) query over the scope's memories, so it neither loads every record + * per turn nor caps how many a user can have. + * + * - **recall** injects the top matches as a system-prompt block, each labelled with where it came + * from, and offers the `save_memory` tool. + * - **save** captures the turn's user message (idempotent — the id is a hash of the text). + * + * ```ts + * import { memoryMiddleware } from "@tanstack/ai-memory"; + * import { upstashMemory } from "@upstash/agentkit-tanstack-ai/memory"; + * + * chat({ + * adapter, messages, + * middleware: [memoryMiddleware({ adapter: upstashMemory(), scope: { threadId, userId } })], + * }); + * ``` + */ +export function upstashMemory(config: UpstashMemoryConfig = {}): MemoryAdapter { + const redis = config.redis ?? RedisClient.fromEnv(); + addTelemetry(redis, config.enableTelemetry); + const { prefix, indexName, minScore, enableTelemetry } = config; + const memory = new AgentMemory({ + ...{ indexName, minScore, enableTelemetry }, + redis, + prefix: prefix ?? "agentkit:tanstackMemory", + metadataSchema: METADATA, + }); + const scopeBy = config.scopeBy ?? "user"; + const topK = config.topK ?? 5; + const maxChars = config.maxMemoryCharacters ?? 2_000; + const capture = config.captureUserMessages ?? true; + const toolName = config.saveToolName ?? "save_memory"; + + const waitForIndexing = config.waitForIndexing ?? true; + // Writes do not create the index, and waiting on a missing index is a silent no-op — a doc written + // before the index exists can miss the create-time backfill. So the first write provisions it + // (any read does, reactively), once per adapter. + let provisioned: Promise | undefined; + + const add = async (userId: string, text: string, source: Source) => { + const trimmed = text.trim().slice(0, maxChars); + provisioned ??= memory.count({ userId }).catch((error) => { + provisioned = undefined; + throw error; + }); + await provisioned; + // Identical text collapses onto one record, so capture is idempotent across turns and retries. + const record = await memory.add({ + userId, + text: trimmed, + id: stableHash(trimmed).slice(0, 12), + metadata: { source }, + }); + if (waitForIndexing) await memory.searchIndex.waitIndexing(); + return record; + }; + + const saveToolFor = (userId: string): Tool => + toolDefinition({ + name: toolName, + description: + "Save a durable fact about the user to long-term memory so it can be recalled in future " + + "conversations (preferences, identity, goals, ...).", + inputSchema: z.object({ + text: z.string().describe("A concise, durable fact about the user to remember for later."), + }), + }).server(async ({ text }: { text: string }) => { + const record = await add(userId, text, "agent"); + return { id: record.id, saved: true }; + }); + + return { + id: "upstash", + name: "Upstash Redis Search", + + async recall(scope, query): Promise { + const userId = memoryScopeKey(scope, scopeBy); + const hits = await memory.recall({ + userId, + query, + topK, + }); + const lines = hits.map((h) => { + const source = h.metadata?.source as Source | undefined; + return `- ${h.text}${source && LABEL[source] ? ` (${LABEL[source]})` : ""}`; + }); + const withTool = config.saveTool ?? true; + return { + systemPrompt: lines.length + ? `Relevant long-term memory about this user:\n${lines.join("\n")}` + : "", + fragments: hits.map((h) => ({ text: h.text, source: h.id })), + ...(withTool + ? { + tools: [saveToolFor(userId)], + toolGuidance: `Call \`${toolName}\` to remember durable facts about the user (preferences, identity, goals) for future conversations.`, + } + : {}), + raw: hits, + }; + }, + + async save(scope, turn): Promise { + if (!capture || !turn.user.trim()) return []; + const started = Date.now(); + try { + await add(memoryScopeKey(scope, scopeBy), turn.user, "userMessage"); + return [{ ok: true, latencyMs: Date.now() - started }]; + } catch (error) { + return [{ ok: false, error: error instanceof Error ? error.message : String(error) }]; + } + }, + + async listFacts(scope): Promise { + // Every fact in the scope: `list` has no cursor, so size the page to the scope's count. + const userId = memoryScopeKey(scope, scopeBy); + const total = await memory.count({ userId }); + if (total === 0) return []; + const records = await memory.list({ userId, limit: total }); + return records.map((r) => ({ + id: r.id, + text: r.text, + ...(r.metadata?.source ? { source: String(r.metadata.source) } : {}), + createdAt: new Date(r.createdAt).toISOString(), + })); + }, + }; +} diff --git a/packages/tanstack-ai/src/middleware/middleware.test.ts b/packages/tanstack-ai/src/middleware/middleware.test.ts new file mode 100644 index 0000000..f294186 --- /dev/null +++ b/packages/tanstack-ai/src/middleware/middleware.test.ts @@ -0,0 +1,159 @@ +import { afterAll, describe, expect, it } from "vitest"; +import { chat, toolDefinition } from "@tanstack/ai"; +import type { StreamChunk } from "@tanstack/ai"; +import { z } from "zod"; +import { Ratelimit } from "@upstash/agentkit-sdk"; +import { rateLimit, toolCache } from "./middleware.js"; +import { scriptedAdapter } from "../testing/test-adapter.js"; +import { + cleanupKeys, + hasRedisCreds, + testRedis, + uniquePrefix, + uniqueUserId, +} from "../testing/test-support.js"; + +async function drain(stream: unknown): Promise { + const out: StreamChunk[] = []; + for await (const c of stream as AsyncIterable) out.push(c); + return out; +} + +describe.skipIf(!hasRedisCreds)("toolCache middleware (live Redis, real chat loop)", () => { + const redis = testRedis(); + const prefix = uniquePrefix("tstc"); + afterAll(async () => { + await cleanupKeys(redis, prefix); + }); + + function setup(result: () => unknown) { + const counter = { weather: 0, email: 0 }; + const weather = toolDefinition({ + name: "get_weather", + description: "weather", + inputSchema: z.object({ city: z.string() }), + }).server(async ({ city }) => { + counter.weather++; + return { city, ...(result() as object) }; + }); + const email = toolDefinition({ + name: "send_email", + description: "email", + inputSchema: z.object({ to: z.string() }), + }).server(async () => { + counter.email++; + return { sent: true }; + }); + return { counter, tools: [weather, email] }; + } + + const turns = (city: string) => [ + { + toolCalls: [ + { id: `w-${city}-${Math.random()}`, name: "get_weather", args: { city } }, + { id: `e-${Math.random()}`, name: "send_email", args: { to: "a@b.c" } }, + ], + }, + { text: "done" }, + ]; + + it("serves a repeated call from Redis and skips the tool; never caches tools not allowlisted", async () => { + const { counter, tools } = setup(() => ({ temp: 21 })); + const userId = uniqueUserId("tc"); + const mw = toolCache({ tools: ["get_weather"], userId, redis, prefix }); + const messages = [{ role: "user" as const, content: "weather + email" }]; + + await drain( + chat({ + adapter: scriptedAdapter(turns("Paris")) as never, + messages, + tools, + middleware: [mw], + }), + ); + const second = scriptedAdapter(turns("Paris")); + await drain(chat({ adapter: second as never, messages, tools, middleware: [mw] })); + + expect(counter.weather).toBe(1); // second run was a cache hit + expect(counter.email).toBe(2); // not allowlisted: always runs + // The model still received the cached result as the tool's output. + const toolMsg = (second.calls[1]!.messages as { role: string; content: string }[]).find( + (m) => m.role === "tool" && m.content.includes("Paris"), + ); + expect(JSON.parse(toolMsg!.content)).toEqual({ city: "Paris", temp: 21 }); + }); + + it("keys by arguments and by user", async () => { + const { counter, tools } = setup(() => ({ temp: 1 })); + const messages = [{ role: "user" as const, content: "w" }]; + const run = (userId: string, city: string) => + drain( + chat({ + adapter: scriptedAdapter(turns(city)) as never, + messages, + tools, + middleware: [toolCache({ tools: ["get_weather"], userId, redis, prefix })], + }), + ); + const alice = uniqueUserId("alice"); + await run(alice, "Rome"); + await run(alice, "Oslo"); + await run(uniqueUserId("bob"), "Rome"); + await run(alice, "Rome"); + expect(counter.weather).toBe(3); + }); + + it("does not cache a failed call", async () => { + let fail = true; + let runs = 0; + const flaky = toolDefinition({ + name: "get_weather", + description: "w", + inputSchema: z.object({ city: z.string() }), + }).server(async () => { + runs++; + if (fail) throw new Error("upstream down"); + return { ok: true }; + }); + const mw = toolCache({ tools: ["get_weather"], userId: uniqueUserId("f"), redis, prefix }); + const script = () => + scriptedAdapter([ + { toolCalls: [{ id: `f-${Math.random()}`, name: "get_weather", args: { city: "X" } }] }, + { text: "ok" }, + ]); + const messages = [{ role: "user" as const, content: "w" }]; + await drain(chat({ adapter: script() as never, messages, tools: [flaky], middleware: [mw] })); + fail = false; + await drain(chat({ adapter: script() as never, messages, tools: [flaky], middleware: [mw] })); + await drain(chat({ adapter: script() as never, messages, tools: [flaky], middleware: [mw] })); + expect(runs).toBe(2); // failure, then a real success, then a hit + }); +}); + +describe.skipIf(!hasRedisCreds)("rateLimit middleware (live Redis, real chat loop)", () => { + const redis = testRedis(); + const prefix = uniquePrefix("tsrl"); + afterAll(async () => { + await cleanupKeys(redis, prefix); + }); + + it("lets runs through up to the limit, then fails the run before the model is called", async () => { + const identifier = uniqueUserId("rl"); + const mw = rateLimit({ redis, prefix, limiter: Ratelimit.fixedWindow(2, "60 s"), identifier }); + const messages = [{ role: "user" as const, content: "hi" }]; + const adapters = [1, 2, 3].map(() => scriptedAdapter([{ text: "hello" }])); + const outcomes: string[] = []; + for (const adapter of adapters) { + try { + const chunks = await drain(chat({ adapter: adapter as never, messages, middleware: [mw] })); + const err = chunks.find((c) => (c as { type: string }).type === "RUN_ERROR"); + outcomes.push(err ? `error:${(err as { message?: string }).message}` : "ok"); + } catch (error) { + outcomes.push(`threw:${(error as Error).name}`); + } + } + expect(outcomes.slice(0, 2)).toEqual(["ok", "ok"]); + expect(outcomes[2]).toMatch(/Rate limit exceeded|RateLimitExceededError/); + expect(adapters[2]!.calls).toHaveLength(0); // the model was never called + }); +}); diff --git a/packages/tanstack-ai/src/middleware/middleware.ts b/packages/tanstack-ai/src/middleware/middleware.ts new file mode 100644 index 0000000..c073a0c --- /dev/null +++ b/packages/tanstack-ai/src/middleware/middleware.ts @@ -0,0 +1,123 @@ +import type { Redis } from "@upstash/redis"; +import { Redis as RedisClient } from "@upstash/redis"; +import { ToolCache, createRateLimit } from "@upstash/agentkit-sdk"; +import type { RateLimitConfig, ToolCacheConfig } from "@upstash/agentkit-sdk"; +import type { ChatMiddleware, ChatMiddlewareContext } from "@tanstack/ai"; +import { addTelemetry } from "../telemetry.js"; + +/** + * The core {@link ToolCacheConfig} (`prefix`, `ttlSeconds`), with `redis` optional, plus which tools + * to cache and who the entries belong to. + */ +export type ToolCacheMiddlewareConfig = Omit & { + /** + * Names of the tools whose results may be cached. Required on purpose: only deterministic, + * side-effect-free tools belong here (a cached `send_email` would silently stop sending). + */ + tools: string[]; + /** The user entries are scoped to — a string, or derived from the run context. */ + userId: string | ((ctx: ChatMiddlewareContext) => string); + /** Upstash Redis client. Defaults to `Redis.fromEnv()`. */ + redis?: Redis; +}; + +/** + * Chat middleware that memoizes tool results in Redis, keyed by `userId` + tool name + a stable hash + * of the arguments. A hit skips the tool entirely (`onBeforeToolCall` → `skip`); a successful miss is + * stored after it runs. Failed calls are never cached, and a result served from the cache is not + * re-written. + * + * ```ts + * chat({ + * adapter, messages, tools: [getWeather], + * middleware: [toolCache({ tools: ["get_weather"], userId, ttlSeconds: 600 })], + * }); + * ``` + */ +export function toolCache(config: ToolCacheMiddlewareConfig): ChatMiddleware { + const redis = config.redis ?? RedisClient.fromEnv(); + addTelemetry(redis, config.enableTelemetry); + const { tools, userId, ...cacheConfig } = config; + const cache = new ToolCache({ ...cacheConfig, redis }); + const allowed = new Set(tools); + const resolveUserId = (ctx: ChatMiddlewareContext) => + typeof userId === "function" ? userId(ctx) : userId; + // Calls this middleware has seen: the args to key the write with, or "hit" when served. Keyed by + // request + call id, because one middleware instance can serve concurrent `chat()` calls and + // providers only guarantee call ids are unique within a request. + const pending = new Map(); + const callKey = (ctx: ChatMiddlewareContext, toolCallId: string) => + `${ctx.requestId}:${toolCallId}`; + + return { + name: "upstash-tool-cache", + async onBeforeToolCall(ctx, hook) { + if (!allowed.has(hook.toolName)) return; + const hit = await cache.get(resolveUserId(ctx), hook.toolName, hook.args); + if (hit) { + pending.set(callKey(ctx, hook.toolCallId), "hit"); + return { type: "skip", result: hit.value }; + } + pending.set(callKey(ctx, hook.toolCallId), { args: hook.args }); + return; + }, + async onAfterToolCall(ctx, info) { + const seen = pending.get(callKey(ctx, info.toolCallId)); + pending.delete(callKey(ctx, info.toolCallId)); + if (!seen || seen === "hit" || !info.ok || !allowed.has(info.toolName)) return; + await cache.set(resolveUserId(ctx), info.toolName, seen.args, info.result); + }, + }; +} + +/** Thrown from `onStart` when the caller is over its limit; surfaces as the run's error. */ +export class RateLimitExceededError extends Error { + constructor( + readonly identifier: string, + readonly limit: number, + readonly remaining: number, + /** Epoch ms when the window resets. */ + readonly reset: number, + ) { + super( + `Rate limit exceeded for "${identifier}". Try again after ${new Date(reset).toISOString()}.`, + ); + this.name = "RateLimitExceededError"; + } +} + +export interface RateLimitMiddlewareConfig extends RateLimitConfig { + /** Who is being limited — usually the authenticated user id, derived server-side. */ + identifier: string | ((ctx: ChatMiddlewareContext) => string | Promise); +} + +/** + * Chat middleware that spends one unit of an Upstash Ratelimit per run, before the model is called, + * and fails the run with {@link RateLimitExceededError} when the identifier is over its limit. + * + * To answer with an HTTP 429 instead of a streamed error, call `createRateLimit(...).limit(id)` in + * your route before `chat()` — both are exported. + * + * ```ts + * middleware: [rateLimit({ limiter: Ratelimit.slidingWindow(10, "60 s"), identifier: userId })] + * ``` + */ +export function rateLimit(config: RateLimitMiddlewareConfig): ChatMiddleware { + const { identifier, ...limitConfig } = config; + // Resolve the client here so the default one is tagged too, not only a client passed in. + const redis = config.redis ?? RedisClient.fromEnv(); + addTelemetry(redis, config.enableTelemetry); + const limiter = createRateLimit({ ...limitConfig, redis }); + return { + name: "upstash-rate-limit", + async onStart(ctx) { + // Subagent child runs share the parent's budget; only the top-level run spends a unit. + if (ctx.parentRunId !== undefined) return; + const id = typeof identifier === "function" ? await identifier(ctx) : identifier; + const result = await limiter.limit(id); + if (!result.success) { + throw new RateLimitExceededError(id, result.limit, result.remaining, result.reset); + } + }, + }; +} diff --git a/packages/tanstack-ai/src/package-exports.test.ts b/packages/tanstack-ai/src/package-exports.test.ts new file mode 100644 index 0000000..b0066af --- /dev/null +++ b/packages/tanstack-ai/src/package-exports.test.ts @@ -0,0 +1,27 @@ +import { readFileSync } from "node:fs"; +import { describe, expect, it } from "vitest"; + +// Every tsup entry must be reachable through `package.json#exports`, and every export must point at +// a built entry. A subpath that is built but not exported only fails for consumers ("Can't resolve +// '@upstash/agentkit-tanstack-ai/persistence'"), never in this package's own tests. +const pkg = JSON.parse(readFileSync(new URL("../package.json", import.meta.url), "utf8")) as { + exports: Record; +}; + +// Loaded at runtime: the config sits outside `src`, so a static import would leave tsc's rootDir. +const tsupConfig = (await import(new URL("../tsup.config.ts", import.meta.url).href)) as { + default: { entry: Record }; +}; +const entries = Object.keys(tsupConfig.default.entry); + +describe("package exports", () => { + it("exports every built entry, and nothing that is not built", () => { + const expected = Object.fromEntries( + entries.map((name) => [ + name === "index" ? "." : `./${name}`, + { types: `./dist/${name}.d.ts`, import: `./dist/${name}.js` }, + ]), + ); + expect(pkg.exports).toEqual(expected); + }); +}); diff --git a/packages/tanstack-ai/src/persistence/blob-store.ts b/packages/tanstack-ai/src/persistence/blob-store.ts new file mode 100644 index 0000000..0165e2b --- /dev/null +++ b/packages/tanstack-ai/src/persistence/blob-store.ts @@ -0,0 +1,330 @@ +/* global Blob */ +import { randomUUID } from "node:crypto"; +import type { Redis } from "@upstash/redis"; +import type { + BlobBody, + BlobGetOptions, + BlobListOptions, + BlobListPage, + BlobObject, + BlobPutOptions, + BlobRange, + BlobRecord, + BlobStore, +} from "@tanstack/ai-persistence"; +import { addTelemetry } from "../telemetry.js"; +import { KEY_LOCKING, loadDocs } from "./records.js"; + +/** + * Commit a put: point the key at its newly uploaded object version, write the record (keeping the + * original `createdAt` across overwrites) and index the key, all at once. Returns the stored record + * and the version it replaced, whose bytes the caller then deletes. + * KEYS: record document, key index, current-version pointer. ARGV: record JSON, key, version path. + */ +const PUT_RECORD = `${KEY_LOCKING}local previous = redis.call("GET", KEYS[3]) or "" +redis.call("SET", KEYS[3], ARGV[3]) +local created = redis.call("JSON.GET", KEYS[1], "$.createdAt") +redis.call("JSON.SET", KEYS[1], "$", ARGV[1]) +if created and created ~= "[]" then + redis.call("JSON.SET", KEYS[1], "$.createdAt", string.sub(created, 2, -2)) +end +redis.call("ZADD", KEYS[2], 0, ARGV[2]) +return {redis.call("JSON.GET", KEYS[1]), previous}`; + +/** Remove a key's record, pointer and index entry at once; returns the version path to delete. */ +const DELETE_RECORD = `${KEY_LOCKING}local previous = redis.call("GET", KEYS[3]) or "" +redis.call("DEL", KEYS[1], KEYS[3]) +redis.call("ZREM", KEYS[2], ARGV[1]) +return previous`; + +/** + * The slice of an `@upstash/blob` `Bucket` this store uses — pass `Bucket.fromEnv()`. Structural, + * so `@upstash/blob` stays an optional peer and tests can supply a stand-in. + */ +export interface BlobBucketLike { + put( + path: string, + body: Blob | ArrayBuffer | ArrayBufferView | string | ReadableStream, + options?: { contentType?: string }, + ): Promise<{ etag: string; size: number }>; + get(path: string): Promise<{ body: ReadableStream }>; + del(target: string): Promise; + /** Used for byte-range reads (a `Range` request against a short-lived signed URL). */ + signedReadUrl?(path: string): Promise<{ url: string }>; +} + +export interface UpstashBlobStoreConfig { + /** The Upstash Blob bucket holding the bytes, e.g. `Bucket.fromEnv()` from `@upstash/blob`. */ + bucket: BlobBucketLike; + /** Upstash Redis client holding each blob's record and the key index. */ + redis: Redis; + /** + * Report the sdk name + version to Upstash as a header on the requests made by your redis client. + * Can also be disabled with the `UPSTASH_DISABLE_TELEMETRY` env var. Defaults to `true`. + */ + enableTelemetry?: boolean; + /** Redis key prefix for the records. */ + prefix: string; + /** + * Path prefix for objects in the bucket. + * @default "agentkit/tanstack/" + */ + pathPrefix?: string; +} + +const encoder = new TextEncoder(); +const decoder = new TextDecoder(); + +async function readAll(stream: ReadableStream): Promise { + const chunks: Uint8Array[] = []; + let total = 0; + const reader = stream.getReader(); + try { + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + chunks.push(value); + total += value.byteLength; + } + } finally { + reader.releaseLock(); + } + const out = new Uint8Array(total); + let offset = 0; + for (const chunk of chunks) { + out.set(chunk, offset); + offset += chunk.byteLength; + } + return out; +} + +async function toBytes(body: BlobBody): Promise { + if (typeof body === "string") return encoder.encode(body); + if (body instanceof ArrayBuffer) return new Uint8Array(body.slice(0)); + if (ArrayBuffer.isView(body)) { + return new Uint8Array(body.buffer.slice(body.byteOffset, body.byteOffset + body.byteLength)); + } + if (typeof Blob !== "undefined" && body instanceof Blob) + return new Uint8Array(await body.arrayBuffer()); + if (typeof ReadableStream !== "undefined" && body instanceof ReadableStream) return readAll(body); + throw new TypeError("Unsupported blob body."); +} + +/** Same rules as TanStack's `resolveBlobRange`: clamp the length, reject an offset outside the object. */ +function resolveRange(size: number, range: BlobRange): { offset: number; length: number } { + const { offset } = range; + if (!Number.isInteger(offset) || offset < 0 || offset >= size) { + throw new RangeError(`Blob range offset ${offset} is outside the object (size ${size}).`); + } + const remaining = size - offset; + if (range.length === undefined) return { offset, length: remaining }; + if (!Number.isInteger(range.length) || range.length < 0) { + throw new RangeError(`Blob range length ${range.length} is not valid.`); + } + return { offset, length: Math.min(range.length, remaining) }; +} + +function toArrayBuffer(bytes: Uint8Array): ArrayBuffer { + const buffer = new ArrayBuffer(bytes.byteLength); + new Uint8Array(buffer).set(bytes); + return buffer; +} + +function isNotFound(error: unknown): boolean { + const e = error as { code?: unknown; status?: unknown; message?: unknown } | null; + return ( + e?.code === "not_found" || e?.status === 404 || /not.?found|404/i.test(String(e?.message ?? "")) + ); +} + +/** + * A TanStack AI `BlobStore` split across the two stores each part fits: the **bytes** go to Upstash + * Blob (S3-compatible object storage), and each blob's **record** (size, etag, content type, custom + * metadata, created/updated times) goes to Redis next to a lexically ordered key index. + * + * That split is what makes the contract exact: `head` and `list` are Redis reads (Blob listings omit + * content type and metadata), `list` pages in key order with a key cursor, and `createdAt` survives an + * overwrite. Byte ranges are read with a `Range` request on a short-lived signed URL, so a video can be + * served in slices without downloading it whole. + * + * Writes put the bytes first and the record second, deletes remove the record first — a crash in + * between leaves at most an orphaned object, never a record pointing at nothing. + */ +export function upstashBlobStore(config: UpstashBlobStoreConfig): BlobStore { + const { bucket, redis, prefix } = config; + addTelemetry(redis, config.enableTelemetry); + const pathPrefix = config.pathPrefix ?? "agentkit/tanstack/"; + const recordKey = (key: string) => `${prefix}:blob:${key}`; + const versionKey = (key: string) => `${prefix}:blobVersion:${key}`; + const indexKey = `${prefix}:blobKeys`; + // One path segment per key, whatever characters the key holds. + // Every put uploads to a fresh, immutable version path, so concurrent writers never share an + // object: whichever commits last wins the pointer, and a record never describes another + // writer's bytes. The replaced version is deleted only after the commit. + const newVersionPath = (key: string) => `${pathPrefix}${encodeURIComponent(key)}/${randomUUID()}`; + const dropVersion = async (path: unknown) => { + if (typeof path !== "string" || path === "") return; + try { + await bucket.del(path); + } catch (error) { + if (!isNotFound(error)) throw error; + } + }; + + const assertKey = (key: string) => { + if (typeof key !== "string" || key === "") { + throw new Error("@upstash/agentkit-tanstack-ai: blob `key` must be a non-empty string."); + } + }; + + const head = async (key: string): Promise => { + assertKey(key); + return (await redis.json.get(recordKey(key))) ?? null; + }; + + async function readBytes(path: string, record: BlobRecord, range?: BlobRange) { + const size = record.size ?? 0; + const served = range ? resolveRange(size, range) : { offset: 0, length: size }; + if (served.length === 0) return { bytes: new Uint8Array(0), served }; + if (range && bucket.signedReadUrl) { + const { url } = await bucket.signedReadUrl(path); + const end = served.offset + served.length - 1; + const res = await fetch(url, { headers: { Range: `bytes=${served.offset}-${end}` } }); + if (res.status === 404) return null; + if (!res.ok) throw new Error(`Upstash Blob range read failed with ${res.status}.`); + const bytes = new Uint8Array(await res.arrayBuffer()); + // A server that ignores Range answers 200 with the whole object; slice it ourselves. + return { + bytes: + res.status === 206 ? bytes : bytes.subarray(served.offset, served.offset + served.length), + served, + }; + } + const all = await readAll((await bucket.get(path)).body); + return { bytes: all.subarray(served.offset, served.offset + served.length), served }; + } + + return { + async put(key: string, body: BlobBody, options?: BlobPutOptions): Promise { + assertKey(key); + const bytes = await toBytes(body); + const contentType = + options?.contentType ?? + (typeof Blob !== "undefined" && body instanceof Blob ? body.type || undefined : undefined); + const path = newVersionPath(key); + const stored = await bucket.put(path, bytes, contentType ? { contentType } : {}); + const now = Date.now(); + // A full replace: omitted optional fields are cleared, and only createdAt survives an overwrite. + const record: BlobRecord = { + key, + size: bytes.byteLength, + etag: stored.etag, + ...(contentType ? { contentType } : {}), + ...(options?.customMetadata ? { customMetadata: { ...options.customMetadata } } : {}), + createdAt: now, + updatedAt: now, + }; + const [stored_, previous] = (await redis.eval( + PUT_RECORD, + [recordKey(key), indexKey, versionKey(key)], + [JSON.stringify(record), key, path], + )) as [BlobRecord, string]; + if (previous !== path) await dropVersion(previous); + return stored_; + }, + + async get(key: string, options?: BlobGetOptions): Promise { + assertKey(key); + let read: Awaited> | undefined; + let record: BlobRecord | null = null; + // The record and its version pointer are read together, so the bytes fetched are the ones the + // record describes. A concurrent overwrite can still delete that version between the snapshot + // and the fetch; then the key has moved on, so read it again rather than report it missing. + for (let attempt = 0; attempt < 3; attempt++) { + const tx = redis.multi(); + tx.json.get(recordKey(key)); + tx.get(versionKey(key)); + const snapshot = (await tx.exec()) as [BlobRecord | null, string | null]; + record = snapshot[0]; + const path = snapshot[1]; + if (!record || !path) return null; + try { + read = await readBytes(path, record, options?.range); + } catch (error) { + if (error instanceof RangeError) throw error; + if (!isNotFound(error)) throw error; + read = null; + } + if (read) break; + // Gone for good only if the key still points at the version that vanished. + if ((await redis.get(versionKey(key))) === path) return null; + } + if (!read || !record) return null; + const { bytes, served } = read; + return { + ...record, + ...(options?.range ? { range: served } : {}), + body: new ReadableStream({ + start(controller) { + controller.enqueue(bytes.slice()); + controller.close(); + }, + }), + arrayBuffer: () => Promise.resolve(toArrayBuffer(bytes)), + text: () => Promise.resolve(decoder.decode(bytes)), + }; + }, + + head, + + async delete(key: string): Promise { + assertKey(key); + const previous = await redis.eval( + DELETE_RECORD, + [recordKey(key), indexKey, versionKey(key)], + [key], + ); + await dropVersion(previous); + }, + + async list(options?: BlobListOptions): Promise { + const limit = options?.limit; + if (limit === 0) return { objects: [], truncated: false }; + const prefixFilter = options?.prefix ?? ""; + const keys: string[] = []; + // Walk the lexical index from max(prefix, cursor), keeping keys under the prefix, until we have + // one more than the page (to know whether it is truncated) or leave the prefix range. + let min: `(${string}` | `[${string}` = + options?.cursor !== undefined && options.cursor >= prefixFilter + ? `(${options.cursor}` + : `[${prefixFilter}`; + const PAGE = 200; + for (;;) { + const batch = (await redis.zrange(indexKey, min, "+", { + byLex: true, + offset: 0, + count: PAGE, + })) as string[]; + let left = false; + for (const key of batch) { + if (!key.startsWith(prefixFilter)) { + if (key > prefixFilter) left = true; + if (left) break; + continue; + } + keys.push(key); + if (limit !== undefined && keys.length > limit) break; + } + if (left || batch.length < PAGE || (limit !== undefined && keys.length > limit)) break; + min = `(${batch[batch.length - 1]}`; + } + const truncated = limit !== undefined && keys.length > limit; + const pageKeys = truncated ? keys.slice(0, limit) : keys; + const objects = await loadDocs(redis, pageKeys.map(recordKey)); + return { + objects, + ...(truncated ? { cursor: pageKeys[pageKeys.length - 1], truncated } : {}), + }; + }, + }; +} diff --git a/packages/tanstack-ai/src/persistence/generation-stores.ts b/packages/tanstack-ai/src/persistence/generation-stores.ts new file mode 100644 index 0000000..3f3f16e --- /dev/null +++ b/packages/tanstack-ai/src/persistence/generation-stores.ts @@ -0,0 +1,189 @@ +import type { Redis } from "@upstash/redis"; +import type { + ArtifactRecord, + ArtifactStore, + GenerationRunRecord, + GenerationRunStore, +} from "@tanstack/ai-persistence"; +import { + CREATE, + KEY_LOCKING, + PATCH, + assertId, + checkRun, + loadDocs, + toMergePatches, +} from "./records.js"; + +/** + * Replace an artifact document and move it between indexes if its run or thread changed. Which + * indexes it is in is kept in a small side hash of raw index-key strings (KEYS[2]) so the document + * stays exactly the record. Under `allow-key-locking` the old index keys must be declared, so the + * caller reads them first and the script compare-and-swaps, returning 0 (retry) if a concurrent save + * moved the artifact in between. + * KEYS: document, side hash, new run index, new thread index, [old run index, old thread index]. + * ARGV: expected old run index ("" = none), expected old thread index, document JSON, createdAt, id. + */ +const SAVE_ARTIFACT = `${KEY_LOCKING}local curRun = redis.call("HGET", KEYS[2], "run") or "" +local curThread = redis.call("HGET", KEYS[2], "thread") or "" +if curRun ~= ARGV[1] or curThread ~= ARGV[2] then return 0 end +if KEYS[5] and KEYS[5] ~= KEYS[3] then redis.call("ZREM", KEYS[5], ARGV[5]) end +if KEYS[6] and KEYS[6] ~= KEYS[4] then redis.call("ZREM", KEYS[6], ARGV[5]) end +redis.call("JSON.SET", KEYS[1], "$", ARGV[3]) +redis.call("HSET", KEYS[2], "run", KEYS[3], "thread", KEYS[4]) +redis.call("ZADD", KEYS[3], ARGV[4], ARGV[5]) +redis.call("ZADD", KEYS[4], ARGV[4], ARGV[5]) +return 1`; + +/** + * Delete an artifact and unindex it, but only from the indexes it is actually in: the caller reads + * the side hash, declares those index keys, and the script compare-and-swaps (0 = a concurrent save + * changed it; retry), so a racing save can never leave a dangling index entry. An artifact that is + * not indexed ("" expected, no index keys) goes through the same check, so a first save that lands + * between the read and the delete is noticed rather than half-deleted. + * KEYS: document, side hash, [run index, thread index]. ARGV: expected run index, expected thread + * index ("" = none), id. + */ +const DELETE_ARTIFACT = `${KEY_LOCKING}local curRun = redis.call("HGET", KEYS[2], "run") or "" +local curThread = redis.call("HGET", KEYS[2], "thread") or "" +if curRun ~= ARGV[1] or curThread ~= ARGV[2] then return 0 end +redis.call("DEL", KEYS[1], KEYS[2]) +if KEYS[3] then redis.call("ZREM", KEYS[3], ARGV[3]) end +if KEYS[4] then redis.call("ZREM", KEYS[4], ARGV[3]) end +return 1`; + +/** + * `GenerationRunStore`: one record per one-shot generation job (image, video, speech, + * transcription), written by `withGenerationPersistence`. A JSON document per run and a per-thread + * sorted set by `startedAt`, so `findLatestForThread` is one index read. + */ +export function redisGenerationRunStore(redis: Redis, prefix: string): GenerationRunStore { + const run = (id: string) => `${prefix}:generationRun:${id}`; + const thread = (id: string) => `${prefix}:threadGenerationRuns:${id}`; + return { + async createOrResume(input) { + assertId(input.runId, "runId"); + assertId(input.threadId, "threadId"); + const record: GenerationRunRecord = { + runId: input.runId, + threadId: input.threadId, + activity: input.activity, + provider: input.provider, + model: input.model, + status: input.status ?? "running", + startedAt: input.startedAt, + }; + const stored = await redis.eval( + CREATE, + [run(input.runId), thread(input.threadId)], + [JSON.stringify(record), String(input.startedAt), input.runId], + ); + return checkRun(stored as GenerationRunRecord)!; + }, + + async update(runId, fields) { + assertId(runId, "runId"); + await redis.eval(PATCH, [run(runId)], toMergePatches(fields)); + }, + + async get(runId) { + assertId(runId, "runId"); + return checkRun(await redis.json.get(run(runId))); + }, + + async findLatestForThread(threadId) { + assertId(threadId, "threadId"); + const [id] = await redis.zrange(thread(threadId), 0, 0, { rev: true }); + return id === undefined ? null : checkRun(await redis.json.get(run(id))); + }, + }; +} + +/** + * `ArtifactStore`: metadata rows for generated outputs (name, MIME type, size, the `blobKey` of the + * bytes). Indexed by run and by thread, sorted by `createdAt` then id (sorted-set ties sort by bytes, + * matching the reference store). The bytes themselves live in the blob store. + */ +export function redisArtifactStore(redis: Redis, prefix: string): ArtifactStore { + const k = { + artifact: (id: string) => `${prefix}:artifact:${id}`, + where: (id: string) => `${prefix}:artifactIndexes:${id}`, + run: (id: string) => `${prefix}:runArtifacts:${id}`, + thread: (id: string) => `${prefix}:threadArtifacts:${id}`, + }; + const load = async (ids: string[]) => loadDocs(redis, ids.map(k.artifact)); + + const readIndexes = async (id: string) => { + const where = await redis.hmget>(k.where(id), "run", "thread"); + return { run: where?.run ?? null, thread: where?.thread ?? null }; + }; + + async function remove(id: string): Promise { + for (let attempt = 0; attempt < 5; attempt++) { + const { run, thread } = await readIndexes(id); + const indexed = Boolean(run && thread); + const deleted = await redis.eval( + DELETE_ARTIFACT, + [k.artifact(id), k.where(id), ...(indexed ? [run!, thread!] : [])], + [indexed ? run! : "", indexed ? thread! : "", id], + ); + if (Number(deleted) === 1) return; + } + throw new Error( + `@upstash/agentkit-tanstack-ai: artifact ${JSON.stringify(id)} kept changing during delete.`, + ); + } + + const store: ArtifactStore = { + async save(record) { + assertId(record.artifactId, "artifactId"); + assertId(record.runId, "runId"); + assertId(record.threadId, "threadId"); + const id = record.artifactId; + for (let attempt = 0; attempt < 5; attempt++) { + const { run: oldRun, thread: oldThread } = await readIndexes(id); + const saved = await redis.eval( + SAVE_ARTIFACT, + [ + k.artifact(id), + k.where(id), + k.run(record.runId), + k.thread(record.threadId), + ...(oldRun && oldThread ? [oldRun, oldThread] : []), + ], + [oldRun ?? "", oldThread ?? "", JSON.stringify(record), String(record.createdAt), id], + ); + if (Number(saved) === 1) return; + } + throw new Error( + `@upstash/agentkit-tanstack-ai: artifact ${JSON.stringify(id)} kept changing during save.`, + ); + }, + + async get(artifactId) { + assertId(artifactId, "artifactId"); + return (await redis.json.get(k.artifact(artifactId))) ?? null; + }, + + async list(runId) { + assertId(runId, "runId"); + return load(await redis.zrange(k.run(runId), 0, -1)); + }, + + async listForThread(threadId) { + assertId(threadId, "threadId"); + return load(await redis.zrange(k.thread(threadId), 0, -1)); + }, + + async delete(artifactId) { + assertId(artifactId, "artifactId"); + await remove(artifactId); + }, + + async deleteForRun(runId) { + assertId(runId, "runId"); + for (const id of await redis.zrange(k.run(runId), 0, -1)) await remove(id); + }, + }; + return store; +} diff --git a/packages/tanstack-ai/src/persistence/index.ts b/packages/tanstack-ai/src/persistence/index.ts new file mode 100644 index 0000000..db3dd5c --- /dev/null +++ b/packages/tanstack-ai/src/persistence/index.ts @@ -0,0 +1,8 @@ +// `@upstash/agentkit-tanstack-ai/persistence` — every TanStack AI persistence store on Upstash Redis +// (and Upstash Blob for generated bytes). Needs `@tanstack/ai-persistence`. +export { upstashPersistence } from "./persistence.js"; +export type { UpstashPersistenceConfig, UpstashPersistenceStores } from "./persistence.js"; + +// Blob store (bytes in Upstash Blob, records in Redis) — included by upstashPersistence({ bucket }). +export { upstashBlobStore } from "./blob-store.js"; +export type { BlobBucketLike, UpstashBlobStoreConfig } from "./blob-store.js"; diff --git a/packages/tanstack-ai/src/persistence/persistence.test.ts b/packages/tanstack-ai/src/persistence/persistence.test.ts new file mode 100644 index 0000000..c4fe9d3 --- /dev/null +++ b/packages/tanstack-ai/src/persistence/persistence.test.ts @@ -0,0 +1,279 @@ +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { runPersistenceConformance } from "@tanstack/ai-persistence/testkit"; +import type { PersistenceConformanceCheck } from "@tanstack/ai-persistence/testkit"; +import type { BlobStore } from "@tanstack/ai-persistence"; +import { Bucket } from "@upstash/blob"; +import { upstashPersistence } from "./persistence.js"; +import { testBucket, type TestBucket } from "../testing/test-bucket.js"; +import { cleanupKeys, hasRedisCreds, testRedis, uniquePrefix } from "../testing/test-support.js"; + +// Every opt-in check the kit offers (as of @tanstack/ai-persistence 0.7.1); none is skipped. +const OPT_IN_CHECKS: PersistenceConformanceCheck[] = [ + "messages.metadata", + "runs.listByThread.state", +]; + +const hasBlobToken = Boolean(process.env.UPSTASH_BLOB_TOKEN); + +// TanStack AI's own backend conformance suite, run against a real Upstash Redis, with every store +// present and none skipped: messages, runs, interrupts, metadata, generationRuns, artifacts, blobs. +// Blob bytes go to a stand-in bucket whose signed URLs are served over real HTTP with Range support. +describe.skipIf(!hasRedisCreds)("upstashPersistence (live Redis)", () => { + const redis = testRedis(); + const prefix = uniquePrefix("tsp"); + let bucket: TestBucket; + + beforeAll(async () => { + bucket = await testBucket(); + }); + + afterAll(async () => { + await bucket.close(); + await cleanupKeys(redis, prefix); + }); + + let n = 0; + runPersistenceConformance( + "upstashPersistence", + // A fresh keyspace per case: the suite assumes each persistence starts empty. + () => + upstashPersistence({ redis, prefix: `${prefix}:${++n}`, bucket, blobPathPrefix: `t/${n}/` }), + { checks: OPT_IN_CHECKS }, + ); + + it("serves byte ranges with an HTTP Range request, and cleans up bucket objects on delete", async () => { + const { blobs } = upstashPersistence({ + redis, + prefix: `${prefix}:range`, + bucket, + blobPathPrefix: "r/", + }).stores; + await blobs!.put("video.bin", "0123456789", { contentType: "application/octet-stream" }); + const before = bucket.rangeHits; + const slice = await blobs!.get("video.bin", { range: { offset: 3, length: 4 } }); + expect(await slice!.text()).toBe("3456"); + expect(slice!.range).toEqual({ offset: 3, length: 4 }); + expect(bucket.rangeHits).toBe(before + 1); + const version = (p: string) => p.startsWith("r/video.bin/"); + expect(bucket.paths().filter(version)).toHaveLength(1); + await blobs!.delete("video.bin"); + expect(bucket.paths().filter(version)).toHaveLength(0); + expect(await blobs!.head("video.bin")).toBeNull(); + }); + + it("concurrent puts of one key never leave a record describing another writer's bytes", async () => { + const { blobs } = upstashPersistence({ + redis, + prefix: `${prefix}:race`, + bucket, + blobPathPrefix: "race/", + }).stores; + const bodies = ["a", "bb", "ccc", "dddd", "eeeee"]; + await Promise.all(bodies.map((body) => blobs!.put("k", body))); + const got = await blobs!.get("k"); + const text = await got!.text(); + expect(bodies).toContain(text); + expect(got!.size).toBe(text.length); // record and bytes are the same version + // Every replaced version was cleaned up: exactly one object remains for the key. + expect(bucket.paths().filter((p) => p.startsWith("race/k/"))).toHaveLength(1); + }); + + it("deleting an artifact while it is re-saved elsewhere leaves no dangling index entry", async () => { + const { artifacts } = upstashPersistence({ redis, prefix: `${prefix}:artdel` }).stores; + const base = { artifactId: "x", name: "n", mimeType: "text/plain", size: 1, createdAt: 1 }; + await artifacts.save({ ...base, runId: "r1", threadId: "t1" }); + await Promise.all([ + artifacts.delete("x"), + artifacts.save({ ...base, runId: "r2", threadId: "t2" }), + ]); + const record = await artifacts.get("x"); + const listed = [ + ...(await artifacts.list("r1")), + ...(await artifacts.list("r2")), + ...(await artifacts.listForThread("t1")), + ...(await artifacts.listForThread("t2")), + ]; + if (record) { + // The save won: it is indexed exactly where the record says. + expect(listed.map((a) => a.runId + "/" + a.threadId)).toEqual(["r2/t2", "r2/t2"]); + } else { + // The delete won: nothing points at it anywhere. + expect(listed).toEqual([]); + expect(await redis.zcard(`${prefix}:artdel:runArtifacts:r2`)).toBe(0); + } + }); + + it("get never reports a live key missing when an overwrite deletes the version it was reading", async () => { + // A bucket whose first read of the original version is preceded by a concurrent overwrite, which + // commits a new version and deletes the one this get had just resolved. + let raced = false; + let blobs: BlobStore; + const racingBucket: typeof bucket = { + ...bucket, + put: bucket.put.bind(bucket), + del: bucket.del.bind(bucket), + get: async (path: string) => { + if (!raced && path.startsWith("racy/k/")) { + raced = true; + await blobs.put("k", "second"); + } + return bucket.get(path); + }, + } as typeof bucket; + blobs = upstashPersistence({ + redis, + prefix: `${prefix}:racy`, + bucket: racingBucket, + blobPathPrefix: "racy/", + }).stores.blobs!; + await blobs.put("k", "first"); + const got = await blobs.get("k"); + expect(raced).toBe(true); + expect(got).not.toBeNull(); + expect(await got!.text()).toBe("second"); + }); + + it("omits the blobs store when no bucket is given", () => { + const stores = upstashPersistence({ redis, prefix: `${prefix}:nob` }).stores; + expect(stores.blobs).toBeUndefined(); + expect(stores.generationRuns).toBeDefined(); + expect(stores.artifacts).toBeDefined(); + }); +}); + +// The same suite against a real Upstash Blob bucket. Needs UPSTASH_BLOB_TOKEN (skipped without it). +describe.skipIf(!hasRedisCreds || !hasBlobToken)( + "upstashPersistence (live Redis + live Upstash Blob)", + () => { + const redis = testRedis(); + const prefix = uniquePrefix("tspblob"); + const bucket = hasBlobToken ? Bucket.fromEnv() : (undefined as never); + const pathPrefix = `agentkit-test/${prefix.replace(/:/g, "-")}/`; + + afterAll(async () => { + await cleanupKeys(redis, prefix); + await bucket.del({ prefix: pathPrefix, all: true }).catch(() => undefined); + }); + + let n = 0; + runPersistenceConformance( + "upstashPersistence + Upstash Blob", + () => + upstashPersistence({ + redis, + prefix: `${prefix}:${++n}`, + bucket, + blobPathPrefix: `${pathPrefix}${n}/`, + }), + { checks: OPT_IN_CHECKS }, + ); + }, +); + +// Every Lua script runs with `allow-key-locking`, under which touching an undeclared key is an error. +// These drive the two paths whose key sets vary: a run with a parent index, and an artifact moved +// between runs and threads (whose old indexes must be declared, not read inside the script). +describe.skipIf(!hasRedisCreds)("key-locking scripts (live Redis)", () => { + const redis = testRedis(); + const prefix = uniquePrefix("tskl"); + afterAll(async () => { + await cleanupKeys(redis, prefix); + }); + + it("creates runs with and without a parent index", async () => { + const { runs } = upstashPersistence({ redis, prefix }).stores; + await runs.createOrResume({ runId: "p", threadId: "t", startedAt: 1 }); + await runs.createOrResume({ + runId: "c", + threadId: "subagent:c", + startedAt: 2, + parentRunId: "p", + }); + expect((await runs.listByParentRun!("p")).map((r) => r.runId)).toEqual(["c"]); + expect(await redis.exists(`${prefix}:parentRuns:_`)).toBe(0); + }); + + it("keeps the detached-runs index in step with the run in the same script", async () => { + const { runs } = upstashPersistence({ redis, prefix }).stores; + const detached = `${prefix}:detachedRuns`; + const since = 1_727_000_000_123; + await runs.createOrResume({ runId: "d", threadId: "t", startedAt: 1 }); + + await runs.update("d", { detachedSince: since }); + expect(await redis.zscore(detached, "d")).toBe(since); + + await runs.update("d", { detachedSince: undefined }); + expect(await redis.zscore(detached, "d")).toBeNull(); + + await runs.update("d", { detachedSince: since }); + await runs.update("d", { status: "completed", finishedAt: since + 1 }); + expect(await redis.zscore(detached, "d")).toBeNull(); + + await runs.update("missing", { detachedSince: since }); + expect(await redis.zscore(detached, "missing")).toBeNull(); + expect(await runs.get("missing")).toBeNull(); + }); + + it("a detach is indexed by the same request that writes it", async () => { + // A client that dies after its first request: any follow-up call fails, as it would if the + // process crashed or ran out of retries between the patch and a separate index write. + let used = false; + const gone = () => { + throw new Error("client gone"); + }; + const oneShot = new Proxy(redis, { + get(target, prop) { + if (prop === "eval") { + return (...args: Parameters) => { + if (used) gone(); + used = true; + return target.eval(...args); + }; + } + if (used && (prop === "zadd" || prop === "zrem" || prop === "json")) gone(); + return Reflect.get(target, prop, target); + }, + }); + const { runs } = upstashPersistence({ redis, prefix }).stores; + await runs.createOrResume({ runId: "crash", threadId: "t", startedAt: 1 }); + + const dying = upstashPersistence({ redis: oneShot, prefix }).stores.runs; + await dying.update("crash", { detachedSince: 5_000 }); + + expect((await runs.get("crash"))?.detachedSince).toBe(5_000); + expect( + (await runs.listReclaimable!({ now: 10_000, ttlMs: 1_000 })).map((r) => r.runId), + ).toContain("crash"); + }); + + it("moves a re-saved artifact between run and thread indexes", async () => { + const { artifacts } = upstashPersistence({ redis, prefix }).stores; + const base = { + artifactId: "a1", + name: "img.png", + mimeType: "image/png", + size: 3, + createdAt: 5, + }; + await artifacts.save({ ...base, runId: "r1", threadId: "t1" }); + await artifacts.save({ ...base, runId: "r2", threadId: "t2" }); + expect(await artifacts.list("r1")).toEqual([]); + expect(await artifacts.listForThread("t1")).toEqual([]); + expect((await artifacts.list("r2")).map((a) => a.artifactId)).toEqual(["a1"]); + expect((await artifacts.listForThread("t2")).map((a) => a.runId)).toEqual(["r2"]); + }); + + it("concurrent saves of one artifact settle on a single indexed copy", async () => { + const { artifacts } = upstashPersistence({ redis, prefix }).stores; + const base = { artifactId: "a2", name: "x", mimeType: "text/plain", size: 1, createdAt: 1 }; + await Promise.all( + ["ra", "rb", "rc"].map((runId) => artifacts.save({ ...base, runId, threadId: `t-${runId}` })), + ); + const where = await Promise.all(["ra", "rb", "rc"].map((r) => artifacts.list(r))); + expect(where.flat()).toHaveLength(1); + const final = (await artifacts.get("a2"))!; + expect((await artifacts.listForThread(final.threadId)).map((a) => a.artifactId)).toEqual([ + "a2", + ]); + }); +}); diff --git a/packages/tanstack-ai/src/persistence/persistence.ts b/packages/tanstack-ai/src/persistence/persistence.ts new file mode 100644 index 0000000..22d8647 --- /dev/null +++ b/packages/tanstack-ai/src/persistence/persistence.ts @@ -0,0 +1,355 @@ +import type { Redis } from "@upstash/redis"; +import { Redis as RedisClient } from "@upstash/redis"; +import type { ModelMessage, RunRecord, RunStore } from "@tanstack/ai"; +import type { + AIPersistence, + ArtifactStore, + BlobStore, + ChatPersistenceStores, + GenerationRunStore, + InterruptRecord, + InterruptStore, + MessageStore, + MetadataStore, +} from "@tanstack/ai-persistence"; +import { addTelemetry } from "../telemetry.js"; +import { + CREATE, + KEY_LOCKING, + PATCH, + assertId, + checkRun, + loadDocs, + toMergePatches, +} from "./records.js"; +import { redisArtifactStore, redisGenerationRunStore } from "./generation-stores.js"; +import { upstashBlobStore } from "./blob-store.js"; +import type { BlobBucketLike } from "./blob-store.js"; + +export interface UpstashPersistenceConfig { + /** Upstash Redis client. Defaults to `Redis.fromEnv()`. */ + redis?: Redis; + /** Base key prefix. Defaults to `agentkit:tanstack`. */ + prefix?: string; + /** + * Expire a thread's stored transcript this many seconds after its last save. Omit to keep + * transcripts forever. Run, interrupt and metadata records are never expired automatically. + */ + messagesTtlSeconds?: number; + /** + * An Upstash Blob bucket (e.g. `Bucket.fromEnv()` from `@upstash/blob`). When given, a `blobs` + * store is included: generated bytes (images, audio, video) go to the bucket, their records to + * Redis. Without it the persistence has no `blobs` store. + */ + bucket?: BlobBucketLike; + /** + * Path prefix for objects written to `bucket`. + * @default "agentkit/tanstack/" + */ + blobPathPrefix?: string; + /** + * Report the sdk name + version to Upstash as a header on the requests made by your redis client. + * Can also be disabled with the `UPSTASH_DISABLE_TELEMETRY` env var. Defaults to `true`. + */ + enableTelemetry?: boolean; +} + +/** The stores {@link upstashPersistence} returns; `blobs` is present when a `bucket` is passed. */ +export type UpstashPersistenceStores = ChatPersistenceStores & { + generationRuns: GenerationRunStore; + artifacts: ArtifactStore; + blobs?: BlobStore; +}; + +/** + * Settle a batch of interrupts atomically: every document must exist and be pending, or nothing + * changes. ARGV[i] is the merge patch for KEYS[i]. + */ +const COMMIT_BATCH = `${KEY_LOCKING}for i = 1, #KEYS do + local status = redis.call("JSON.GET", KEYS[i], "$.status") + if not status then return redis.error_reply("missing:" .. i) end + if status ~= '["pending"]' then return redis.error_reply("nonpending:" .. i) end +end +for i = 1, #KEYS do redis.call("JSON.MERGE", KEYS[i], "$", ARGV[i]) end +return #KEYS`; + +/** + * {@link PATCH} for a run that also keeps the detached-runs index in step with the patched record, in + * the same script, so a crash or a racing update can never leave a detached run out of the index. + * KEYS: run document, detached index. ARGV: the two merge patches, runId. + */ +const UPDATE_RUN = `${KEY_LOCKING}if redis.call("EXISTS", KEYS[1]) == 0 then return 0 end +redis.call("JSON.MERGE", KEYS[1], "$", ARGV[1]) +if ARGV[2] ~= "{}" then redis.call("JSON.MERGE", KEYS[1], "$", ARGV[2]) end +local status = redis.call("JSON.GET", KEYS[1], "$.status") +local since = string.sub(redis.call("JSON.GET", KEYS[1], "$.detachedSince"), 2, -2) +if status == '["running"]' and tonumber(since) then + redis.call("ZADD", KEYS[2], since, ARGV[3]) +else + redis.call("ZREM", KEYS[2], ARGV[3]) +end +return 1`; + +/** + * Every TanStack AI persistence store on Upstash Redis — pass the result to `withPersistence()` or + * `withGenerationPersistence()`. + * + * - **messages**: the thread transcript, one JSON document per thread (`saveThread` overwrites). + * - **runs**: a JSON document per run plus sorted-set indexes by thread, by parent run and by detach + * time, so `findActiveRun`, `listByThread`, `listByParentRun` and `listReclaimable` are index reads. + * - **interrupts**: human-in-the-loop pauses, indexed by thread and run; `commitBatch` is atomic. + * - **metadata**: one document per `(namespace, key)`, both parts escaped so they never collide. + * - **generationRuns** / **artifacts**: one-shot generation jobs (image, video, speech) and records of + * what they produced. + * - **blobs** (when `bucket` is given): the produced bytes in Upstash Blob, records in Redis. + * + * Records are RedisJSON documents written by single commands or single Lua scripts, so concurrent + * instances cannot interleave a write. Checked against TanStack's `runPersistenceConformance` suite. + * + * ```ts + * import { withPersistence } from "@tanstack/ai-persistence"; + * import { upstashPersistence } from "@upstash/agentkit-tanstack-ai/persistence"; + * + * chat({ adapter, messages, threadId, middleware: [withPersistence(upstashPersistence())] }); + * ``` + */ +export function upstashPersistence( + config: UpstashPersistenceConfig = {}, +): AIPersistence { + const redis = config.redis ?? RedisClient.fromEnv(); + addTelemetry(redis, config.enableTelemetry); + const p = config.prefix ?? "agentkit:tanstack"; + const k = { + thread: (id: string) => `${p}:thread:${id}`, + run: (id: string) => `${p}:run:${id}`, + threadRuns: (id: string) => `${p}:threadRuns:${id}`, + parentRuns: (id: string) => `${p}:parentRuns:${id}`, + detached: `${p}:detachedRuns`, + interrupt: (id: string) => `${p}:interrupt:${id}`, + threadInterrupts: (id: string) => `${p}:threadInterrupts:${id}`, + runInterrupts: (id: string) => `${p}:runInterrupts:${id}`, + // encodeURIComponent never emits ":", so (namespace, key) pairs cannot collide. + meta: (ns: string, key: string) => + `${p}:meta:${encodeURIComponent(ns)}:${encodeURIComponent(key)}`, + }; + + const loadRuns = async (ids: string[]) => + (await loadDocs(redis, ids.map(k.run))).map((r) => checkRun(r)!); + const loadInterrupts = (ids: string[]) => loadDocs(redis, ids.map(k.interrupt)); + const patch = (key: string, fields: object) => redis.eval(PATCH, [key], toMergePatches(fields)); + + const messages: MessageStore = { + loadThread: (async (threadId: string) => { + assertId(threadId, "threadId"); + return (await redis.json.get(k.thread(threadId))) ?? []; + }) as MessageStore["loadThread"], + async saveThread(threadId, list) { + assertId(threadId, "threadId"); + const tx = redis.multi(); + tx.json.set(k.thread(threadId), "$", list as unknown as Record); + if (config.messagesTtlSeconds) tx.expire(k.thread(threadId), config.messagesTtlSeconds); + await tx.exec(); + }, + }; + + const runs: RunStore = { + async createOrResume(input) { + assertId(input.runId, "runId"); + assertId(input.threadId, "threadId"); + const record: RunRecord = { + runId: input.runId, + threadId: input.threadId, + status: input.status ?? "running", + startedAt: input.startedAt, + ...(input.parentRunId !== undefined ? { parentRunId: input.parentRunId } : {}), + ...(input.subagentRunId !== undefined ? { subagentRunId: input.subagentRunId } : {}), + ...(input.name !== undefined ? { name: input.name } : {}), + }; + const indexes = [k.threadRuns(input.threadId)]; + if (input.parentRunId !== undefined) indexes.push(k.parentRuns(input.parentRunId)); + const stored = await redis.eval( + CREATE, + [k.run(input.runId), ...indexes], + [JSON.stringify(record), String(input.startedAt), input.runId], + ); + return checkRun(stored as RunRecord)!; + }, + + async update(runId, fields) { + assertId(runId, "runId"); + if ("detachedSince" in fields || "status" in fields) { + await redis.eval( + UPDATE_RUN, + [k.run(runId), k.detached], + [...toMergePatches(fields), runId], + ); + } else { + await patch(k.run(runId), fields); + } + }, + + async get(runId) { + assertId(runId, "runId"); + return checkRun(await redis.json.get(k.run(runId))); + }, + + async findActiveRun(threadId) { + assertId(threadId, "threadId"); + const PAGE = 50; + for (let start = 0; ; start += PAGE) { + const ids = await redis.zrange(k.threadRuns(threadId), start, start + PAGE - 1, { + rev: true, + }); + const active = (await loadRuns(ids)).find((r) => r.status === "running"); + if (active) return active; + if (ids.length < PAGE) return null; + } + }, + + async listByThread(threadId) { + assertId(threadId, "threadId"); + return loadRuns(await redis.zrange(k.threadRuns(threadId), 0, -1)); + }, + + async listByParentRun(parentRunId) { + assertId(parentRunId, "parentRunId"); + return loadRuns(await redis.zrange(k.parentRuns(parentRunId), 0, -1)); + }, + + async listReclaimable({ now, ttlMs }) { + const cutoff = now - ttlMs; + const ids = await redis.zrange(k.detached, "-inf", cutoff, { byScore: true }); + return (await loadRuns(ids)).filter( + (r) => r.status === "running" && r.detachedSince !== undefined && r.detachedSince <= cutoff, + ); + }, + }; + + const interrupts: InterruptStore = { + async create(record) { + assertId(record.interruptId, "interruptId"); + await redis.eval( + CREATE, + [ + k.interrupt(record.interruptId), + k.threadInterrupts(record.threadId), + k.runInterrupts(record.runId), + ], + [ + JSON.stringify({ ...record, status: "pending" }), + String(record.requestedAt), + record.interruptId, + ], + ); + }, + + async resolve(interruptId, response) { + assertId(interruptId, "interruptId"); + await patch(k.interrupt(interruptId), { + status: "resolved", + resolvedAt: Date.now(), + response, + }); + }, + + async cancel(interruptId) { + assertId(interruptId, "interruptId"); + await patch(k.interrupt(interruptId), { status: "cancelled", resolvedAt: Date.now() }); + }, + + async commitBatch(entries) { + if (entries.length === 0) return; + const seen = new Set(); + for (const entry of entries) { + assertId(entry.interruptId, "interruptId"); + if (seen.has(entry.interruptId)) { + throw new Error(`Interrupt batch contains duplicate id: ${entry.interruptId}.`); + } + seen.add(entry.interruptId); + } + const resolvedAt = Date.now(); + // One merge patch per entry. A cancel clears any earlier response (`null` deletes the field). + const patches = entries.map((e) => + JSON.stringify( + e.status === "resolved" + ? { status: e.status, resolvedAt, response: e.response ?? null } + : { status: e.status, resolvedAt, response: null }, + ), + ); + try { + await redis.eval( + COMMIT_BATCH, + entries.map((e) => k.interrupt(e.interruptId)), + patches, + ); + } catch (error) { + const match = /(missing|nonpending):(\d+)/.exec(String((error as Error)?.message ?? error)); + if (!match) throw error; + const id = entries[Number(match[2]) - 1]?.interruptId; + throw new Error( + match[1] === "missing" + ? `Interrupt batch references missing id: ${id}.` + : `Interrupt batch references non-pending id: ${id}.`, + ); + } + }, + + async get(interruptId) { + assertId(interruptId, "interruptId"); + return (await redis.json.get(k.interrupt(interruptId))) ?? null; + }, + + async list(threadId) { + assertId(threadId, "threadId"); + return loadInterrupts(await redis.zrange(k.threadInterrupts(threadId), 0, -1)); + }, + + async listPending(threadId) { + return (await interrupts.list(threadId)).filter((r) => r.status === "pending"); + }, + + async listByRun(runId) { + assertId(runId, "runId"); + return loadInterrupts(await redis.zrange(k.runInterrupts(runId), 0, -1)); + }, + + async listPendingByRun(runId) { + return (await interrupts.listByRun(runId)).filter((r) => r.status === "pending"); + }, + }; + + const metadata: MetadataStore = { + // Values are wrapped (`{ v }`) so any JSON value, including a bare string, round-trips exactly. + async get(namespace, key) { + const doc = await redis.json.get<{ v: unknown }>(k.meta(namespace, key)); + return doc ? doc.v : null; + }, + async set(namespace, key, value) { + await redis.json.set(k.meta(namespace, key), "$", { v: value ?? null }); + }, + async delete(namespace, key) { + await redis.del(k.meta(namespace, key)); + }, + }; + + const blobs = config.bucket + ? upstashBlobStore({ + bucket: config.bucket, + redis, + prefix: p, + ...(config.blobPathPrefix !== undefined ? { pathPrefix: config.blobPathPrefix } : {}), + }) + : undefined; + + return { + stores: { + messages, + runs, + interrupts, + metadata, + generationRuns: redisGenerationRunStore(redis, p), + artifacts: redisArtifactStore(redis, p), + ...(blobs ? { blobs } : {}), + }, + } as AIPersistence; +} diff --git a/packages/tanstack-ai/src/persistence/records.ts b/packages/tanstack-ai/src/persistence/records.ts new file mode 100644 index 0000000..2416164 --- /dev/null +++ b/packages/tanstack-ai/src/persistence/records.ts @@ -0,0 +1,89 @@ +import type { Redis } from "@upstash/redis"; +import type { RunStatus } from "@tanstack/ai"; + +/** + * Shared plumbing for the persistence stores. Records are RedisJSON documents, so values keep their + * JSON types natively (no per-field encoding) and patches are applied server-side. + * + * Every script runs with `allow-key-locking`: Upstash locks only the keys a script declares, and an + * undeclared key is an error, so each script touches exactly its `KEYS`. + */ +export const KEY_LOCKING = "#!lua flags=allow-key-locking\n"; + +/** + * Create a document only if it does not exist yet, add it to every index in KEYS[2..], and return the + * stored document either way (the idempotent `createOrResume` contract). + * KEYS: document, indexes… ARGV: document JSON, index score, index member. + */ +export const CREATE = `${KEY_LOCKING}if redis.call("JSON.SET", KEYS[1], "$", ARGV[1], "NX") then + for i = 2, #KEYS do redis.call("ZADD", KEYS[i], ARGV[2], ARGV[3]) end +end +return redis.call("JSON.GET", KEYS[1])`; + +/** + * Patch an existing document with two JSON merge patches (see {@link toMergePatches}). A missing + * document is a no-op: `JSON.MERGE` would otherwise create it. + */ +export const PATCH = `${KEY_LOCKING}if redis.call("EXISTS", KEYS[1]) == 0 then return 0 end +redis.call("JSON.MERGE", KEYS[1], "$", ARGV[1]) +if ARGV[2] ~= "{}" then redis.call("JSON.MERGE", KEYS[1], "$", ARGV[2]) end +return 1`; + +const isPlainObject = (v: unknown): v is Record => + typeof v === "object" && v !== null && !Array.isArray(v); + +/** + * Turn a `{ ...existing, ...patch }`-style patch into merge patches for {@link PATCH}. A merge patch + * deletes a field set to `null` (how a key present with `undefined` — "clear it" — is expressed) but + * merges nested objects instead of replacing them, so object-valued fields are cleared by the first + * patch and set by the second, giving replace semantics. + */ +export function toMergePatches(patch: object): [string, string] { + const first: Record = {}; + const second: Record = {}; + for (const [field, value] of Object.entries(patch)) { + if (value === undefined) first[field] = null; + else if (isPlainObject(value)) { + first[field] = null; + second[field] = value; + } else first[field] = value; + } + return [JSON.stringify(first), JSON.stringify(second)]; +} + +/** Read several documents in one round trip, in order, dropping missing ones. */ +export async function loadDocs(redis: Redis, keys: string[]): Promise { + if (keys.length === 0) return []; + const pipe = redis.pipeline(); + for (const key of keys) pipe.json.get(key); + const rows = (await pipe.exec()) as (T | null)[]; + return rows.filter((row): row is T => row !== null && row !== undefined); +} + +/** Reject ids that would break a key: empty, or carrying CR/LF. */ +export function assertId(value: string, name: string): void { + if (typeof value !== "string" || value === "" || /[\r\n]/.test(value)) { + throw new Error( + `@upstash/agentkit-tanstack-ai: \`${name}\` must be a non-empty string without CR/LF.`, + ); + } +} + +const RUN_STATUSES: ReadonlySet = new Set([ + "running", + "interrupted", + "completed", + "failed", + "aborted", +]); + +/** Validate a stored run's status at read time — TanStack's readers act destructively on it. */ +export function checkRun(record: T | null): T | null { + if (record === null || record === undefined) return null; + if (typeof record.status !== "string" || !RUN_STATUSES.has(record.status)) { + throw new Error( + `@upstash/agentkit-tanstack-ai: run ${JSON.stringify(record.runId)} has an invalid status.`, + ); + } + return record; +} diff --git a/packages/tanstack-ai/src/search/search-tools.test.ts b/packages/tanstack-ai/src/search/search-tools.test.ts new file mode 100644 index 0000000..4daa669 --- /dev/null +++ b/packages/tanstack-ai/src/search/search-tools.test.ts @@ -0,0 +1,74 @@ +import { s } from "@upstash/redis"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { chat } from "@tanstack/ai"; +import type { StreamChunk } from "@tanstack/ai"; +import { createSearchTools } from "./search-tools.js"; +import { scriptedAdapter } from "../testing/test-adapter.js"; +import { hasRedisCreds, testRedis, uniquePrefix } from "../testing/test-support.js"; + +async function drain(stream: unknown): Promise { + const out: StreamChunk[] = []; + for await (const c of stream as AsyncIterable) out.push(c); + return out; +} + +const schema = s.object({ name: s.string(), price: s.number(), category: s.string().noTokenize() }); + +describe.skipIf(!hasRedisCreds)("createSearchTools (live Redis, real chat loop)", () => { + const redis = testRedis(); + const name = uniquePrefix("tssearch").replace(/[^a-zA-Z0-9_]/g, "_"); + const prefix = `${name}:`; + const tools = createSearchTools({ schema, redis, indexName: name, prefix }); + const byName = Object.fromEntries(tools.map((t) => [(t as { name: string }).name, t])); + + beforeAll(async () => { + // Provision before seeding, then seed and wait (provision -> write -> waitIndexing -> read). + await (byName.count as unknown as { execute: (i: unknown) => Promise }).execute({ + filter: { category: { $eq: "none" } }, + }); + await redis.json.set(`${prefix}1`, "$", { + name: "Wireless headphones", + price: 99, + category: "audio", + }); + await redis.json.set(`${prefix}2`, "$", { name: "Desk lamp", price: 25, category: "home" }); + await redis.search.index({ name }).waitIndexing(); + }); + + afterAll(async () => { + try { + await redis.search.index({ name }).drop(); + } catch { + /* may not exist */ + } + await redis.del(`${prefix}1`, `${prefix}2`); + }); + + it("exposes search / aggregate / count with schema-aware descriptions", () => { + expect(Object.keys(byName).sort()).toEqual(["aggregate", "count", "search"]); + expect((byName.search as { description: string }).description).toContain("`price` (F64)"); + }); + + it("the model's search call runs a typo-tolerant query and the result reaches the next turn", async () => { + const adapter = scriptedAdapter([ + { + toolCalls: [ + { id: "q1", name: "search", args: { filter: { name: { $smart: "hedphones" } } } }, + ], + }, + { text: "Found them." }, + ]); + await drain( + chat({ + adapter: adapter as never, + messages: [{ role: "user", content: "headphones?" }], + tools, + }), + ); + const toolMsg = (adapter.calls[1]!.messages as { role: string; content: string }[]).find( + (m) => m.role === "tool", + ); + expect(toolMsg!.content).toContain("Wireless headphones"); + expect(toolMsg!.content).not.toContain("Desk lamp"); + }); +}); diff --git a/packages/tanstack-ai/src/search/search-tools.ts b/packages/tanstack-ai/src/search/search-tools.ts new file mode 100644 index 0000000..5ff045f --- /dev/null +++ b/packages/tanstack-ai/src/search/search-tools.ts @@ -0,0 +1,52 @@ +import type { Redis } from "@upstash/redis"; +import { Redis as RedisClient } from "@upstash/redis"; +import { createSearchToolDefs } from "@upstash/agentkit-sdk"; +import type { AnySearchSchema, SearchToolDefsConfig } from "@upstash/agentkit-sdk"; +import { toolDefinition } from "@tanstack/ai"; +import type { Tool } from "@tanstack/ai"; +import { addTelemetry } from "../telemetry.js"; + +/** + * The core {@link SearchToolDefsConfig} (`schema`, `indexName`, `prefix`, `defaultLimit`, + * `enableTelemetry`), with `redis` optional, plus the tool names. + */ +export type CreateSearchToolsConfig = Omit< + SearchToolDefsConfig, + "redis" +> & { + /** Upstash Redis client. Defaults to `Redis.fromEnv()`. */ + redis?: Redis; + /** Tool names. Defaults to `search`, `aggregate` and `count`. */ + names?: { search?: string; aggregate?: string; count?: string }; +}; + +/** + * Schema-driven Redis Search tools for TanStack AI — `search`, `aggregate` and `count` over one + * Upstash Redis Search index, for RAG over your own documents. The tool descriptions are generated + * from the schema (fields, types, and the operators each accepts), and the index is created on the + * first read. Returns server tools ready for `chat({ tools })`. + * + * ```ts + * const tools = createSearchTools({ + * indexName: "products", + * schema: s.object({ name: s.string(), price: s.number(), inStock: s.boolean() }), + * }); + * chat({ adapter, messages, tools }); + * ``` + */ +export function createSearchTools( + config: CreateSearchToolsConfig, +): Tool[] { + const redis = config.redis ?? RedisClient.fromEnv(); + addTelemetry(redis, config.enableTelemetry); + const { names: toolNames, ...defsConfig } = config; + const defs = createSearchToolDefs({ ...defsConfig, redis }); + const names = { search: "search", aggregate: "aggregate", count: "count", ...toolNames }; + return (["search", "aggregate", "count"] as const).map((key) => + toolDefinition({ + name: names[key], + description: defs[key].description, + inputSchema: defs[key].inputSchema, + }).server(async (input: unknown) => defs[key].execute(input as Record)), + ); +} diff --git a/packages/tanstack-ai/src/stream/event-log.test.ts b/packages/tanstack-ai/src/stream/event-log.test.ts new file mode 100644 index 0000000..13a7f85 --- /dev/null +++ b/packages/tanstack-ai/src/stream/event-log.test.ts @@ -0,0 +1,116 @@ +import { afterAll, describe, expect, it } from "vitest"; +import { EventLog, EventLogClosedError } from "./event-log.js"; +import { cleanupKeys, hasRedisCreds, testRedis, uniquePrefix } from "../testing/test-support.js"; + +const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)); + +describe.skipIf(!hasRedisCreds)("EventLog (live Redis)", () => { + const redis = testRedis(); + const prefix = uniquePrefix("log"); + + afterAll(async () => { + await cleanupKeys(redis, prefix); + }); + + it("round-trips values exactly, including JSON-looking strings and numbers", async () => { + const log = new EventLog({ redis, prefix }); + const events = [{ type: "A", n: 1 }, "123", '{"a":1}', 42, null, [1, "x"]]; + const ids = await log.append("exact", events); + expect(ids).toHaveLength(events.length); + expect((await log.snapshot("exact")).map((e) => e.event)).toEqual(events); + }); + + it("snapshot is [] for an unknown log and never waits", async () => { + const log = new EventLog({ redis, prefix }); + expect(await log.snapshot("nope")).toEqual([]); + expect(await log.tail("nope")).toBeNull(); + }); + + it("read resumes strictly after a position and stops once closed", async () => { + const log = new EventLog({ redis, prefix, pollIntervalMs: 50 }); + const ids = await log.append("resume", [1, 2, 3]); + const seen: number[] = []; + const reader = (async () => { + for await (const e of log.read("resume", { after: ids[0] })) seen.push(e.event); + })(); + await sleep(150); + await log.append("resume", [4, 5]); + await log.close("resume"); + await reader; + expect(seen).toEqual([2, 3, 4, 5]); + expect(await log.isClosed("resume")).toBe(true); + }); + + it("two readers see the same live stream (fan-out)", async () => { + const log = new EventLog({ redis, prefix, pollIntervalMs: 50 }); + const collect = async () => { + const out: string[] = []; + for await (const e of log.read("fan")) out.push(e.event); + return out; + }; + const a = collect(); + const b = collect(); + await log.append("fan", ["x"]); + await sleep(120); + await log.append("fan", ["y", "z"]); + await log.close("fan"); + expect(await a).toEqual(["x", "y", "z"]); + expect(await b).toEqual(["x", "y", "z"]); + }); + + it("stops tailing when the signal aborts", async () => { + const log = new EventLog({ redis, prefix, pollIntervalMs: 50 }); + await log.append("abort", [1]); + const controller = new AbortController(); + const seen: number[] = []; + const reader = (async () => { + for await (const e of log.read("abort", { signal: controller.signal })) seen.push(e.event); + })(); + await sleep(150); + controller.abort(); + await reader; + expect(seen).toEqual([1]); + }); + + it("rejects a join on a log that never produces", async () => { + const log = new EventLog({ redis, prefix, pollIntervalMs: 50 }); + const read = async () => { + for await (const _ of log.read("ghost", { firstEntryTimeoutMs: 200 })) void _; + }; + await expect(read()).rejects.toThrow(/no entries/); + }); + + it("pages through logs larger than one XRANGE page", async () => { + const log = new EventLog({ redis, prefix }); + const big = Array.from({ length: 1_200 }, (_, i) => i); + await log.append("big", big.slice(0, 600)); + await log.append("big", big.slice(600)); + expect((await log.snapshot("big")).map((e) => e.event)).toEqual(big); + await log.close("big"); + const read: number[] = []; + for await (const e of log.read("big")) read.push(e.event); + expect(read).toEqual(big); + }); + + it("refuses appends after close, so nothing lands after readers' final drain", async () => { + const log = new EventLog({ redis, prefix, ttlSeconds: 60 }); + await log.append("closed", [1]); + await log.close("closed"); + await expect(log.append("closed", [2])).rejects.toBeInstanceOf(EventLogClosedError); + expect((await log.snapshot("closed")).map((e) => e.event)).toEqual([1]); + // The flag and the stream keep the same expiry, so the flag cannot lapse first. + const [streamTtl, flagTtl] = await Promise.all([ + redis.ttl(`${prefix}:events:closed`), + redis.ttl(`${prefix}:closed:closed`), + ]); + expect(Math.abs(streamTtl - flagTtl)).toBeLessThanOrEqual(1); + }); + + it("expires logs whose producer never closed", async () => { + const log = new EventLog({ redis, prefix, ttlSeconds: 60 }); + await log.append("ttl", [1]); + const ttl = await redis.ttl(`${prefix}:events:ttl`); + expect(ttl).toBeGreaterThan(0); + expect(ttl).toBeLessThanOrEqual(60); + }); +}); diff --git a/packages/tanstack-ai/src/stream/event-log.ts b/packages/tanstack-ai/src/stream/event-log.ts new file mode 100644 index 0000000..21f451d --- /dev/null +++ b/packages/tanstack-ai/src/stream/event-log.ts @@ -0,0 +1,255 @@ +import type { Redis } from "@upstash/redis"; +import { addTelemetry } from "../telemetry.js"; +import { KEY_LOCKING } from "../persistence/records.js"; + +/** + * Every stored event carries this prefix. `@upstash/redis` auto-deserializes replies, so a field + * holding valid JSON (`"hi"`, `123`) would come back as a different type; the prefix makes every + * stored value unparseable as JSON and guarantees a byte-exact round trip. + */ +const MARKER = "agentkit-event-v1:"; + +/** + * Append a batch in one script: refuse a closed log (an event appended after `close` would land + * after readers' final drain and never be delivered), add every entry, refresh the stream's expiry. + * Under `allow-key-locking` both keys are declared. + * KEYS: stream, closed flag. ARGV: ttlSeconds (0 = none), encoded events… Returns the entry ids. + */ +const APPEND = `${KEY_LOCKING}if redis.call("EXISTS", KEYS[2]) == 1 then + return redis.error_reply("EVENTLOG_CLOSED") +end +local ids = {} +for i = 2, #ARGV do ids[#ids + 1] = redis.call("XADD", KEYS[1], "*", "e", ARGV[i]) end +if tonumber(ARGV[1]) > 0 then redis.call("EXPIRE", KEYS[1], ARGV[1]) end +return ids`; + +/** Thrown by `append` on a log that has been closed. */ +export class EventLogClosedError extends Error { + constructor(readonly logId: string) { + super(`EventLog: log ${JSON.stringify(logId)} is closed; nothing more can be appended.`); + this.name = "EventLogClosedError"; + } +} + +export interface EventLogConfig { + /** Upstash Redis client. */ + redis: Redis; + /** Base key prefix. Defaults to `agentkit:log`. */ + prefix?: string; + /** + * Expiry applied to a log on every append and on close, in seconds, so logs whose producer died + * without closing still go away. `0` disables expiry. + * @default 86400 + */ + ttlSeconds?: number; + /** + * How often a tailing {@link EventLog.read} polls for new entries, in ms. Blocking reads span + * requests, which the REST API does not keep, so tailing is a poll. + * @default 250 + */ + pollIntervalMs?: number; + /** + * Report the sdk name + version to Upstash as a header on the requests made by your redis client. + * Can also be disabled with the `UPSTASH_DISABLE_TELEMETRY` env var. Defaults to `true`. + */ + enableTelemetry?: boolean; +} + +/** One stored event and its position. `id` is the Redis Stream entry id, opaque to callers. */ +export interface LogEntry { + id: string; + event: T; +} + +function assertLogId(logId: string): void { + if (typeof logId !== "string" || logId === "" || /[\r\n]/.test(logId)) { + throw new Error("EventLog: `logId` must be a non-empty string without CR/LF."); + } +} + +const PAGE = 500; + +/** + * An append-only, resumable event log on Upstash Redis Streams — the storage under resumable + * streaming (reload mid-answer, a second device joining), fan-out to several readers, and replay. + * + * - `append` writes a batch atomically (one script) and returns one position per event, in order; + * it refuses a closed log with {@link EventLogClosedError}. + * - `read` yields everything **after** a position and keeps tailing until the log is closed. + * - `snapshot` returns what is stored right now and never waits. + * - `close` terminalizes the log; readers drain what is left and stop. + * + * Keys: `:events:` (the stream) and `:closed:` (the terminal flag). + */ +export class EventLog { + private redis: Redis; + private prefix: string; + private ttlSeconds: number; + private pollIntervalMs: number; + + constructor(config: EventLogConfig) { + this.redis = config.redis; + addTelemetry(config.redis, config.enableTelemetry); + this.prefix = config.prefix ?? "agentkit:log"; + this.ttlSeconds = config.ttlSeconds ?? 86_400; + this.pollIntervalMs = config.pollIntervalMs ?? 250; + } + + private streamKey(logId: string): string { + return `${this.prefix}:events:${logId}`; + } + + private closedKey(logId: string): string { + return `${this.prefix}:closed:${logId}`; + } + + private decode(fields: unknown): T { + const raw = (fields as { e?: unknown } | null)?.e; + if (typeof raw !== "string" || !raw.startsWith(MARKER)) { + throw new Error("EventLog: stored entry is not an EventLog event."); + } + return JSON.parse(raw.slice(MARKER.length)) as T; + } + + /** Append a batch atomically; returns one position per event, in the same order. */ + async append(logId: string, events: T[]): Promise { + assertLogId(logId); + if (events.length === 0) return []; + try { + const ids = (await this.redis.eval( + APPEND, + [this.streamKey(logId), this.closedKey(logId)], + [String(this.ttlSeconds), ...events.map((event) => MARKER + JSON.stringify(event))], + )) as unknown[]; + return ids.map(String); + } catch (error) { + if (/EVENTLOG_CLOSED/.test(String((error as Error)?.message ?? error))) { + throw new EventLogClosedError(logId); + } + throw error; + } + } + + /** One page of entries strictly after `after` (`"0"` = from the start). */ + private async page(logId: string, after: string, count = PAGE): Promise[]> { + const start = after === "0" ? "-" : `(${after}`; + const rows = (await this.redis.xrange(this.streamKey(logId), start, "+", count)) as Record< + string, + unknown + >; + // Stream ids are not integer-like strings, so object key order is insertion (= stream) order. + return Object.entries(rows ?? {}).map(([id, fields]) => ({ id, event: this.decode(fields) })); + } + + /** Everything stored right now, in append order. Never waits; `[]` for an unknown log. */ + async snapshot(logId: string): Promise[]> { + assertLogId(logId); + const out: LogEntry[] = []; + let after = "0"; + for (;;) { + const rows = await this.page(logId, after); + out.push(...rows); + if (rows.length < PAGE) return out; + after = rows[rows.length - 1]!.id; + } + } + + /** The position of the last stored entry, or `null` for an empty/unknown log. */ + async tail(logId: string): Promise { + assertLogId(logId); + const rows = (await this.redis.xrevrange(this.streamKey(logId), "+", "-", 1)) as Record< + string, + unknown + >; + return Object.keys(rows ?? {})[0] ?? null; + } + + /** Mark the log complete. Readers drain what is stored and then stop. Idempotent. */ + async close(logId: string): Promise { + assertLogId(logId); + const tx = this.redis.multi(); + if (this.ttlSeconds > 0) { + tx.set(this.closedKey(logId), "1", { ex: this.ttlSeconds }); + tx.expire(this.streamKey(logId), this.ttlSeconds); + } else { + tx.set(this.closedKey(logId), "1"); + } + await tx.exec(); + } + + /** Whether {@link EventLog.close} has been called for this log. */ + async isClosed(logId: string): Promise { + assertLogId(logId); + return (await this.redis.exists(this.closedKey(logId))) === 1; + } + + /** Delete the log and its closed flag. */ + async delete(logId: string): Promise { + assertLogId(logId); + await this.redis.del(this.streamKey(logId), this.closedKey(logId)); + } + + /** + * Yield every entry strictly after `after` (default: from the start), then keep tailing until the + * log is closed or `signal` aborts. Closing is checked only once the reader has caught up, and a + * final drain runs after it is seen, so nothing appended before `close` is skipped. + * + * `firstEntryTimeoutMs` bounds the wait for a log that has no entries at all yet (a joiner that + * attached to a run that never produced) — it rejects instead of parking forever. + */ + async *read( + logId: string, + opts: { + after?: string; + signal?: AbortSignal; + pollIntervalMs?: number; + firstEntryTimeoutMs?: number; + } = {}, + ): AsyncGenerator> { + assertLogId(logId); + const poll = opts.pollIntervalMs ?? this.pollIntervalMs; + let after = opts.after ?? "0"; + let seenAny = after !== "0"; + const startedAt = Date.now(); + for (;;) { + if (opts.signal?.aborted) return; + const rows = await this.page(logId, after); + for (const row of rows) { + seenAny = true; + after = row.id; + yield row; + if (opts.signal?.aborted) return; + } + if (rows.length === PAGE) continue; + if (await this.isClosed(logId)) { + // Anything appended between the page above and the close is picked up here. + for (;;) { + const rest = await this.page(logId, after); + for (const row of rest) { + after = row.id; + yield row; + } + if (rest.length < PAGE) return; + } + } + if ( + !seenAny && + opts.firstEntryTimeoutMs !== undefined && + Date.now() - startedAt >= opts.firstEntryTimeoutMs + ) { + throw new Error( + `EventLog: log ${JSON.stringify(logId)} produced no entries within ${opts.firstEntryTimeoutMs}ms.`, + ); + } + await new Promise((resolve) => { + const timer = setTimeout(done, poll); + function done() { + clearTimeout(timer); + opts.signal?.removeEventListener("abort", done); + resolve(); + } + opts.signal?.addEventListener("abort", done, { once: true }); + }); + } + } +} diff --git a/packages/tanstack-ai/src/stream/stream.test.ts b/packages/tanstack-ai/src/stream/stream.test.ts new file mode 100644 index 0000000..c81e655 --- /dev/null +++ b/packages/tanstack-ai/src/stream/stream.test.ts @@ -0,0 +1,146 @@ +import { randomUUID } from "node:crypto"; +import { afterAll, describe, expect, it } from "vitest"; +import { replayRunStream, toServerSentEventsResponse } from "@tanstack/ai"; +import type { StreamChunk } from "@tanstack/ai"; +import { Redis } from "@upstash/redis"; +import { upstashStream } from "./stream.js"; +import { cleanupKeys, hasRedisCreds, testRedis, uniquePrefix } from "../testing/test-support.js"; + +const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)); + +function text(delta: string): StreamChunk { + return { + type: "TEXT_MESSAGE_CONTENT", + messageId: "m1", + delta, + timestamp: Date.now(), + } as unknown as StreamChunk; +} + +const deltas = (chunks: StreamChunk[]) => + chunks.map((c) => (c as { delta?: string }).delta).filter((d): d is string => d !== undefined); + +/** A "second server instance": its own client, so nothing is shared in-process. */ +const otherInstance = () => Redis.fromEnv(); + +describe.skipIf(!hasRedisCreds)("upstashStream (live Redis)", () => { + const redis = testRedis(); + const prefix = uniquePrefix("tss"); + const cfg = { redis, prefix, pollIntervalMs: 50 }; + + afterAll(async () => { + await cleanupKeys(redis, prefix); + }); + + it("persists a produced SSE response and replays it from another instance", async () => { + const runId = randomUUID(); + async function* produce() { + for (const d of ["Hel", "lo", " world"]) yield text(d); + } + const res = toServerSentEventsResponse(produce(), { + durability: { adapter: upstashStream({ runId }, cfg) }, + }); + const body = await res.text(); + expect(body).toContain("Hel"); + + const joiner = upstashStream({ runId }, { ...cfg, redis: otherInstance() }); + const replayed: StreamChunk[] = []; + for await (const c of replayRunStream(joiner)) replayed.push(c); + expect(deltas(replayed)).toEqual(["Hel", "lo", " world"]); + }); + + it("resumes strictly after a client's last offset (reload mid-answer)", async () => { + const runId = randomUUID(); + const producer = upstashStream({ runId }, cfg); + const offsets = await producer.append([text("a"), text("b"), text("c")]); + await producer.close(); + + // The browser reconnects with Last-Event-ID set to the offset it saw last. + const request = new Request("https://app.test/api/chat", { + headers: { "Last-Event-ID": offsets[0]! }, + }); + const resumed = upstashStream(request, { ...cfg, redis: otherInstance() }); + expect(resumed.resumeFrom()).toBe(offsets[0]); + const seen: string[] = []; + for await (const e of resumed.read(resumed.resumeFrom()!)) { + seen.push((e.chunk as { delta: string }).delta); + } + expect(seen).toEqual(["b", "c"]); + }); + + it("a joiner attached mid-run tails live chunks until the producer closes", async () => { + const runId = randomUUID(); + const producer = upstashStream({ runId }, cfg); + await producer.append([text("1")]); + + const joiner = upstashStream({ runId }, { ...cfg, redis: otherInstance() }); + const seen: StreamChunk[] = []; + const reading = (async () => { + for await (const c of replayRunStream(joiner)) seen.push(c); + })(); + + await sleep(150); + await producer.append([text("2"), text("3")]); + await sleep(150); + await producer.append([text("4")]); + await producer.close(); + await reading; + expect(deltas(seen)).toEqual(["1", "2", "3", "4"]); + }); + + it("reads the run id from X-Run-Id on a producer request", () => { + const request = new Request("https://app.test/api/chat", { headers: { "X-Run-Id": "run-42" } }); + const s = upstashStream(request, cfg); + expect(s.resumeFrom()).toBeNull(); + }); + + it("snapshot returns the stored prefix without waiting on an open log", async () => { + const runId = randomUUID(); + const producer = upstashStream({ runId }, cfg); + expect(await producer.snapshot()).toEqual([]); + const offsets = await producer.append([text("x"), text("y")]); + const snap = await producer.snapshot(); + expect(snap.map((e) => e.offset)).toEqual(offsets); + expect(deltas(snap.map((e) => e.chunk))).toEqual(["x", "y"]); + }); + + it("fails loudly for a concrete offset of an unknown run", async () => { + const runId = randomUUID(); + const producer = upstashStream({ runId }, cfg); + const [offset] = await producer.append([text("z")]); + await redis.del(`${prefix}:events:${runId}`); // expired + const stale = upstashStream({ runId, offset: offset! }, cfg); + const drain = async () => { + for await (const _ of stale.read(offset!)) void _; + }; + await expect(drain()).rejects.toThrow(/Unknown or expired/); + }); + + it("rejects a from-start join on a run that never produces", async () => { + const s = upstashStream({ runId: randomUUID() }, { ...cfg, firstChunkDeadlineMs: 200 }); + const drain = async () => { + for await (const _ of s.read("-1")) void _; + }; + await expect(drain()).rejects.toThrow(/no entries/); + }); + + it("object form stays bound to its runId when given another run's offset", async () => { + const a = upstashStream({ runId: "bind-a-" + randomUUID() }, cfg); + const [offset] = await a.append([text("q")]); + const b = upstashStream({ runId: "bind-b-" + randomUUID(), offset: offset! }, cfg); + const drain = async () => { + for await (const _ of b.read(b.resumeFrom()!)) void _; + }; + await expect(drain()).rejects.toThrow(/belongs to run/); + }); + + it("rejects an offset that belongs to a different run", async () => { + const a = upstashStream({ runId: "run-a-" + randomUUID() }, cfg); + const [offset] = await a.append([text("q")]); + const b = upstashStream({ runId: "run-b-" + randomUUID() }, cfg); + const drain = async () => { + for await (const _ of b.read(offset!)) void _; + }; + await expect(drain()).rejects.toThrow(/belongs to run/); + }); +}); diff --git a/packages/tanstack-ai/src/stream/stream.ts b/packages/tanstack-ai/src/stream/stream.ts new file mode 100644 index 0000000..e0703dd --- /dev/null +++ b/packages/tanstack-ai/src/stream/stream.ts @@ -0,0 +1,162 @@ +import { randomUUID } from "node:crypto"; +import type { Redis } from "@upstash/redis"; +import { Redis as RedisClient } from "@upstash/redis"; +import { EventLog } from "./event-log.js"; +import type { EventLogConfig } from "./event-log.js"; +import { resolveResumeRunId } from "@tanstack/ai"; +import type { StreamChunk, StreamDurability } from "@tanstack/ai"; +import { addTelemetry } from "../telemetry.js"; +import { assertId } from "../persistence/records.js"; + +const OFFSET_PREFIX = "upstash:v1:"; + +/** + * The core {@link EventLogConfig} (TTL, poll interval), with `redis` optional, plus the join deadline. + * Defaults here: `prefix` `agentkit:stream`, `ttlSeconds` 86400 (how long a run stays resumable), + * `pollIntervalMs` 150. + */ +export type UpstashStreamConfig = Omit & { + /** Upstash Redis client. Defaults to `Redis.fromEnv()`. */ + redis?: Redis; + /** + * How long a from-start join (`-1` / `now`) waits for a run that has not produced anything yet + * before failing. Raise it when the producer is queued (e.g. started by a background job). + * @default 2000 + */ + firstChunkDeadlineMs?: number; +}; + +/** Explicit construction when there is no `Request` (server functions, background workers). */ +export interface UpstashStreamInit { + /** The run this adapter produces or attaches to. */ + runId: string; + /** Resume offset captured by the consumer. Defaults to `null` (a producer / from-start reader). */ + offset?: string | null; +} + +function assertRunId(runId: string): string { + assertId(runId, "runId"); + return runId; +} + +function encodeOffset(runId: string, entryId: string): string { + return `${OFFSET_PREFIX}${encodeURIComponent(runId)}:${entryId}`; +} + +function decodeOffset(offset: string): { runId: string; entryId: string } { + if (!offset.startsWith(OFFSET_PREFIX)) + throw new Error(`Invalid upstash stream offset: ${offset}`); + const rest = offset.slice(OFFSET_PREFIX.length); + const sep = rest.lastIndexOf(":"); + const entryId = rest.slice(sep + 1); + if (sep === -1 || !/^\d+-\d+$/.test(entryId)) { + throw new Error(`Invalid upstash stream offset: ${offset}`); + } + return { runId: decodeURIComponent(rest.slice(0, sep)), entryId }; +} + +function readResumeOffset(request: Request): string | null { + const header = request.headers.get("Last-Event-ID"); + if (header) return header; + try { + return new URL(request.url).searchParams.get("offset"); + } catch { + return null; + } +} + +/** + * Resumable delivery for TanStack AI on Upstash Redis Streams — the multi-instance counterpart of + * `memoryStream()`. Every chunk the producer emits is appended to a per-run stream before it is + * delivered, so a client that reloads, drops its connection, or opens the same thread on another + * device replays from its last offset and keeps tailing the live run, even when that request lands + * on a different server instance than the producer. + * + * Offsets are opaque to callers (`upstash:v1::`). Construct from the incoming + * `Request` (reads `Last-Event-ID` / `?offset` and `X-Run-Id` / `?runId`, same precedence as core), + * or from `{ runId, offset }`. + * + * ```ts + * import { chat, toServerSentEventsResponse } from "@tanstack/ai"; + * import { upstashStream } from "@upstash/agentkit-tanstack-ai"; + * + * export async function POST(request: Request) { + * const stream = chat({ adapter, messages, threadId }); + * return toServerSentEventsResponse(stream, { durability: { adapter: upstashStream(request) } }); + * } + * ``` + */ +export function upstashStream( + source: Request | UpstashStreamInit, + config: UpstashStreamConfig = {}, +): StreamDurability { + const redis = config.redis ?? RedisClient.fromEnv(); + addTelemetry(redis, config.enableTelemetry); + const { firstChunkDeadlineMs = 2_000, ...logConfig } = config; + const log = new EventLog({ + ...logConfig, + redis, + prefix: config.prefix ?? "agentkit:stream", + pollIntervalMs: config.pollIntervalMs ?? 150, + }); + + const resumeOffset = + source instanceof Request ? readResumeOffset(source) : (source.offset ?? null); + // Object form is bound to its explicit `runId` (as `memoryStream` is): a mismatched offset is + // rejected by `read`, never silently re-targeted. Only a Request derives the run from its offset. + let runId: string; + if (!(source instanceof Request)) { + runId = assertRunId(source.runId); + } else if (resumeOffset !== null && resumeOffset !== "-1" && resumeOffset !== "now") { + runId = assertRunId(decodeOffset(resumeOffset).runId); + } else { + const requested = resolveResumeRunId(source); + runId = requested === null ? randomUUID() : assertRunId(requested); + } + + return { + resumeFrom: () => resumeOffset, + + append: async (chunks) => { + const ids = await log.append(runId, chunks); + return ids.map((id) => encodeOffset(runId, id)); + }, + + snapshot: async () => + (await log.snapshot(runId)).map((entry) => ({ + offset: encodeOffset(runId, entry.id), + chunk: entry.event, + })), + + close: () => log.close(runId), + + read: async function* (offset, signal) { + const isFromStartJoin = offset === "-1" || offset === "now"; + let after = "0"; + if (offset === "now") { + after = (await log.tail(runId)) ?? "0"; + } else if (!isFromStartJoin) { + const decoded = decodeOffset(offset); + if (decoded.runId !== runId) { + throw new Error( + `Upstash stream offset belongs to run ${JSON.stringify(decoded.runId)}, not ${JSON.stringify(runId)}`, + ); + } + // A concrete offset for a run with nothing stored means it expired (or never existed): fail + // loudly rather than park a reader on a log that will never be written again. + if ((await log.tail(runId)) === null && !(await log.isClosed(runId))) { + throw new Error(`Unknown or expired upstash stream run: ${JSON.stringify(runId)}`); + } + after = decoded.entryId; + } + const hasData = (await log.tail(runId)) !== null; + for await (const entry of log.read(runId, { + after, + ...(signal ? { signal } : {}), + ...(isFromStartJoin && !hasData ? { firstEntryTimeoutMs: firstChunkDeadlineMs } : {}), + })) { + yield { offset: encodeOffset(runId, entry.id), chunk: entry.event }; + } + }, + }; +} diff --git a/packages/tanstack-ai/src/telemetry.test.ts b/packages/tanstack-ai/src/telemetry.test.ts new file mode 100644 index 0000000..1d03024 --- /dev/null +++ b/packages/tanstack-ai/src/telemetry.test.ts @@ -0,0 +1,173 @@ +import { afterEach, describe, expect, test } from "vitest"; +import { Redis } from "@upstash/redis"; +import { SDK_TELEMETRY } from "@upstash/agentkit-sdk"; +import { TANSTACK_AI_TELEMETRY } from "./telemetry.js"; +import { VERSION } from "./version.js"; +import { upstashPersistence } from "./persistence/persistence.js"; +import { upstashBlobStore } from "./persistence/blob-store.js"; +import { upstashMemory } from "./memory/memory.js"; +import { upstashStream } from "./stream/stream.js"; +import { EventLog } from "./stream/event-log.js"; +import { upstashLocks } from "./locks/locks.js"; +import { RedisLock } from "./locks/redis-lock.js"; +import { rateLimit, toolCache } from "./middleware/middleware.js"; +import { createSearchTools } from "./search/search-tools.js"; +import { Ratelimit } from "@upstash/agentkit-sdk"; +import { s } from "@upstash/redis"; + +/** + * Proof the tags ride on the wire, not just that the client was told about them: the Upstash client + * calls the global `fetch`, so stubbing it captures the real outgoing request headers. + */ +const realFetch = globalThis.fetch; +let sent: Record[] = []; + +function spyOnFetch(): void { + sent = []; + globalThis.fetch = (async (_url: unknown, init?: { headers?: Record }) => { + sent.push({ ...init?.headers }); + return new Response(JSON.stringify({ result: "OK" }), { status: 200 }); + }) as unknown as typeof fetch; +} + +/** A client pointed at nowhere, one request per command (no auto-pipelining) for the stub above. */ +function stubbedRedis(): Redis { + return new Redis({ + url: "https://telemetry.test.upstash.io", + token: "test-token", + responseEncoding: false, + retry: false, + enableAutoPipelining: false, + }); +} + +/** The tags on the one request `send` makes. */ +async function tagsAfter(redis: Redis): Promise { + await redis.set("agentkit:telemetry-test", "value"); + expect(sent.length).toBe(1); + return (sent[0]?.["Upstash-Telemetry-Sdk"] ?? "").split(",").filter(Boolean); +} + +const bucket = { + put: async () => ({ etag: "e", size: 0 }), + get: async () => ({ body: new ReadableStream() }), + del: async () => undefined, +}; + +// Every public factory, each on its own client, with the core tag it should also carry (if any). +const factories: Array<[string, (redis: Redis, enableTelemetry?: boolean) => unknown, boolean]> = [ + [ + "upstashPersistence", + (redis, e) => upstashPersistence({ redis, ...(e === false ? { enableTelemetry: e } : {}) }), + false, + ], + [ + "upstashBlobStore", + (redis, e) => + upstashBlobStore({ + redis, + bucket, + prefix: "p", + ...(e === false ? { enableTelemetry: e } : {}), + }), + false, + ], + [ + "upstashStream", + (redis, e) => + upstashStream({ runId: "r" }, { redis, ...(e === false ? { enableTelemetry: e } : {}) }), + false, + ], + [ + "EventLog", + (redis, e) => new EventLog({ redis, ...(e === false ? { enableTelemetry: e } : {}) }), + false, + ], + [ + "upstashLocks", + (redis, e) => upstashLocks({ redis, ...(e === false ? { enableTelemetry: e } : {}) }), + false, + ], + [ + "RedisLock", + (redis, e) => new RedisLock({ redis, ...(e === false ? { enableTelemetry: e } : {}) }), + false, + ], + [ + "upstashMemory", + (redis, e) => upstashMemory({ redis, ...(e === false ? { enableTelemetry: e } : {}) }), + true, + ], + [ + "toolCache", + (redis, e) => + toolCache({ + redis, + tools: ["t"], + userId: "u", + ...(e === false ? { enableTelemetry: e } : {}), + }), + true, + ], + [ + "rateLimit", + (redis, e) => + rateLimit({ + redis, + limiter: Ratelimit.fixedWindow(1, "1 s"), + identifier: "u", + ...(e === false ? { enableTelemetry: e } : {}), + }), + true, + ], + [ + "createSearchTools", + (redis, e) => + createSearchTools({ + redis, + schema: s.object({ a: s.string() }), + ...(e === false ? { enableTelemetry: e } : {}), + }), + true, + ], +]; + +describe("telemetry on the wire", () => { + afterEach(() => { + globalThis.fetch = realFetch; + }); + + test("the tag is this package's name and version", () => { + expect(TANSTACK_AI_TELEMETRY).toBe(`@upstash/agentkit-tanstack-ai@${VERSION}`); + }); + + test.each(factories)("%s tags its client", async (_name, create, withCore) => { + spyOnFetch(); + const redis = stubbedRedis(); + create(redis); + const tags = await tagsAfter(redis); + expect(tags[0]).toMatch(/^@upstash\/redis@/); + expect(tags).toContain(TANSTACK_AI_TELEMETRY); + if (withCore) expect(tags).toContain(SDK_TELEMETRY); + }); + + test("a client shared by every factory carries each tag once", async () => { + spyOnFetch(); + const redis = stubbedRedis(); + for (const [, create] of factories) create(redis); + const tags = await tagsAfter(redis); + expect(tags.filter((t) => t === TANSTACK_AI_TELEMETRY)).toHaveLength(1); + expect(tags.filter((t) => t === SDK_TELEMETRY)).toHaveLength(1); + }); + + test.each(factories)( + "%s with enableTelemetry: false adds no agentkit tag", + async (_name, create) => { + spyOnFetch(); + const redis = stubbedRedis(); + create(redis, false); + const tags = await tagsAfter(redis); + expect(tags.some((t) => t.includes("agentkit"))).toBe(false); + }, + ); +}); diff --git a/packages/tanstack-ai/src/telemetry.ts b/packages/tanstack-ai/src/telemetry.ts new file mode 100644 index 0000000..a266156 --- /dev/null +++ b/packages/tanstack-ai/src/telemetry.ts @@ -0,0 +1,14 @@ +import { addTelemetry as tagClient } from "@upstash/agentkit-sdk"; +import { VERSION } from "./version.js"; + +/** The telemetry tag of this package, appended to the redis client's `Upstash-Telemetry-Sdk` header. */ +export const TANSTACK_AI_TELEMETRY = `@upstash/agentkit-tanstack-ai@${VERSION}`; + +/** + * Tag the redis client with this adapter's sdk name + version. The core primitives built underneath + * add their own `@upstash/agentkit-sdk` tag, so the header reports both layers. Opt out with + * `enableTelemetry: false`, the same option on the redis client, or `UPSTASH_DISABLE_TELEMETRY`. + */ +export function addTelemetry(redis: unknown, enableTelemetry?: boolean): void { + tagClient(redis, { sdk: TANSTACK_AI_TELEMETRY, enabled: enableTelemetry }); +} diff --git a/packages/tanstack-ai/src/testing/test-adapter.ts b/packages/tanstack-ai/src/testing/test-adapter.ts new file mode 100644 index 0000000..6ec9362 --- /dev/null +++ b/packages/tanstack-ai/src/testing/test-adapter.ts @@ -0,0 +1,87 @@ +/** + * Test-only: a scripted TanStack AI text adapter, so middleware and memory are exercised through a + * real `chat()` agent loop without a model provider. `turns[i]` is what the model "says" on call i: + * either a list of tool calls or a final text answer. The last entry repeats if the loop runs longer. + */ +import type { StreamChunk } from "@tanstack/ai"; + +export type ScriptedTurn = + | { toolCalls: Array<{ id: string; name: string; args: unknown }> } + | { text: string }; + +export interface ScriptedAdapter { + kind: "text"; + name: string; + model: string; + "~types": never; + /** Every `chatStream` call's options, for asserting on what the loop sent the model. */ + calls: Array<{ messages: unknown[]; systemPrompts?: unknown; tools?: unknown[] }>; + chatStream: (options: unknown) => AsyncIterable; + structuredOutput: () => Promise; +} + +export function scriptedAdapter(turns: ScriptedTurn[]): ScriptedAdapter { + const calls: ScriptedAdapter["calls"] = []; + return { + kind: "text", + name: "scripted", + model: "scripted-1", + "~types": undefined as never, + calls, + async *chatStream(options: unknown) { + calls.push(options as ScriptedAdapter["calls"][number]); + const turn = turns[Math.min(calls.length - 1, turns.length - 1)]!; + const timestamp = Date.now(); + const n = calls.length; + if ("toolCalls" in turn) { + for (const call of turn.toolCalls) { + yield { + type: "TOOL_CALL_START", + toolCallId: call.id, + toolCallName: call.name, + toolName: call.name, + timestamp, + } as unknown as StreamChunk; + yield { + type: "TOOL_CALL_ARGS", + toolCallId: call.id, + delta: JSON.stringify(call.args), + timestamp, + } as unknown as StreamChunk; + yield { + type: "TOOL_CALL_END", + toolCallId: call.id, + toolCallName: call.name, + toolName: call.name, + input: call.args, + timestamp, + } as unknown as StreamChunk; + } + yield { + type: "RUN_FINISHED", + finishReason: "tool_calls", + timestamp, + } as unknown as StreamChunk; + } else { + const messageId = `msg-${n}`; + yield { + type: "TEXT_MESSAGE_START", + messageId, + role: "assistant", + timestamp, + } as unknown as StreamChunk; + yield { + type: "TEXT_MESSAGE_CONTENT", + messageId, + delta: turn.text, + timestamp, + } as unknown as StreamChunk; + yield { type: "TEXT_MESSAGE_END", messageId, timestamp } as unknown as StreamChunk; + yield { type: "RUN_FINISHED", finishReason: "stop", timestamp } as unknown as StreamChunk; + } + }, + structuredOutput: async () => { + throw new Error("scriptedAdapter: structuredOutput is not scripted"); + }, + }; +} diff --git a/packages/tanstack-ai/src/testing/test-bucket.ts b/packages/tanstack-ai/src/testing/test-bucket.ts new file mode 100644 index 0000000..da67f38 --- /dev/null +++ b/packages/tanstack-ai/src/testing/test-bucket.ts @@ -0,0 +1,86 @@ +/** + * Test-only stand-in for an `@upstash/blob` `Bucket`: objects in memory, and `signedReadUrl` served by + * a local HTTP server that honours `Range` (answering 206), so the store's range-read path is + * exercised over real HTTP rather than short-circuited. + */ +import { createServer, type Server } from "node:http"; +import { randomUUID } from "node:crypto"; +import type { BlobBucketLike } from "../persistence/blob-store.js"; + +export interface TestBucket extends BlobBucketLike { + close(): Promise; + /** Paths currently stored, for asserting on what the store wrote. */ + paths(): string[]; + /** How many Range requests the HTTP server answered with 206. */ + rangeHits: number; +} + +export async function testBucket(): Promise { + const objects = new Map(); + const tokens = new Map(); + let rangeHits = 0; + const server: Server = createServer((req, res) => { + const path = tokens.get((req.url ?? "").slice(1)); + const object = path ? objects.get(path) : undefined; + if (!object) { + res.writeHead(404).end(); + return; + } + const match = /^bytes=(\d+)-(\d+)$/.exec(String(req.headers.range ?? "")); + if (match) { + const start = Number(match[1]); + const end = Math.min(Number(match[2]), object.bytes.byteLength - 1); + rangeHits++; + res.writeHead(206, { "content-range": `bytes ${start}-${end}/${object.bytes.byteLength}` }); + res.end(Buffer.from(object.bytes.subarray(start, end + 1))); + return; + } + res.writeHead(200).end(Buffer.from(object.bytes)); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const { port } = server.address() as { port: number }; + + const toBytes = async (body: unknown): Promise => { + if (body instanceof Uint8Array) return new Uint8Array(body); + if (typeof body === "string") return new TextEncoder().encode(body); + if (body instanceof ArrayBuffer) return new Uint8Array(body); + return new Uint8Array( + await new Response(body as ConstructorParameters[0]).arrayBuffer(), + ); + }; + + const bucket: TestBucket = { + rangeHits: 0, + async put(path, body) { + const bytes = await toBytes(body); + const etag = randomUUID(); + objects.set(path, { bytes, etag }); + return { etag, size: bytes.byteLength }; + }, + async get(path) { + const object = objects.get(path); + if (!object) throw Object.assign(new Error("not found"), { code: "not_found" }); + const bytes = object.bytes.slice(); + return { + body: new ReadableStream({ + start(c) { + c.enqueue(bytes); + c.close(); + }, + }), + }; + }, + async del(path) { + objects.delete(path); + }, + async signedReadUrl(path) { + const token = randomUUID(); + tokens.set(token, path); + return { url: `http://127.0.0.1:${port}/${token}` }; + }, + paths: () => [...objects.keys()], + close: () => new Promise((resolve) => server.close(() => resolve())), + }; + Object.defineProperty(bucket, "rangeHits", { get: () => rangeHits }); + return bucket; +} diff --git a/packages/tanstack-ai/src/testing/test-support.ts b/packages/tanstack-ai/src/testing/test-support.ts new file mode 100644 index 0000000..1f2c65f --- /dev/null +++ b/packages/tanstack-ai/src/testing/test-support.ts @@ -0,0 +1,45 @@ +/** + * Test-only helpers (never imported by `index.ts`). The adapter tests run against a real Upstash + * Redis instance — only LLM calls are mocked. Credentials come from the repo-root `.env`; when + * absent, `hasRedisCreds` is false and the suites skip themselves. + */ +import { randomUUID } from "node:crypto"; +import { config } from "dotenv"; +import { Redis } from "@upstash/redis"; + +config(); + +export const hasRedisCreds = Boolean( + process.env.UPSTASH_REDIS_REST_URL && process.env.UPSTASH_REDIS_REST_TOKEN, +); + +/** Whether a real OpenAI key is available for end-to-end model tests. */ +export const hasOpenAIKey = Boolean(process.env.OPENAI_API_KEY); + +/** The model used by tests (always gpt-4o). */ +export const TEST_MODEL = "gpt-4o"; + +/** A real Upstash Redis client from env. Only call when `hasRedisCreds` is true. */ +export function testRedis(): Redis { + return Redis.fromEnv(); +} + +/** A collision-proof key prefix so parallel test runs never share keys or indexes. */ +export function uniquePrefix(label: string): string { + return `test:${label}:${randomUUID().slice(0, 8)}`; +} + +/** A collision-proof, colon-free userId (core key-part validation rejects ':', the key separator). */ +export function uniqueUserId(label: string): string { + return `test-${label}-${randomUUID().slice(0, 8)}`; +} + +/** Delete every key under a key prefix (best-effort cleanup in afterAll hooks). */ +export async function cleanupKeys(redis: Redis, prefix: string): Promise { + let cursor = "0"; + do { + const [next, keys] = await redis.scan(cursor, { match: `${prefix}*`, count: 200 }); + cursor = next; + if (keys.length) await redis.del(...keys); + } while (cursor !== "0"); +} diff --git a/packages/tanstack-ai/src/version.ts b/packages/tanstack-ai/src/version.ts new file mode 100644 index 0000000..5df3807 --- /dev/null +++ b/packages/tanstack-ai/src/version.ts @@ -0,0 +1,2 @@ +// Generated by scripts/sync-version.mjs (run by `pnpm ci:version`) — do not edit by hand. +export const VERSION = "0.0.0"; diff --git a/packages/tanstack-ai/tsconfig.json b/packages/tanstack-ai/tsconfig.json new file mode 100644 index 0000000..5285d28 --- /dev/null +++ b/packages/tanstack-ai/tsconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "rootDir": "src", + "outDir": "dist" + }, + "include": ["src"] +} diff --git a/packages/tanstack-ai/tsup.config.ts b/packages/tanstack-ai/tsup.config.ts new file mode 100644 index 0000000..6644d1e --- /dev/null +++ b/packages/tanstack-ai/tsup.config.ts @@ -0,0 +1,14 @@ +import { defineConfig } from "tsup"; + +export default defineConfig({ + entry: { + index: "src/index.ts", + persistence: "src/persistence/index.ts", + memory: "src/memory/index.ts", + }, + format: ["esm"], + dts: true, + clean: true, + sourcemap: true, + treeshake: true, +}); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 2546539..290c633 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -234,6 +234,64 @@ importers: specifier: 7.0.2 version: 7.0.2 + examples/tanstack-ai-demo: + dependencies: + '@tanstack/ai': + specifier: 0.63.0 + version: 0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3) + '@tanstack/ai-client': + specifier: 0.36.0 + version: 0.36.0(@opentelemetry/api@1.9.1)(zod@4.4.3) + '@tanstack/ai-memory': + specifier: 0.2.8 + version: 0.2.8(@tanstack/ai@0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(vitest@2.1.9(@types/node@20.19.43)(lightningcss@1.32.0)) + '@tanstack/ai-openai': + specifier: 0.25.1 + version: 0.25.1(@tanstack/ai@0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(ws@8.21.3)(zod@4.4.3) + '@tanstack/ai-persistence': + specifier: 0.7.1 + version: 0.7.1(@tanstack/ai@0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(vitest@2.1.9(@types/node@20.19.43)(lightningcss@1.32.0)) + '@tanstack/ai-react': + specifier: 0.29.3 + version: 0.29.3(@opentelemetry/api@1.9.1)(@tanstack/ai@0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(@types/react@19.2.15)(react-dom@19.2.6(react@19.2.6))(react@19.2.6)(zod@4.4.3) + '@upstash/agentkit-tanstack-ai': + specifier: workspace:* + version: link:../../packages/tanstack-ai + '@upstash/redis': + specifier: ^1.38.4 + version: 1.38.4 + next: + specifier: 16.2.9 + version: 16.2.9(@opentelemetry/api@1.9.1)(react-dom@19.2.6(react@19.2.6))(react@19.2.6) + react: + specifier: 19.2.6 + version: 19.2.6 + react-dom: + specifier: 19.2.6 + version: 19.2.6(react@19.2.6) + zod: + specifier: ^4 + version: 4.4.3 + devDependencies: + '@types/node': + specifier: ^20 + version: 20.19.43 + '@types/react': + specifier: 19.2.15 + version: 19.2.15 + '@types/react-dom': + specifier: 19.2.3 + version: 19.2.3(@types/react@19.2.15) + dotenv: + specifier: ^16.4.5 + version: 16.6.1 + typescript: + specifier: ^5 + version: 5.9.3 + vitest: + specifier: ^2.0.0 + version: 2.1.9(@types/node@20.19.43)(lightningcss@1.32.0) + packages/ai-sdk: dependencies: '@upstash/agentkit-sdk': @@ -325,8 +383,44 @@ importers: specifier: ^16.4.5 version: 16.6.1 + packages/tanstack-ai: + dependencies: + '@upstash/agentkit-sdk': + specifier: workspace:* + version: link:../sdk + '@upstash/redis': + specifier: ^1.38.4 + version: 1.38.4 + zod: + specifier: ^3.23.8 || ^4 + version: 4.4.3 + devDependencies: + '@tanstack/ai': + specifier: 0.63.0 + version: 0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3) + '@tanstack/ai-memory': + specifier: 0.2.8 + version: 0.2.8(@tanstack/ai@0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(vitest@2.1.9(@types/node@20.19.43)(lightningcss@1.32.0)) + '@tanstack/ai-persistence': + specifier: 0.7.1 + version: 0.7.1(@tanstack/ai@0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(vitest@2.1.9(@types/node@20.19.43)(lightningcss@1.32.0)) + '@upstash/blob': + specifier: 0.0.7 + version: 0.0.7(react@19.2.6) + dotenv: + specifier: ^16.4.5 + version: 16.6.1 + packages: + '@ag-ui/core@1.0.0': + resolution: {integrity: sha512-yCRhsQvb4+lmGMrDsi4TOJxqKpm/kb3H64wwL+Z90jNVrkhy17RErfzsNgybfbjNKLM0nenEDYDVJsZHJJSigQ==} + peerDependencies: + zod: ^3.25.18 || ^4.0.0 + peerDependenciesMeta: + zod: + optional: true + '@ai-sdk/gateway@4.0.87': resolution: {integrity: sha512-6wJmX/D5OO7tvDJbtFdXdBuRt8AkCnW+FLK+W35UHe/tHgolAa4hclXtJcMSYnlTDy/t5DNRgdesVaoMXtwzow==} engines: {node: '>=22'} @@ -2363,6 +2457,94 @@ packages: '@tailwindcss/postcss@4.3.0': resolution: {integrity: sha512-Jm05Tjx+9yCLGv5qw1c+84Psds8MnyrEQYCB+FFk2lgGiUjlRqdxke4mVTuYrj2xnVZqKim2Apr5ySuQRYAw/w==} + '@tanstack/ai-client@0.36.0': + resolution: {integrity: sha512-fXqf2r+PeV79jZwY6ZJ3QwL5yKC56qB60nOSDbHYJfmCcW/sOkFN8/Z4SIrQdwsDpDYlAJNuXzSs4xx2M3s6VA==} + + '@tanstack/ai-event-client@0.13.0': + resolution: {integrity: sha512-qGN7saScQHqDd06DE5CNG+wl1O9v6anwsGCCeP4N7zFbZd5k9/2C1u4cdJ2+Qxl2aA8drusLFZMhY/ljhBJqLQ==} + + '@tanstack/ai-memory@0.2.8': + resolution: {integrity: sha512-eG4jZwwqiKoKPq5C+aCWywC1f0LBD7g4q4vw4o8rNB/sHxGilw2HcUmAWg6bXdr6HtApdUskaNsZGtS9oZJrHA==} + peerDependencies: + '@honcho-ai/sdk': '>=2.0.0' + '@tanstack/ai': ^0.63.0 + '@vectorize-io/hindsight-client': '>=0.6.0' + ioredis: '>=5.0.0' + redis: '>=4.0.0' + vitest: ^4.1.10 + peerDependenciesMeta: + '@honcho-ai/sdk': + optional: true + '@vectorize-io/hindsight-client': + optional: true + ioredis: + optional: true + redis: + optional: true + vitest: + optional: true + + '@tanstack/ai-openai@0.25.1': + resolution: {integrity: sha512-SmYpGbHS7Uor49STSqhT5RIK0x3nnOx51koc16xvclB+5+lN/33Psty2/hfhaeRBkHk2t7q1F4fSjUKr/E6Kig==} + peerDependencies: + '@tanstack/ai': ^0.63.0 + + '@tanstack/ai-persistence@0.7.1': + resolution: {integrity: sha512-d/m0tjZkKigPwFexqMs5QCi7STsDGoHGBWP4GvAY5J8tAx3txRD+VWYMoJoiiIvU9H95uGq7hyBPAkfHjlh3Cw==} + peerDependencies: + '@tanstack/ai': ^0.63.0 + vitest: ^4.1.10 + peerDependenciesMeta: + vitest: + optional: true + + '@tanstack/ai-react@0.29.3': + resolution: {integrity: sha512-JLU4Pp/Qi3r8pBquugDKyEwJ36W2irwKCH391Ln4k3tdtk4a0rI8jsz3xepdWXkgWr/DhJAlKDzZ6sYZPJWVNg==} + peerDependencies: + '@mcp-ui/client': ^7 + '@tanstack/ai': ^0.63.0 + '@types/react': '>=18.0.0' + react: '>=18.0.0' + react-dom: '>=18.0.0' + peerDependenciesMeta: + '@mcp-ui/client': + optional: true + react-dom: + optional: true + + '@tanstack/ai-utils@0.4.1': + resolution: {integrity: sha512-B3PGn2WYiRtivZCt7MUN2iXeoK9fwZhgdbrPIz5hPqVR8r8sNDZMmMqtiOkp7sxzPvVlCBItFVHw7QpZ9g6e9w==} + + '@tanstack/ai@0.63.0': + resolution: {integrity: sha512-4S4hOOc/2LvxNkMzwuBaECtchPQsxrlLNqnmi/WjcXmX8gyboy8UNPwwS/9XzFFJ4VwSEW2Md+h+OtTGmf6Y/g==} + engines: {node: '>=18'} + peerDependencies: + '@opentelemetry/api': '>=1.9.0' + peerDependenciesMeta: + '@opentelemetry/api': + optional: true + + '@tanstack/devtools-event-client@0.4.4': + resolution: {integrity: sha512-6T5Yop/793YI+H+5J8Hsyj4kCih9sl4t3ElLgKioW5hk3ocn+ZdSJ94tT7vL7uabxSugWYBZlOTMPzEw2puvQw==} + engines: {node: '>=18'} + hasBin: true + + '@tanstack/markdown@0.0.13': + resolution: {integrity: sha512-ZSQP2ulBvlPzdpZeRheBJnaELo7e4ET/D/XrfgweHyqhhLMN5vUJJAUl7ufSRd7GlfO/vwdzjKPh0E0JwqJkDA==} + peerDependencies: + octane: '>=0.1.12' + react: '>=18' + peerDependenciesMeta: + octane: + optional: true + react: + optional: true + + '@tanstack/openai-base@0.12.1': + resolution: {integrity: sha512-mkt86u9kIW4oScA738ntrpTMXTqkTxZl5XU54lcUkdl5xT3Kx4fM4HT9x9x49twiUNHJEP7y44m7gK7FkZbiMA==} + peerDependencies: + '@tanstack/ai': ^0.63.0 + '@types/d3-array@3.2.2': resolution: {integrity: sha512-hOLWVbm7uRza0BYXpIIW5pxfrKe0W+D5lrFiAEYR+pb6w3N2SwSMaJbXdUfSEv+dT4MfHBLtn5js0LAWaO6otw==} @@ -2737,6 +2919,15 @@ packages: '@upsetjs/venn.js@2.0.0': resolution: {integrity: sha512-WbBhLrooyePuQ1VZxrJjtLvTc4NVfpOyKx0sKqioq9bX1C1m7Jgykkn8gLrtwumBioXIqam8DLxp88Adbue6Hw==} + '@upstash/blob@0.0.7': + resolution: {integrity: sha512-Q2+GqZdnZxjzdHrROd/dSGVmR7im4FJ6SSKh5qfOLyYiZmtzMiLjd+yG0vwq2urtZ+UswpF4MlLQ//UEIUNkCw==} + engines: {node: '>=20'} + peerDependencies: + react: '>=18' + peerDependenciesMeta: + react: + optional: true + '@upstash/box@0.7.5': resolution: {integrity: sha512-PfUn8Ppz+gkI76vftM1VnFvLxJs4SgFy3CR3898ahl4YXReXdWrXwTkP3QjzHjVd9Lj3e1dkQ8iLod9fTCmALg==} engines: {node: '>=18.0.0'} @@ -4172,6 +4363,26 @@ packages: oniguruma-to-es@4.3.6: resolution: {integrity: sha512-csuQ9x3Yr0cEIs/Zgx/OEt9iBw9vqIunAPQkx19R/fiMq2oGVTgcMqO/V3Ybqefr1TBvosI6jU539ksaBULJyA==} + openai@6.49.0: + resolution: {integrity: sha512-aYCc0C6L864eR6WSYIwQGyXriw/nIyZx0ObvhzOEVuk0zoBDpynjSbrionWI7q65B5H8jJX0DXR9snEzM6bfPg==} + peerDependencies: + '@aws-sdk/credential-provider-node': '>=3.972.0 <4' + '@smithy/hash-node': '>=4.3.0 <5' + '@smithy/signature-v4': '>=5.4.0 <6' + ws: ^8.18.0 + zod: ^3.25 || ^4.0 + peerDependenciesMeta: + '@aws-sdk/credential-provider-node': + optional: true + '@smithy/hash-node': + optional: true + '@smithy/signature-v4': + optional: true + ws: + optional: true + zod: + optional: true + optionator@0.9.4: resolution: {integrity: sha512-6IpQ7mKUxRcZNLIObR0hz7lxsapSSIYNZJwXPGeF0mTVqGKFIXj1DQcMoT22S3ROcLyY/rz0PWaWZ9ayWmad9g==} engines: {node: '>= 0.8.0'} @@ -4227,6 +4438,9 @@ packages: parse5@7.3.0: resolution: {integrity: sha512-IInvU7fabl34qmi9gY8XOVxhYyMyuH2xUNpb2q8/Y+7552KlejkRvqvD19nMoUW/uQGGbqNpA6Tufu5FL5BZgw==} + partial-json@0.1.7: + resolution: {integrity: sha512-Njv/59hHaokb/hRUjce3Hdv12wd60MtM9Z5Olmn+nehe0QDAsRtRbJPvJ0Z91TusF0SuZRIvnM+S4l6EIP8leA==} + path-data-parser@0.1.0: resolution: {integrity: sha512-NOnmBpt5Y2RWbuv0LMzsayp3lVylAHLPUTut412ZA3l+C4uw4ZVkQbjShYCQ8TCpUMdPapr4YjUqLYD6v68j+w==} @@ -4977,6 +5191,10 @@ packages: snapshots: + '@ag-ui/core@1.0.0(zod@4.4.3)': + optionalDependencies: + zod: 4.4.3 + '@ai-sdk/gateway@4.0.87(zod@4.4.3)': dependencies: '@ai-sdk/provider': 4.0.17 @@ -6792,6 +7010,92 @@ snapshots: postcss: 8.5.15 tailwindcss: 4.3.0 + '@tanstack/ai-client@0.36.0(@opentelemetry/api@1.9.1)(zod@4.4.3)': + dependencies: + '@tanstack/ai': 0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3) + '@tanstack/ai-event-client': 0.13.0 + '@tanstack/ai-utils': 0.4.1 + transitivePeerDependencies: + - '@opentelemetry/api' + - zod + + '@tanstack/ai-event-client@0.13.0': + dependencies: + '@tanstack/devtools-event-client': 0.4.4 + + '@tanstack/ai-memory@0.2.8(@tanstack/ai@0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(vitest@2.1.9(@types/node@20.19.43)(lightningcss@1.32.0))': + dependencies: + '@tanstack/ai': 0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3) + '@tanstack/ai-event-client': 0.13.0 + optionalDependencies: + vitest: 2.1.9(@types/node@20.19.43)(lightningcss@1.32.0) + + '@tanstack/ai-openai@0.25.1(@tanstack/ai@0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(ws@8.21.3)(zod@4.4.3)': + dependencies: + '@tanstack/ai': 0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3) + '@tanstack/ai-utils': 0.4.1 + '@tanstack/openai-base': 0.12.1(@tanstack/ai@0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(ws@8.21.3)(zod@4.4.3) + openai: 6.49.0(ws@8.21.3)(zod@4.4.3) + transitivePeerDependencies: + - '@aws-sdk/credential-provider-node' + - '@smithy/hash-node' + - '@smithy/signature-v4' + - ws + - zod + + '@tanstack/ai-persistence@0.7.1(@tanstack/ai@0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(vitest@2.1.9(@types/node@20.19.43)(lightningcss@1.32.0))': + dependencies: + '@tanstack/ai': 0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3) + '@tanstack/ai-utils': 0.4.1 + optionalDependencies: + vitest: 2.1.9(@types/node@20.19.43)(lightningcss@1.32.0) + + '@tanstack/ai-react@0.29.3(@opentelemetry/api@1.9.1)(@tanstack/ai@0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(@types/react@19.2.15)(react-dom@19.2.6(react@19.2.6))(react@19.2.6)(zod@4.4.3)': + dependencies: + '@tanstack/ai': 0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3) + '@tanstack/ai-client': 0.36.0(@opentelemetry/api@1.9.1)(zod@4.4.3) + '@tanstack/markdown': 0.0.13(react@19.2.6) + '@types/react': 19.2.15 + react: 19.2.6 + optionalDependencies: + react-dom: 19.2.6(react@19.2.6) + transitivePeerDependencies: + - '@opentelemetry/api' + - octane + - zod + + '@tanstack/ai-utils@0.4.1': {} + + '@tanstack/ai@0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3)': + dependencies: + '@ag-ui/core': 1.0.0(zod@4.4.3) + '@standard-schema/spec': 1.1.0 + '@tanstack/ai-event-client': 0.13.0 + '@tanstack/ai-utils': 0.4.1 + partial-json: 0.1.7 + optionalDependencies: + '@opentelemetry/api': 1.9.1 + transitivePeerDependencies: + - zod + + '@tanstack/devtools-event-client@0.4.4': {} + + '@tanstack/markdown@0.0.13(react@19.2.6)': + optionalDependencies: + react: 19.2.6 + + '@tanstack/openai-base@0.12.1(@tanstack/ai@0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(ws@8.21.3)(zod@4.4.3)': + dependencies: + '@tanstack/ai': 0.63.0(@opentelemetry/api@1.9.1)(zod@4.4.3) + '@tanstack/ai-utils': 0.4.1 + openai: 6.49.0(ws@8.21.3)(zod@4.4.3) + transitivePeerDependencies: + - '@aws-sdk/credential-provider-node' + - '@smithy/hash-node' + - '@smithy/signature-v4' + - ws + - zod + '@types/d3-array@3.2.2': {} '@types/d3-axis@3.0.6': @@ -7137,6 +7441,10 @@ snapshots: d3-selection: 3.0.0 d3-transition: 3.0.1(d3-selection@3.0.0) + '@upstash/blob@0.0.7(react@19.2.6)': + optionalDependencies: + react: 19.2.6 + '@upstash/box@0.7.5(zod@4.4.3)': dependencies: ws: 8.21.3 @@ -8877,6 +9185,11 @@ snapshots: regex: 6.1.0 regex-recursion: 6.0.2 + openai@6.49.0(ws@8.21.3)(zod@4.4.3): + optionalDependencies: + ws: 8.21.3 + zod: 4.4.3 + optionator@0.9.4: dependencies: deep-is: 0.1.4 @@ -8938,6 +9251,8 @@ snapshots: dependencies: entities: 6.0.1 + partial-json@0.1.7: {} + path-data-parser@0.1.0: {} path-exists@4.0.0: {} diff --git a/scripts/sync-version.mjs b/scripts/sync-version.mjs index 1136177..eb99153 100644 --- a/scripts/sync-version.mjs +++ b/scripts/sync-version.mjs @@ -18,6 +18,7 @@ const TARGETS = { "packages/ai-sdk": "src/version.ts", "packages/eve": "src/version.ts", "packages/eve-extension": "extension/lib/version.ts", + "packages/tanstack-ai": "src/version.ts", }; const root = resolve(dirname(fileURLToPath(import.meta.url)), "..");