From cf2aa6f78034a16ff99762b40a717c14f39667f1 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 11 Nov 2026 10:42:54 +0000 Subject: [PATCH 1/7] feat(tanstack-ai): sandbox checkpoint/instance stores and MCP task/event stores Add two entry points that implement more TanStack AI 1.x backend contracts on Upstash Redis: - ./sandbox: upstashSandboxCheckpointStore (SandboxCheckpointStore with forkFromCheckpoint) and upstashSandboxInstanceStore. They pass runSandboxCheckpointStoreConformance, runSandboxCheckpointForkConformance and runSandboxInstanceStoreConformance. The run store also passes runDurableRunFieldsConformance. - ./mcp: upstashMcpTaskStore (createMCPServer taskStore) and upstashMcpEventStore (MCP Events direct mode). upstashPersistence({ bucket }) now returns a required blobs store in its type, so it plugs straight into createSandboxSnapshots. Optional peers: @tanstack/ai-sandbox and @tanstack/ai-mcp >=1.0.0. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01CMK4etJMDDBFkjEyWDVCdy --- .changeset/tanstack-sandbox-mcp-stores.md | 14 + CLAUDE.md | 32 +- packages/tanstack-ai/README.md | 150 +++++- packages/tanstack-ai/package.json | 29 +- packages/tanstack-ai/src/mcp/event-store.ts | 146 ++++++ packages/tanstack-ai/src/mcp/index.ts | 6 + packages/tanstack-ai/src/mcp/mcp.test.ts | 141 ++++++ packages/tanstack-ai/src/mcp/task-store.ts | 60 +++ .../src/persistence/persistence.ts | 6 + .../src/sandbox/checkpoint-store.ts | 469 ++++++++++++++++++ packages/tanstack-ai/src/sandbox/index.ts | 7 + .../tanstack-ai/src/sandbox/instance-store.ts | 57 +++ .../tanstack-ai/src/sandbox/sandbox.test.ts | 119 +++++ packages/tanstack-ai/src/sandbox/validate.ts | 172 +++++++ .../tanstack-ai/src/sandbox/wiring.test.ts | 25 + packages/tanstack-ai/tsup.config.ts | 2 + pnpm-lock.yaml | 76 +++ 17 files changed, 1481 insertions(+), 30 deletions(-) create mode 100644 .changeset/tanstack-sandbox-mcp-stores.md create mode 100644 packages/tanstack-ai/src/mcp/event-store.ts create mode 100644 packages/tanstack-ai/src/mcp/index.ts create mode 100644 packages/tanstack-ai/src/mcp/mcp.test.ts create mode 100644 packages/tanstack-ai/src/mcp/task-store.ts create mode 100644 packages/tanstack-ai/src/sandbox/checkpoint-store.ts create mode 100644 packages/tanstack-ai/src/sandbox/index.ts create mode 100644 packages/tanstack-ai/src/sandbox/instance-store.ts create mode 100644 packages/tanstack-ai/src/sandbox/sandbox.test.ts create mode 100644 packages/tanstack-ai/src/sandbox/validate.ts create mode 100644 packages/tanstack-ai/src/sandbox/wiring.test.ts diff --git a/.changeset/tanstack-sandbox-mcp-stores.md b/.changeset/tanstack-sandbox-mcp-stores.md new file mode 100644 index 00000000..d2303680 --- /dev/null +++ b/.changeset/tanstack-sandbox-mcp-stores.md @@ -0,0 +1,14 @@ +--- +"@upstash/agentkit-tanstack-ai": minor +--- + +Add stores for TanStack AI's sandbox and MCP packages, all on Upstash Redis. + +- `@upstash/agentkit-tanstack-ai/sandbox` (needs `@tanstack/ai-sandbox` >= 1.0.0): + - `upstashSandboxCheckpointStore()` is a `SandboxCheckpointStore` for portable sandbox snapshots. It supports `forkFromCheckpoint`, which branches a thread, files and transcript included. + - `upstashSandboxInstanceStore()` lets any server resume a thread's sandbox. + - Both pass TanStack's sandbox conformance suites. +- `@upstash/agentkit-tanstack-ai/mcp` (needs `@tanstack/ai-mcp` >= 1.0.0): + - `upstashMcpTaskStore()` is the `taskStore` for `createMCPServer`. + - `upstashMcpEventStore()` stores MCP Events webhook subscriptions, shared across instances. +- `upstashPersistence({ bucket })` is now typed with a required `blobs` store, so you can pass it straight to `createSandboxSnapshots()`. diff --git a/CLAUDE.md b/CLAUDE.md index b9cc284d..13fb0e29 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1601,7 +1601,9 @@ find node_modules -path "*/zod/package.json" | while read f; do echo "$f $(node - **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`). + `@tanstack/ai-persistence`), `./memory` (`src/memory/index.ts`, needs `@tanstack/ai-memory`), `./byok` + (`src/byok/`), `./sandbox` (`src/sandbox/index.ts`, needs `@tanstack/ai-sandbox`) and `./mcp` + (`src/mcp/index.ts`, needs `@tanstack/ai-mcp`). 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 @@ -1706,11 +1708,33 @@ find node_modules -path "*/zod/package.json" | while read f; do echo "$f $(node `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`. + need_tools/replay loop), QStash-backed background runs writing to `upstashStream`, a QStash schedule for + `reapDetachedRuns`, `@tanstack/ai-code-mode`'s `AgentStore` and `@tanstack/ai-code-mode-snippets`' + `SnippetStorage` (both small keyed stores; `ai-code-mode` peers the pre-1.0 `@tanstack/ai-harness`). + **Deliberately not built:** `@tanstack/ai-mcp`'s `McpSessionStore` (MCP Apps) — its + `McpServerDescriptor` carries functions (`toolFilter`, `needsApproval`, `toolName`) and transport + objects, so it cannot round-trip through Redis. - **ai 7.0.169 + @tanstack/ai 1.4.0 + @tanstack/ai-openai 1.1.1 (2026-11-11):** devDependency/example pins only (no published-content change, no changeset); `@ai-sdk/*` ranges moved with `pnpm -r update "@ai-sdk/*"` (one `ai@7.0.169` in the lockfile) and `pnpm dedupe` dropped a stale second `@tanstack/ai@1.1.0` that `@tanstack/ai-client@1.1.0` (`^1.1.0`) kept in the lockfile. eve stays on 0.91.0: no stable release since, npm `latest` is `1.0.0-rc.9` (not adopted); published `@upstash/agentkit-eve-extension@0.19.0` + `@upstash/agentkit-eve@0.19.0` install on it, `eve build` clean and the mocked `eve eval` passes 4/4. A `skipLibCheck: false` consumer of `ai` errors inside `@ai-sdk/provider` (`TS7016 … module 'json-schema'`, no `@types/json-schema` dependency) on 7.0.163 and 7.0.169 alike: upstream, fixed by `npm i -D @types/json-schema`. The free throwaway Redis caps search indexes at 1: run index-using test files one at a time with `FLUSHDB` in between (`FLUSHDB` drops the indexes too; `SEARCH.DROPINDEX` is not available). - **ai 7.0.163 (2026-11-05):** pins bumped from 7.0.159 (no code change); `examples/ai-sdk-demo`'s `@ai-sdk/react` raised to `^4.0.167` because 4.0.160 depends on an *exact* `ai@7.0.156`, which split the `UIMessage` type from the demo's `ai@7.0.163` and failed `next build` (`SteeringMessageUIPart`). Re-checked 2026-11-05: published `@upstash/agentkit-eve-extension@0.19.0` / `@upstash/agentkit-eve@0.19.0` build clean on `eve@1.0.0-rc.5` (the 0.91.0 build stays); TanStack AI 1.x (ai 1.1.0, ai-persistence 1.0.1, ai-memory 1.1.0) still not adopted (same `threadId`/`blobKey` test-file errors). - **TanStack AI 1.x applied (2026-11-05):** `@tanstack/ai` 1.1.0, `ai-persistence` 1.0.1, `ai-memory` 1.1.0 (devDeps, exact) and the tanstack-ai-demo's `ai-client` 1.1.0 / `ai-react` 1.0.2 / `ai-openai` 1.1.0. `src` needed no change; only tests did: `chat()` calls pass `threadId` (`middleware.test.ts`, `search-tools.test.ts`) and `ArtifactRecord` fixtures carry `blobKey` (`persistence.test.ts`). The peer floors (`>=0.65.0` / `>=0.8.0` / `>=0.2.8`) are untouched: the source is unchanged and still works on 0.x, and published `@upstash/agentkit-tanstack-ai@0.2.0` imports on 1.x, so no changeset. Upstream `@tanstack/ai@1.1.0`'s `spawn.d.ts` still has the broken `import('../../..')` (TS2307 with `skipLibCheck: false`). The 1.0 contracts `SessionIndexStore` and `CredentialStore` are implemented in the follow-up below. +- **Sandbox + MCP stores (2026-11-11, `@upstash/agentkit-tanstack-ai` minor):** `./sandbox` = + `upstashSandboxCheckpointStore` (`SandboxCheckpointStore` + `forkFromCheckpoint`) and + `upstashSandboxInstanceStore`; `./mcp` = `upstashMcpTaskStore` (`TaskStore` for `createMCPServer`) and + `upstashMcpEventStore` (`MCPEventStore`, MCP Events direct mode). Optional peers `@tanstack/ai-sandbox` + / `@tanstack/ai-mcp` `>=1.0.0` (checkpoint/instance `.d.ts` and `MCPEventSubscription` are identical from + 1.0.0 to 1.2.x). Gotchas: (1) writer leases are `"||"` strings compared against + the CALLER's clock in Lua (ARGV), not a Redis TTL — the conformance suite injects a fake `now`; the PX is + lease + 1 day, GC only. (2) The fork must refuse a destination with ANY persistence state (transcript, + `threadRuns`, `threadInterrupts`, `threadGenerationRuns`, `threadArtifacts`) and write the transcript into + `${prefix}:thread:` in the same script — so the checkpoint store must share `upstashPersistence()`'s + `prefix` (thread ids are NOT URI-encoded in those persistence keys; our own `sandbox:*` keys are). + (3) `runSandboxCheckpointForkConformance` creates the fork promise before reading destination state over + the network, so a live-Redis rejection shows up as a vitest "unhandled error"; the test wraps + `forkFromCheckpoint` with a no-op `.catch` (assertions unchanged). (4) Heads, index members, id-index + values and set members are `"~" + id`, because `@upstash/redis` JSON-parses `"007"` to `7`. (5) The + validation helpers mirror `@tanstack/ai-sandbox`'s unexported `validateEntries`/`validateArtifacts` + (`src/sandbox/validate.ts`) — re-diff on sandbox bumps. (6) `upstashPersistence` gained an overload: with + `bucket` the return type has a required `blobs`, which `createSandboxSnapshots` needs + (`src/sandbox/wiring.test.ts` pins it). Also runs `runDurableRunFieldsConformance` on the run store. - **TanStack AI 1.0 contracts implemented (2026-11-05, `@upstash/agentkit-tanstack-ai` minor):** `stores.log` (`persistence/log-store.ts`, see below), `stores.sessions` (`SessionIndexStore`, `persistence/sessions.ts`) and `./byok` (`upstashCredentials`, `byok/credentials.ts`, `CredentialStore`). **Peer floors raised to `@tanstack/ai >=1.0.0` and `@tanstack/ai-persistence >=1.0.0`** (verified by packing 0.71.0/0.10.1: neither has `SessionIndexStore` / `@tanstack/ai/byok`; `ai-memory` floor untouched). Sessions: doc `{r,e}` per `(tenant, thread)` + one ZSET per tenant (score `updatedAt`, member `~` so the client never JSON-parses a numeric-looking id); `list` walks the ZSET in batches of whole score groups and re-reads a tie group in full so the `threadId` UTF-8 order is exact, filtering after loading; `ZRANGE BYSCORE REV` takes `(max, min)` (cast through `number`, the SDK types forbid `(`-exclusive strings). Revisions come from an `INCR` counter, never reused (the credential contract requires that; sessions may reuse). Passes the 19 `sessions` cases of `runPersistenceConformance` (they run only when the store is present). Credentials: one hash per owner (`%none` segment = undefined tenant/user, distinct from `""`), value `{r,t,x?,c | iv,ct}` with type/expiry in the clear so `list` never decrypts; optional AES-256-GCM `encryptionKey`, AAD = owner key + id (a copied ciphertext does not decrypt). `storedProviderKeys` lives in `@tanstack/ai/byok/server`, not `/byok`. **`stores.log` (`LogStore`):** per log a meta hash (head, first, epoch, owner, harness, exp, rel, wake, aff, del, laneH, oseq), a records ZSET (score seq, member `:`), snapshot string (`seq` must stay the first JSON key: `truncate` reads it with a Lua pattern), offers hash + info hash + pending ZSET, and a tick counter; shared `logDue` ZSET (when each log needs an owner) and `logHosts` hash. Every op is one Lua script over the same nine declared keys; the store clock is Redis `TIME` inside the script; numbers written from Lua go through `%.0f` (plain `tostring` prints epoch ms in scientific notation); ids returned from Lua are `~`-prefixed because the client JSON-parses numeric-looking strings (`"007"` -> 7). `subscribe` = immediate local callback + a `reported` and an `observed` tick per log so each counter value is reported once (a baseline read that already contains a local write must not mark it reported: that swallowed the notification), plus a 250 ms tick poll for other machines (non-overlapping; the first successful read is the baseline, so a failed first read is retried instead of silencing the subscription). `changed()` skips a counter the poll already reported (deterministic repro in tests: delay the append reply with a Proxy), which fixed a double notification (a poll can see the tick before the write reply arrives; it failed the `calls a listener after each append` conformance case in CI). The baseline read is still async, so a write by another machine landing in the first round trip after `subscribe` can go unreported (listeners are wake-ups, hosts reconcile with `read`). `append` rejects more than 10,000 records (a 50,000-record batch outran the Lua time limit and left part of the writes). Known follow-ups: `listLogsWithoutOwner` re-checks entries skipped by `harnesses` / a live pin on every call; consumed offers, old hosts and `deleteLog` tombstones are never cleaned up; `read` / `listOfferedInputs` without a limit return everything. Conformance needs the `twin` option (a second store set on the same keyspace) for the 4 `log across instances` cases; all 56 log cases and 19 sessions cases pass. Not implemented: `ownership` TTL-based cleanup of finished logs (logs and their keys live until `deleteLog`). - **eve 0.91.0 + ai 7.0.159 (2026-10-31):** `eve ^0.91.0` / `ai 7.0.159` everywhere; the eve-extension dist is rebuilt on 0.91.0 (manifest `builtWithEve: 0.91.0`) and its `eve` peer floor is `>=0.91.0 || >=1.0.0-0` — verified by packing the 0.91.0-built dist into consumers: eve 0.90.0 refuses it (`was built with eve 0.91.0. Upgrade eve to 0.91.0 or later`), 0.91.0 and `1.0.0-rc.1` build clean. Both eve peers carry `|| >=1.0.0-0` because semver never matches a prerelease against `>=0.x`: `npm install eve@latest` (`1.0.0-rc.1`) failed ERESOLVE against the published `@upstash/agentkit-eve@0.18.0` (`peer eve@">=0.87.0"`). **Not adopted:** `eve@1.0.0-rc.1` (a dist built on it stamps `builtWithEve: 1.0.0-rc.1` and would refuse every 0.x eve, so build on stable) and **TanStack AI 1.0** (`@tanstack/ai` 1.0.1, `-persistence` 1.0.1, `-memory` 1.0.0): `chat()` now requires `threadId` and `ArtifactRecord` requires `blobKey`, so `packages/tanstack-ai` test files fail `tsc` and the middleware tests throw `chat() needs a threadId` (source + build are fine); pins stay on 0.71.0 / 0.10.1 / 0.2.16. diff --git a/packages/tanstack-ai/README.md b/packages/tanstack-ai/README.md index 287ce4a0..2eac591a 100644 --- a/packages/tanstack-ai/README.md +++ b/packages/tanstack-ai/README.md @@ -6,28 +6,33 @@ resumable streaming, distributed locks and memory, and ships in-memory implement 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`). +| 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 | +| `createSandboxSnapshots()` / `withSandbox()` stores | `upstashSandboxCheckpointStore()`, `upstashSandboxInstanceStore()` | Portable sandbox snapshots with history and forks, and any server can resume a thread's sandbox | +| `createMCPServer({ taskStore })`, MCP Events `{ store }` | `upstashMcpTaskStore()`, `upstashMcpEventStore()` | Long-running MCP tasks and webhook subscriptions that survive restarts and work across instances | +| 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()`, `upstashMemory()` and the sandbox stores are checked against TanStack AI's +own conformance suites (`runPersistenceConformance`, `runMemoryAdapterContract`, +`runSandboxCheckpointStoreConformance`, `runSandboxCheckpointForkConformance`, +`runSandboxInstanceStoreConformance`, `runDurableRunFieldsConformance`). ## Install ```bash npm install @upstash/agentkit-tanstack-ai @tanstack/ai # plus, for the features you use: -npm install @tanstack/ai-persistence @tanstack/ai-memory +npm install @tanstack/ai-persistence @tanstack/ai-memory @tanstack/ai-sandbox @tanstack/ai-mcp ``` -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. +Persistence, memory, sandbox and MCP stores have their own entry points +(`@upstash/agentkit-tanstack-ai/persistence`, `/memory`, `/sandbox`, `/mcp`, plus `/byok`), 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. @@ -75,7 +80,14 @@ messages, and `delete` removes only the entry. Use it on its own with `upstashSe ```ts const { sessions } = upstashPersistence().stores; -await sessions.upsert({ tenantId, threadId, members: [{ userId, role: "owner" }], title, createdAt: now, updatedAt: now }); +await sessions.upsert({ + tenantId, + threadId, + members: [{ userId, role: "owner" }], + title, + createdAt: now, + updatedAt: now, +}); const { entries, cursor } = await sessions.list({ tenantId, memberId: userId, limit: 20 }); ``` @@ -112,11 +124,11 @@ Each chunk is appended to a Redis Stream before it is delivered. A client that r `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 | +| 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 @@ -172,7 +184,11 @@ import { storedProviderKeys } from "@tanstack/ai/byok/server"; import { upstashCredentials } from "@upstash/agentkit-tanstack-ai/byok"; const credentials = upstashCredentials({ encryptionKey: process.env.CREDENTIALS_KEY }); -await credentials.set({ scope: { tenantId, userId }, id: "openai", credential: { type: "api_key", value } }); +await credentials.set({ + scope: { tenantId, userId }, + id: "openai", + credential: { type: "api_key", value }, +}); const keys = storedProviderKeys(credentials, { tenantId, userId }); // the user's key, else the tenant's ``` @@ -188,7 +204,7 @@ JSON. Things to know: - The ciphertext is bound to the key `prefix`, the owner, the id, and the clear-text type and expiry shown by `list`. **Changing `prefix`, or copying rows to another prefix, makes every encrypted row undecryptable**; editing the type or expiry in Redis makes that row fail to decrypt. -- Credentials saved before you added a key are still read as plaintext. Anyone who can *write* to your Redis +- Credentials saved before you added a key are still read as plaintext. Anyone who can _write_ to your Redis could therefore swap an encrypted credential for a plaintext one. After re-saving your credentials with the key, set `requireEncryption: true` to refuse plaintext rows. - The revision is not part of the ciphertext, so someone with write access to your Redis can replay an older @@ -196,6 +212,96 @@ JSON. Things to know: - Expired `secret_answer` credentials are not removed by the store; delete them after reading, as the harness does. +## Sandbox snapshots and instances + +```ts +import { createSandboxSnapshots, withSandbox } from "@tanstack/ai-sandbox"; +import { withPersistence } from "@tanstack/ai-persistence"; +import { Bucket } from "@upstash/blob"; +import { upstashLocks } from "@upstash/agentkit-tanstack-ai"; +import { upstashPersistence } from "@upstash/agentkit-tanstack-ai/persistence"; +import { + upstashSandboxCheckpointStore, + upstashSandboxInstanceStore, +} from "@upstash/agentkit-tanstack-ai/sandbox"; + +const persistence = upstashPersistence({ bucket: Bucket.fromEnv() }); // file bytes go to Upstash Blob +const instances = upstashSandboxInstanceStore(); +const snapshots = createSandboxSnapshots({ + persistence, + checkpoints: upstashSandboxCheckpointStore(), // same Redis + prefix as `persistence` + instances, + locks: upstashLocks(), +}); + +chat({ + adapter, + messages, + threadId, + middleware: [ + withPersistence(snapshots.persistence), + withSandbox(sandbox, { instances, snapshots }), + ], +}); + +// Branch a conversation, files included, from any earlier checkpoint: +await snapshots.fork({ threadId, checkpointId, destinationThreadId: "thread-2" }); +``` + +Portable snapshots let a later run rebuild a sandbox's files after the provider sandbox is gone. +Together with `upstashPersistence({ bucket })` and `upstashLocks()`, every store they need is on Upstash. +Pair them with TanStack's own Upstash Box provider, `@tanstack/ai-sandbox-upstash-box`, for the sandbox +itself. + +- **`upstashSandboxCheckpointStore()`**: the `SandboxCheckpointStore`. Each checkpoint is one immutable + document; each thread has a head pointer and an index. Appends are compare-and-set on the head, and every + write first checks the writer lease (owner token, fence and expiry), so a writer that lost its lease + changes nothing. Blob keys are reference-counted for cleanup (`listBlobReferences`). It supports + `forkFromCheckpoint`: in one script, it copies a checkpoint into an empty thread and writes that thread's + transcript. That's why it must use the same `redis` and `prefix` as `upstashPersistence()`. Pass the same + `messagesTtlSeconds` too, if you set one. +- **`upstashSandboxInstanceStore()`**: the `SandboxInstanceStore` that maps a sandbox key to the + provider's sandbox id and latest snapshot, so any server can resume the sandbox. Optional `ttlSeconds`. + +Needs `@tanstack/ai-sandbox` >= 1.0.0. + +## MCP task and event stores + +```ts +import { createMCPClient } from "@tanstack/ai-mcp"; +import { createMCPServer } from "@tanstack/ai-mcp/server"; +import { upstashMcpEventStore, upstashMcpTaskStore } from "@upstash/agentkit-tanstack-ai/mcp"; + +// Server: long-running tools as MCP tasks, readable from any instance. +const server = createMCPServer({ + name: "tools", + version: "1.0.0", + tools, + taskStore: upstashMcpTaskStore(), +}); + +// Client: MCP Events webhook subscriptions (direct mode), shared by every instance. +const mcp = await createMCPClient({ + transport: { type: "http", url: "https://mcp.example.com/mcp" }, + experimental_events: { store: upstashMcpEventStore() }, +}); +``` + +- **`upstashMcpTaskStore()`**: the `TaskStore` for `createMCPServer`. Task records survive restarts, and + `tasks/get` / `tasks/result` work on any instance. The optional `ttlSeconds` expires finished records. + Cancelling a task still stops its work only on the instance that is running it, because TanStack keeps + the abort controller in memory. +- **`upstashMcpEventStore()`**: the `MCPEventStore` for MCP Events. Subscriptions, cursors and refresh + deadlines are shared, so a webhook can land on any server. `update` is a compare-and-set, so concurrent + cursor updates never lose a field. The subscription's webhook `secret` is stored as is. + +Needs `@tanstack/ai-mcp` >= 1.0.0. Not provided: `McpSessionStore` (MCP Apps). Its server descriptors +carry functions (`toolFilter`, `needsApproval`, ...) that cannot be stored in Redis. + +Are you building an MCP _server_ with the official `@modelcontextprotocol/server` SDK instead of TanStack? +[`@upstash/mcp-toolkit`](https://github.com/upstash/agentkit/tree/main/packages/mcp-toolkit) gives you +durable tasks and events there, with QStash for background execution and webhook delivery. + ## Tool caching and rate limiting ```ts diff --git a/packages/tanstack-ai/package.json b/packages/tanstack-ai/package.json index 362859d1..ddb20d8d 100644 --- a/packages/tanstack-ai/package.json +++ b/packages/tanstack-ai/package.json @@ -32,6 +32,14 @@ "./byok": { "types": "./dist/byok.d.ts", "import": "./dist/byok.js" + }, + "./sandbox": { + "types": "./dist/sandbox.d.ts", + "import": "./dist/sandbox.js" + }, + "./mcp": { + "types": "./dist/mcp.d.ts", + "import": "./dist/mcp.js" } }, "files": [ @@ -54,7 +62,10 @@ "resumable-stream", "memory", "lock", - "rag" + "rag", + "sandbox", + "checkpoint", + "mcp" ], "dependencies": { "@upstash/agentkit-sdk": "workspace:*", @@ -63,16 +74,20 @@ }, "devDependencies": { "@tanstack/ai": "1.4.0", + "@tanstack/ai-mcp": "1.2.0", "@tanstack/ai-memory": "1.1.0", "@tanstack/ai-persistence": "1.0.1", - "dotenv": "^16.4.5", - "@upstash/blob": "0.0.7" + "@tanstack/ai-sandbox": "1.2.1", + "@upstash/blob": "0.0.7", + "dotenv": "^16.4.5" }, "peerDependencies": { "@tanstack/ai": ">=1.0.0", "@tanstack/ai-memory": ">=0.2.8", "@tanstack/ai-persistence": ">=1.0.0", - "@upstash/blob": ">=0.0.7" + "@upstash/blob": ">=0.0.7", + "@tanstack/ai-sandbox": ">=1.0.0", + "@tanstack/ai-mcp": ">=1.0.0" }, "peerDependenciesMeta": { "@tanstack/ai-memory": { @@ -83,6 +98,12 @@ }, "@upstash/blob": { "optional": true + }, + "@tanstack/ai-sandbox": { + "optional": true + }, + "@tanstack/ai-mcp": { + "optional": true } } } diff --git a/packages/tanstack-ai/src/mcp/event-store.ts b/packages/tanstack-ai/src/mcp/event-store.ts new file mode 100644 index 00000000..99f57d1f --- /dev/null +++ b/packages/tanstack-ai/src/mcp/event-store.ts @@ -0,0 +1,146 @@ +import type { Redis } from "@upstash/redis"; +import { Redis as RedisClient } from "@upstash/redis"; +import type { MCPEventStore, MCPEventSubscription } from "@tanstack/ai-mcp/events"; +import { addTelemetry } from "../telemetry.js"; +import { KEY_LOCKING } from "../persistence/records.js"; + +export interface UpstashMcpEventStoreConfig { + /** Upstash Redis client. Defaults to `Redis.fromEnv()`. */ + redis?: Redis; + /** Base key prefix. Defaults to `agentkit:tanstack`. */ + prefix?: string; + /** Set `false` to opt out of the telemetry header. */ + enableTelemetry?: boolean; +} + +// Index values and set members are "~" + key, so the client never JSON-parses a numeric-looking key. + +/** + * Write a subscription and keep the id index in step. + * KEYS: document, id→key hash, key→id hash, key set. + * ARGV: key, document JSON, expected revision ("" = unconditional), id ("" = none). + * Returns the new revision, or -1 if the revision moved (a concurrent `update`). + */ +const WRITE = `${KEY_LOCKING}if ARGV[3] ~= "" and tostring(redis.call("HGET", KEYS[1], "r")) ~= ARGV[3] then return -1 end +local prev = redis.call("HGET", KEYS[3], ARGV[1]) +if prev and prev ~= ARGV[4] then redis.call("HDEL", KEYS[2], prev) end +local rev = redis.call("HINCRBY", KEYS[1], "r", 1) +redis.call("HSET", KEYS[1], "d", ARGV[2]) +if ARGV[4] ~= "" then + redis.call("HSET", KEYS[2], ARGV[4], "~" .. ARGV[1]) + redis.call("HSET", KEYS[3], ARGV[1], ARGV[4]) +else + redis.call("HDEL", KEYS[3], ARGV[1]) +end +redis.call("SADD", KEYS[4], "~" .. ARGV[1]) +return rev`; + +/** KEYS: document, id→key hash, key→id hash, key set. ARGV: key. */ +const DELETE = `${KEY_LOCKING}local prev = redis.call("HGET", KEYS[3], ARGV[1]) +if prev then redis.call("HDEL", KEYS[2], prev) end +redis.call("HDEL", KEYS[3], ARGV[1]) +redis.call("DEL", KEYS[1]) +redis.call("SREM", KEYS[4], "~" .. ARGV[1]) +return 1`; + +interface Stored { + d: MCPEventSubscription; + r: number; +} + +/** Subscriptions per Redis round trip while `list` walks the store. */ +const LIST_BATCH = 100; + +/** + * `MCPEventStore` (MCP Events in direct mode, `@tanstack/ai-mcp/events`) on Upstash Redis: webhook + * subscriptions, their cursors and refresh deadlines survive restarts and are shared by every + * instance, so a webhook can land on any server. One hash per subscription (document + revision), + * plus an id index; `update` is a compare-and-set on the revision, retried on a race, so two + * instances advancing the same cursor never lose a field. + * + * The subscription's webhook `secret` is stored as is: use a Redis database you would trust with + * other credentials. + * + * @example + * ```ts + * import { createMCPClient } from "@tanstack/ai-mcp"; + * import { upstashMcpEventStore } from "@upstash/agentkit-tanstack-ai/mcp"; + * + * const mcp = await createMCPClient({ + * transport: { type: "http", url: "https://mcp.example.com/mcp" }, + * experimental_events: { store: upstashMcpEventStore() }, + * }); + * ``` + */ +export function upstashMcpEventStore(config: UpstashMcpEventStoreConfig = {}): MCPEventStore { + const redis = config.redis ?? RedisClient.fromEnv(); + addTelemetry(redis, config.enableTelemetry); + const p = config.prefix ?? "agentkit:tanstack"; + const k = { + doc: (key: string) => `${p}:mcp:eventSub:${encodeURIComponent(key)}`, + ids: `${p}:mcp:eventSubIds`, + keyIds: `${p}:mcp:eventSubKeyIds`, + all: `${p}:mcp:eventSubs`, + }; + const keys = (key: string) => [k.doc(key), k.ids, k.keyIds, k.all]; + + const load = async (key: string): Promise => { + const hash = await redis.hgetall<{ d: MCPEventSubscription; r: number }>(k.doc(key)); + return hash && hash.d !== undefined ? { d: hash.d, r: Number(hash.r) } : undefined; + }; + const idOf = (s: MCPEventSubscription) => (typeof s.id === "string" ? s.id : ""); + + return { + async get(key) { + return (await load(key))?.d; + }, + + async getById(id) { + const key = await redis.hget(k.ids, id); + return key === null || key === undefined ? undefined : (await load(String(key).slice(1)))?.d; + }, + + async set(subscription) { + await redis.eval(WRITE, keys(subscription.key), [ + subscription.key, + JSON.stringify(subscription), + "", + idOf(subscription), + ]); + }, + + async update(key, patch) { + for (;;) { + const stored = await load(key); + if (!stored) return; + // `Object.assign` semantics: a field present in the patch replaces the stored one, and a + // field set to `undefined` is dropped (JSON leaves it out). + const next = { ...stored.d, ...patch, key } as MCPEventSubscription; + const id = patch.id === undefined && stored.d.id !== undefined ? stored.d.id : next.id; + const result = (await redis.eval(WRITE, keys(key), [ + key, + JSON.stringify(next), + String(stored.r), + typeof id === "string" ? id : "", + ])) as number; + if (result !== -1) return; + } + }, + + async delete(key) { + await redis.eval(DELETE, keys(key), [key]); + }, + + async *list() { + let cursor = "0"; + do { + const [next, members] = await redis.sscan(k.all, cursor, { count: LIST_BATCH }); + cursor = String(next); + for (const member of members) { + const stored = await load(String(member).slice(1)); + if (stored) yield stored.d; + } + } while (cursor !== "0"); + }, + }; +} diff --git a/packages/tanstack-ai/src/mcp/index.ts b/packages/tanstack-ai/src/mcp/index.ts new file mode 100644 index 00000000..4bfe3ede --- /dev/null +++ b/packages/tanstack-ai/src/mcp/index.ts @@ -0,0 +1,6 @@ +// `@upstash/agentkit-tanstack-ai/mcp` — `@tanstack/ai-mcp` stores on Upstash Redis: MCP server task +// records and MCP Events webhook subscriptions. Needs `@tanstack/ai-mcp`. +export { upstashMcpTaskStore } from "./task-store.js"; +export type { UpstashMcpTaskStoreConfig } from "./task-store.js"; +export { upstashMcpEventStore } from "./event-store.js"; +export type { UpstashMcpEventStoreConfig } from "./event-store.js"; diff --git a/packages/tanstack-ai/src/mcp/mcp.test.ts b/packages/tanstack-ai/src/mcp/mcp.test.ts new file mode 100644 index 00000000..7b2f06c1 --- /dev/null +++ b/packages/tanstack-ai/src/mcp/mcp.test.ts @@ -0,0 +1,141 @@ +import { afterAll, describe, expect, it } from "vitest"; +import type { MCPEventSubscription } from "@tanstack/ai-mcp/events"; +import { upstashMcpEventStore } from "./event-store.js"; +import { upstashMcpTaskStore } from "./task-store.js"; +import { cleanupKeys, hasRedisCreds, testRedis, uniquePrefix } from "../testing/test-support.js"; + +const subscription = (key: string, id?: string): MCPEventSubscription => ({ + key, + ...(id === undefined ? {} : { id }), + name: "issues.opened", + arguments: { repo: "upstash/agentkit", labels: [] }, + definition: { name: "issues.opened" } as MCPEventSubscription["definition"], + delivery: { mode: "webhook", url: "https://example.com/hook", secret: "s3cr3t" }, + refreshBefore: null, + cursor: null, + truncated: false, + status: "pending", +}); + +// No upstream testkit exists for these contracts; the cases mirror `inMemoryMCPEventStore()` and +// `inMemoryTaskStore()` from `@tanstack/ai-mcp`. +describe.skipIf(!hasRedisCreds)("MCP stores (live Redis)", () => { + const redis = testRedis(); + const prefix = uniquePrefix("tmcp"); + let n = 0; + const next = () => `${prefix}:${++n}`; + + afterAll(async () => { + await cleanupKeys(redis, prefix); + }); + + describe("upstashMcpTaskStore", () => { + it("round-trips values, returns null when missing, and deletes", async () => { + const tasks = upstashMcpTaskStore({ redis, prefix: next() }); + expect(await tasks.get("t1")).toBeNull(); + const record = { + taskId: "t1", + status: "working", + ttl: null, + createdAt: "2026-11-11T00:00:00Z", + }; + await tasks.set("t1", record); + expect(await tasks.get("t1")).toEqual(record); + await tasks.set("t1", { ...record, status: "completed" }); + expect(await tasks.get("t1")).toMatchObject({ status: "completed" }); + await tasks.delete("t1"); + expect(await tasks.get("t1")).toBeNull(); + }); + + it("keeps primitive values as they were", async () => { + const tasks = upstashMcpTaskStore({ redis, prefix: next() }); + await tasks.set("s", "123"); + await tasks.set("n", 123); + expect(await tasks.get("s")).toBe("123"); + expect(await tasks.get("n")).toBe(123); + }); + + it("expires records with ttlSeconds", async () => { + const p = next(); + const tasks = upstashMcpTaskStore({ redis, prefix: p, ttlSeconds: 60 }); + await tasks.set("t", { status: "working" }); + const ttl = await redis.ttl(`${p}:mcp:task:t`); + expect(ttl).toBeGreaterThan(0); + expect(ttl).toBeLessThanOrEqual(60); + }); + }); + + describe("upstashMcpEventStore", () => { + it("gets by key and by id, and returns undefined when missing", async () => { + const store = upstashMcpEventStore({ redis, prefix: next() }); + expect(await store.get("k1")).toBeUndefined(); + expect(await store.getById("id1")).toBeUndefined(); + await store.set(subscription("k1", "id1")); + expect(await store.get("k1")).toEqual(subscription("k1", "id1")); + expect(await store.getById("id1")).toEqual(subscription("k1", "id1")); + }); + + it("re-indexes when set replaces a subscription with another id", async () => { + const store = upstashMcpEventStore({ redis, prefix: next() }); + await store.set(subscription("k1", "old")); + await store.set(subscription("k1", "new")); + expect(await store.getById("old")).toBeUndefined(); + expect(await store.getById("new")).toMatchObject({ key: "k1" }); + await store.set(subscription("k1")); + expect(await store.getById("new")).toBeUndefined(); + }); + + it("update merges fields, assigns an id, ignores a missing key", async () => { + const store = upstashMcpEventStore({ redis, prefix: next() }); + await store.update("missing", { cursor: "c" }); + expect(await store.get("missing")).toBeUndefined(); + + await store.set(subscription("k1")); + await store.update("k1", { id: "id1", status: "active", cursor: "c1" }); + expect(await store.get("k1")).toEqual({ + ...subscription("k1", "id1"), + status: "active", + cursor: "c1", + }); + expect(await store.getById("id1")).toMatchObject({ cursor: "c1" }); + + await store.update("k1", { id: "id2" }); + expect(await store.getById("id1")).toBeUndefined(); + expect(await store.getById("id2")).toMatchObject({ key: "k1", cursor: "c1" }); + }); + + it("keeps every field when updates race", async () => { + const store = upstashMcpEventStore({ redis, prefix: next() }); + await store.set(subscription("k1", "id1")); + await Promise.all([ + store.update("k1", { cursor: "c9" }), + store.update("k1", { refreshBefore: "2026-12-01T00:00:00Z" }), + store.update("k1", { status: "active" }), + store.update("k1", { truncated: true }), + ]); + expect(await store.get("k1")).toMatchObject({ + cursor: "c9", + refreshBefore: "2026-12-01T00:00:00Z", + status: "active", + truncated: true, + }); + }); + + it("deletes and lists, with numeric-looking keys and ids intact", async () => { + const store = upstashMcpEventStore({ redis, prefix: next() }); + await store.set(subscription("007", "42")); + await store.set(subscription("k2", "id2")); + expect(await store.getById("42")).toMatchObject({ key: "007" }); + const keys: string[] = []; + for await (const s of store.list!()) keys.push(s.key); + expect(keys.sort()).toEqual(["007", "k2"]); + + await store.delete("007"); + expect(await store.get("007")).toBeUndefined(); + expect(await store.getById("42")).toBeUndefined(); + const after: string[] = []; + for await (const s of store.list!()) after.push(s.key); + expect(after).toEqual(["k2"]); + }); + }); +}); diff --git a/packages/tanstack-ai/src/mcp/task-store.ts b/packages/tanstack-ai/src/mcp/task-store.ts new file mode 100644 index 00000000..7675a12f --- /dev/null +++ b/packages/tanstack-ai/src/mcp/task-store.ts @@ -0,0 +1,60 @@ +import type { Redis } from "@upstash/redis"; +import { Redis as RedisClient } from "@upstash/redis"; +import type { TaskStore } from "@tanstack/ai-mcp/server"; +import { addTelemetry } from "../telemetry.js"; + +export interface UpstashMcpTaskStoreConfig { + /** Upstash Redis client. Defaults to `Redis.fromEnv()`. */ + redis?: Redis; + /** Base key prefix. Defaults to `agentkit:tanstack`. */ + prefix?: string; + /** + * Expire a task record this many seconds after its last write. Omit to keep records until + * `delete`. Clients poll a task until it settles and then read its result, so keep it well above + * your longest task plus the time a client may take to come back for the result. + */ + ttlSeconds?: number; + /** Set `false` to opt out of the telemetry header. */ + enableTelemetry?: boolean; +} + +/** + * `TaskStore` for `createMCPServer({ taskStore })` from `@tanstack/ai-mcp/server` on Upstash Redis: + * MCP task records survive a restart and any server instance can answer `tasks/get` and + * `tasks/result` for a task another instance started. Values are stored as JSON. + * + * Cancelling still aborts the task's work only on the instance that runs it (TanStack keeps the + * `AbortController` in process); another instance records the cancellation. + * + * @example + * ```ts + * import { createMCPServer } from "@tanstack/ai-mcp/server"; + * import { upstashMcpTaskStore } from "@upstash/agentkit-tanstack-ai/mcp"; + * + * const server = createMCPServer({ name: "tools", version: "1.0.0", taskStore: upstashMcpTaskStore() }); + * ``` + */ +export function upstashMcpTaskStore( + config: UpstashMcpTaskStoreConfig = {}, +): TaskStore { + const redis = config.redis ?? RedisClient.fromEnv(); + addTelemetry(redis, config.enableTelemetry); + const p = config.prefix ?? "agentkit:tanstack"; + const key = (id: string) => `${p}:mcp:task:${encodeURIComponent(id)}`; + + return { + async get(id) { + // Wrapped as `{ v }` so a bare string or number round-trips without the client re-parsing it. + const doc = await redis.get<{ v: Value }>(key(id)); + return doc ? doc.v : null; + }, + async set(id, value) { + const doc = JSON.stringify({ v: value ?? null }); + if (config.ttlSeconds) await redis.set(key(id), doc, { ex: config.ttlSeconds }); + else await redis.set(key(id), doc); + }, + async delete(id) { + await redis.del(key(id)); + }, + }; +} diff --git a/packages/tanstack-ai/src/persistence/persistence.ts b/packages/tanstack-ai/src/persistence/persistence.ts index 2efbf0c9..1970b225 100644 --- a/packages/tanstack-ai/src/persistence/persistence.ts +++ b/packages/tanstack-ai/src/persistence/persistence.ts @@ -135,6 +135,12 @@ return 1`; * chat({ adapter, messages, threadId, middleware: [withPersistence(upstashPersistence())] }); * ``` */ +export function upstashPersistence( + config: UpstashPersistenceConfig & { bucket: BlobBucketLike }, +): AIPersistence; +export function upstashPersistence( + config?: UpstashPersistenceConfig, +): AIPersistence; export function upstashPersistence( config: UpstashPersistenceConfig = {}, ): AIPersistence { diff --git a/packages/tanstack-ai/src/sandbox/checkpoint-store.ts b/packages/tanstack-ai/src/sandbox/checkpoint-store.ts new file mode 100644 index 00000000..909b0647 --- /dev/null +++ b/packages/tanstack-ai/src/sandbox/checkpoint-store.ts @@ -0,0 +1,469 @@ +import type { Redis } from "@upstash/redis"; +import { Redis as RedisClient } from "@upstash/redis"; +import { + SandboxCheckpointConflictError, + SandboxCheckpointDuplicateIdError, + SandboxCheckpointError, + SandboxCheckpointInvalidEntryError, + SandboxCheckpointNotHeadError, + SandboxCheckpointParentMismatchError, + SandboxCheckpointWriterConflictError, + SandboxCheckpointWriterLostError, + type ForkCapableSandboxCheckpointStore, + type SandboxCheckpoint, + type SandboxCheckpointWriter, + type SandboxCheckpointWriterLease, +} from "@tanstack/ai-sandbox"; +import { addTelemetry } from "../telemetry.js"; +import { KEY_LOCKING } from "../persistence/records.js"; +import { + assertValidIdentifier, + blobKeysOf, + compareUtf8Bytes, + validateCheckpoint, +} from "./validate.js"; + +export interface UpstashSandboxCheckpointStoreConfig { + /** Upstash Redis client. Defaults to `Redis.fromEnv()`. */ + redis?: Redis; + /** + * Base key prefix. Defaults to `agentkit:tanstack`. Use the same `prefix` (and Redis) as + * `upstashPersistence()`: `forkFromCheckpoint` writes the forked conversation into that + * persistence's message store, in the same script that creates the fork root. + */ + prefix?: string; + /** + * The `messagesTtlSeconds` you pass to `upstashPersistence()`, so a forked transcript expires like + * any other. Omit to keep transcripts forever. + */ + messagesTtlSeconds?: number; + /** Clock for writer leases. Defaults to `Date.now`. */ + now?: () => number; + /** How long a writer lease lasts without a renew. @default 120_000 */ + leaseDurationMs?: number; + /** When the lease holder should renew. Must be below `leaseDurationMs`. @default 45_000 */ + renewAfterMs?: number; + /** Set `false` to opt out of the telemetry header. */ + enableTelemetry?: boolean; +} + +// A writer lease is stored as "||". Expiry is compared against the +// caller's clock (ARGV), not a Redis TTL, so a lease means the same thing to every host and to an +// injected test clock. The key still gets a generous PX so abandoned threads do not keep it forever. +const WRITER_OK = `local function writerOk(key, token, fence, now) + local w = redis.call("GET", key) + if not w then return false end + local t, f, e = string.match(w, "^([^|]*)|([^|]*)|([^|]*)$") + return t == token and f == fence and tonumber(e) > tonumber(now) +end +`; + +/** KEYS: writer, fence counter. ARGV: now, ownerToken, expiresAt, px. Returns the fence, or -1 if held. */ +const ACQUIRE = `${KEY_LOCKING}local w = redis.call("GET", KEYS[1]) +if w then + local e = string.match(w, "|([^|]*)$") + if e and tonumber(e) > tonumber(ARGV[1]) then return -1 end +end +local fence = redis.call("INCR", KEYS[2]) +redis.call("SET", KEYS[1], ARGV[2] .. "|" .. fence .. "|" .. ARGV[3], "PX", ARGV[4]) +return fence`; + +/** KEYS: writer. ARGV: ownerToken, fence, now, new expiresAt, px. Returns 1, or 0 if the lease is gone. */ +const RENEW = `${KEY_LOCKING}${WRITER_OK}if not writerOk(KEYS[1], ARGV[1], ARGV[2], ARGV[3]) then return 0 end +redis.call("SET", KEYS[1], ARGV[1] .. "|" .. ARGV[2] .. "|" .. ARGV[4], "PX", ARGV[5]) +return 1`; + +/** KEYS: writer. ARGV: ownerToken, fence. Deletes the lease only if it is still this one. */ +const RELEASE = `${KEY_LOCKING}local w = redis.call("GET", KEYS[1]) +if w and w:sub(1, #ARGV[1] + #ARGV[2] + 2) == ARGV[1] .. "|" .. ARGV[2] .. "|" then + redis.call("DEL", KEYS[1]) +end +return 1`; + +/** + * KEYS: writer, checkpoint doc, thread head, thread index, blob reference counts. + * ARGV: ownerToken, fence, now, expected head ("" = none), "1" if the parent matches the expected + * head, checkpoint JSON, createdAt, member ("~" + id), blob keys… + * Checks run in the reference store's order: writer, duplicate id, stale head, parent. + */ +const APPEND = `${KEY_LOCKING}${WRITER_OK}if not writerOk(KEYS[1], ARGV[1], ARGV[2], ARGV[3]) then return "WRITER_LOST" end +if redis.call("EXISTS", KEYS[2]) == 1 then return "DUPLICATE" end +local head = redis.call("GET", KEYS[3]) or "" +if head ~= ARGV[4] then return "STALE_HEAD:" .. head end +if ARGV[5] ~= "1" then return "PARENT" end +redis.call("SET", KEYS[2], ARGV[6]) +redis.call("SET", KEYS[3], ARGV[8]) +redis.call("ZADD", KEYS[4], ARGV[7], ARGV[8]) +for i = 9, #ARGV do redis.call("HINCRBY", KEYS[5], ARGV[i], 1) end +return "OK"`; + +/** + * KEYS: writer, thread head, checkpoint doc, thread index, blob reference counts. + * ARGV: ownerToken, fence, now, member ("~" + id). + */ +const DELETE_HEAD = `${KEY_LOCKING}${WRITER_OK}if not writerOk(KEYS[1], ARGV[1], ARGV[2], ARGV[3]) then return "WRITER_LOST" end +if (redis.call("GET", KEYS[2]) or "") ~= ARGV[4] then return "NOT_HEAD" end +local raw = redis.call("GET", KEYS[3]) +if not raw then return "MISSING" end +local cp = cjson.decode(raw) +redis.call("DEL", KEYS[3]) +redis.call("ZREM", KEYS[4], ARGV[4]) +local parent = cp.parentCheckpointId +if type(parent) == "string" and parent ~= "" then + redis.call("SET", KEYS[2], "~" .. parent) +else + redis.call("DEL", KEYS[2]) +end +local seen = {} +local function unref(key) + if type(key) ~= "string" or seen[key] then return end + seen[key] = true + if redis.call("HINCRBY", KEYS[5], key, -1) <= 0 then redis.call("HDEL", KEYS[5], key) end +end +if type(cp.files) == "table" then + for _, e in ipairs(cp.files) do if e.kind == "file" then unref(e.blobKey) end end +end +if type(cp.artifacts) == "table" then + for _, a in ipairs(cp.artifacts) do unref(a.blobKey) end +end +return "OK"`; + +/** + * KEYS: destination writer, source checkpoint doc, destination checkpoint doc, destination head, + * destination index, blob reference counts, then the destination thread's persistence keys that + * must be empty: transcript, runs, interrupts, generation runs, artifacts (KEYS[7..11]). + * ARGV: ownerToken, fence, now, checkpoint JSON, transcript JSON, member, createdAt, ttl ("0" = none), + * blob keys… + */ +const FORK = `${KEY_LOCKING}${WRITER_OK}if not writerOk(KEYS[1], ARGV[1], ARGV[2], ARGV[3]) then return "WRITER_LOST" end +if redis.call("EXISTS", KEYS[2]) == 0 then return "SOURCE_NOT_FOUND" end +if redis.call("EXISTS", KEYS[3]) == 1 or redis.call("EXISTS", KEYS[4]) == 1 + or redis.call("ZCARD", KEYS[5]) > 0 or redis.call("EXISTS", KEYS[7]) == 1 then + return "NOT_EMPTY" +end +for i = 8, 11 do + if redis.call("ZCARD", KEYS[i]) > 0 then return "NOT_EMPTY" end +end +redis.call("SET", KEYS[3], ARGV[4]) +redis.call("SET", KEYS[4], ARGV[6]) +redis.call("ZADD", KEYS[5], ARGV[7], ARGV[6]) +redis.call("JSON.SET", KEYS[7], "$", ARGV[5]) +if ARGV[8] ~= "0" then redis.call("EXPIRE", KEYS[7], ARGV[8]) end +for i = 9, #ARGV do redis.call("HINCRBY", KEYS[6], ARGV[i], 1) end +return "OK"`; + +/** A checkpoint is JSON by contract; copying through JSON also detaches it from the caller's object. */ +const copy = (checkpoint: SandboxCheckpoint): SandboxCheckpoint => + JSON.parse(JSON.stringify(checkpoint)); + +/** Thread heads and index members start with "~", so the client never JSON-parses a numeric id. */ +const member = (id: string) => `~${id}`; + +/** + * `SandboxCheckpointStore` (from `@tanstack/ai-sandbox`) on Upstash Redis, with the optional + * `forkFromCheckpoint` capability. A checkpoint is one immutable JSON document; each thread has a + * head pointer and a sorted index, and blob keys are reference-counted in one hash so a sweeper can + * delete unreferenced sandbox files. Every write is one Lua script that first checks the writer + * lease (owner token + fence + expiry), so a writer that lost its lease cannot change anything. + * + * `forkFromCheckpoint` copies a checkpoint into an empty thread as its root and writes the + * checkpoint's conversation as that thread's transcript in the same script; pass the `prefix` you + * give `upstashPersistence()`. Checked against `runSandboxCheckpointStoreConformance` and + * `runSandboxCheckpointForkConformance` from `@tanstack/ai-sandbox/testkit`. + * + * @example + * ```ts + * import { upstashPersistence } from "@upstash/agentkit-tanstack-ai/persistence"; + * import { upstashSandboxCheckpointStore } from "@upstash/agentkit-tanstack-ai/sandbox"; + * + * const persistence = upstashPersistence(); + * const checkpoints = upstashSandboxCheckpointStore(); // same Redis + prefix as `persistence` + * ``` + */ +export function upstashSandboxCheckpointStore( + config: UpstashSandboxCheckpointStoreConfig = {}, +): ForkCapableSandboxCheckpointStore { + const redis = config.redis ?? RedisClient.fromEnv(); + addTelemetry(redis, config.enableTelemetry); + const p = config.prefix ?? "agentkit:tanstack"; + const now = config.now ?? (() => Date.now()); + const leaseDurationMs = config.leaseDurationMs ?? 120_000; + const renewAfterMs = config.renewAfterMs ?? 45_000; + if (!Number.isFinite(leaseDurationMs) || leaseDurationMs <= 0) { + throw new Error("leaseDurationMs must be finite and positive"); + } + if (!Number.isFinite(renewAfterMs) || renewAfterMs <= 0 || renewAfterMs >= leaseDurationMs) { + throw new Error("renewAfterMs must be finite, positive, and less than leaseDurationMs"); + } + // The lease key outlives its lease by a day: expiry is decided by `expiresAt`, the PX only + // garbage-collects leases nobody released. + const px = String(Math.ceil(leaseDurationMs) + 86_400_000); + const ttl = String(config.messagesTtlSeconds ?? 0); + + const enc = encodeURIComponent; + const k = { + doc: (id: string) => `${p}:sandbox:checkpoint:${enc(id)}`, + head: (thread: string) => `${p}:sandbox:head:${enc(thread)}`, + index: (thread: string) => `${p}:sandbox:checkpoints:${enc(thread)}`, + writer: (thread: string) => `${p}:sandbox:writer:${enc(thread)}`, + fence: (thread: string) => `${p}:sandbox:fence:${enc(thread)}`, + refs: `${p}:sandbox:blobRefs`, + // `upstashPersistence()` keys, which a fork must find empty (and fills the transcript of). + transcript: (thread: string) => `${p}:thread:${thread}`, + threadRuns: (thread: string) => `${p}:threadRuns:${thread}`, + threadInterrupts: (thread: string) => `${p}:threadInterrupts:${thread}`, + threadGenerationRuns: (thread: string) => `${p}:threadGenerationRuns:${thread}`, + threadArtifacts: (thread: string) => `${p}:threadArtifacts:${thread}`, + }; + + const writerLost = (thread: string) => + new SandboxCheckpointWriterLostError( + `Checkpoint writer lease for thread '${thread}' is no longer current`, + ); + const writerArgs = (w: SandboxCheckpointWriter) => [w.ownerToken, String(w.fence), String(now())]; + + const load = async (id: string): Promise => + (await redis.get(k.doc(id))) ?? null; + + return { + async get(id) { + assertValidIdentifier(id, "Checkpoint id"); + return load(id); + }, + + async list(threadId) { + assertValidIdentifier(threadId, "Thread id"); + const members = await redis.zrange(k.index(threadId), 0, -1); + if (members.length === 0) return []; + const docs = await redis.mget<(SandboxCheckpoint | null)[]>( + ...members.map((m) => k.doc(String(m).slice(1))), + ); + return docs + .filter((d): d is SandboxCheckpoint => d !== null && d.threadId === threadId) + .sort((a, b) => a.createdAt - b.createdAt || compareUtf8Bytes(a.id, b.id)); + }, + + async getHead(threadId) { + assertValidIdentifier(threadId, "Thread id"); + const head = await redis.get(k.head(threadId)); + return head ? String(head).slice(1) : null; + }, + + async append({ checkpoint: input, expectedHeadId, writer }) { + const checkpoint = copy(input); + assertValidIdentifier(checkpoint.threadId, "Checkpoint thread id"); + assertValidIdentifier(writer.threadId, "Writer thread id"); + assertValidIdentifier(checkpoint.id, "Checkpoint id"); + if (expectedHeadId !== null) assertValidIdentifier(expectedHeadId, "Expected head id"); + if (checkpoint.parentCheckpointId != null) { + assertValidIdentifier(checkpoint.parentCheckpointId, "Parent checkpoint id"); + } + if (writer.threadId !== checkpoint.threadId) { + throw new SandboxCheckpointWriterLostError( + "Checkpoint writer thread does not match checkpoint thread", + ); + } + validateCheckpoint(checkpoint); + const parentCheckpointId = checkpoint.parentCheckpointId ?? null; + const stored: SandboxCheckpoint = { ...checkpoint, parentCheckpointId }; + const thread = stored.threadId; + + const result = (await redis.eval( + APPEND, + [k.writer(thread), k.doc(stored.id), k.head(thread), k.index(thread), k.refs], + [ + ...writerArgs(writer), + expectedHeadId === null ? "" : member(expectedHeadId), + parentCheckpointId === expectedHeadId ? "1" : "0", + JSON.stringify(stored), + String(stored.createdAt), + member(stored.id), + ...blobKeysOf(stored), + ], + )) as string; + if (result === "WRITER_LOST") throw writerLost(thread); + if (result === "DUPLICATE") { + throw new SandboxCheckpointDuplicateIdError(`Checkpoint '${stored.id}' already exists`); + } + if (result.startsWith("STALE_HEAD:")) { + const actual = result.slice("STALE_HEAD:".length); + throw new SandboxCheckpointConflictError( + `Expected head '${expectedHeadId}', but thread '${thread}' is at '${actual ? actual.slice(1) : null}'`, + ); + } + if (result === "PARENT") { + throw new SandboxCheckpointParentMismatchError( + `Checkpoint '${stored.id}' parent does not match expected head`, + ); + } + return { headId: stored.id }; + }, + + async deleteHead({ threadId, checkpointId, writer }) { + assertValidIdentifier(threadId, "Thread id"); + assertValidIdentifier(checkpointId, "Checkpoint id"); + assertValidIdentifier(writer.threadId, "Writer thread id"); + if (writer.threadId !== threadId) { + throw new SandboxCheckpointWriterLostError( + "Checkpoint writer thread does not match operation thread", + ); + } + const result = (await redis.eval( + DELETE_HEAD, + [k.writer(threadId), k.head(threadId), k.doc(checkpointId), k.index(threadId), k.refs], + [...writerArgs(writer), member(checkpointId)], + )) as string; + if (result === "WRITER_LOST") throw writerLost(threadId); + if (result === "NOT_HEAD") { + throw new SandboxCheckpointNotHeadError( + `Checkpoint '${checkpointId}' is not the current head of thread '${threadId}'`, + ); + } + if (result === "MISSING") { + throw new SandboxCheckpointNotHeadError(`Checkpoint '${checkpointId}' does not exist`); + } + }, + + async acquireWriter(threadId): Promise { + assertValidIdentifier(threadId, "Thread id"); + const ownerToken = globalThis.crypto.randomUUID(); + const lease = { expiresAt: now() + leaseDurationMs }; + const fence = (await redis.eval( + ACQUIRE, + [k.writer(threadId), k.fence(threadId)], + [String(now()), ownerToken, String(lease.expiresAt), px], + )) as number; + if (fence === -1) { + throw new SandboxCheckpointWriterConflictError( + `Thread '${threadId}' already has an active checkpoint writer`, + ); + } + return { + threadId, + ownerToken, + fence: Number(fence), + get expiresAt() { + return lease.expiresAt; + }, + renewAfterMs, + renew: async () => { + const expiresAt = now() + leaseDurationMs; + const ok = (await redis.eval( + RENEW, + [k.writer(threadId)], + [ownerToken, String(fence), String(now()), String(expiresAt), px], + )) as number; + if (ok !== 1) throw writerLost(threadId); + lease.expiresAt = expiresAt; + return { expiresAt }; + }, + release: async () => { + await redis.eval(RELEASE, [k.writer(threadId)], [ownerToken, String(fence)]); + }, + }; + }, + + async listBlobReferences() { + const all = (await redis.hgetall>(k.refs)) ?? {}; + return Object.entries(all) + .map(([key, references]) => ({ key, references: Number(references) })) + .filter((r) => r.references > 0) + .sort((a, b) => compareUtf8Bytes(a.key, b.key)); + }, + + async forkFromCheckpoint(input) { + const { sourceThreadId, sourceCheckpointId, destinationThreadId, destinationCheckpointId } = + input; + const writer = { + threadId: input.writer.threadId, + ownerToken: input.writer.ownerToken, + fence: input.writer.fence, + }; + assertValidIdentifier(sourceThreadId, "Source thread id"); + assertValidIdentifier(sourceCheckpointId, "Source checkpoint id"); + assertValidIdentifier(destinationThreadId, "Destination thread id"); + assertValidIdentifier(destinationCheckpointId, "Destination checkpoint id"); + assertValidIdentifier(writer.threadId, "Writer thread id"); + if (!Number.isFinite(input.createdAt)) { + throw new SandboxCheckpointInvalidEntryError( + "Fork checkpoint createdAt must be a finite number", + ); + } + if (sourceThreadId === destinationThreadId) { + throw new SandboxCheckpointError( + "SANDBOX_SNAPSHOT_FORK_SOURCE_THREAD_MISMATCH", + "Source and destination threads must differ", + ); + } + const source = await load(sourceCheckpointId); + if (!source) { + throw new SandboxCheckpointError( + "SANDBOX_SNAPSHOT_FORK_SOURCE_NOT_FOUND", + "Source checkpoint was not found", + ); + } + if (source.threadId !== sourceThreadId) { + throw new SandboxCheckpointError( + "SANDBOX_SNAPSHOT_FORK_SOURCE_THREAD_MISMATCH", + "Source checkpoint belongs to another thread", + ); + } + if (writer.threadId !== destinationThreadId) { + throw new SandboxCheckpointWriterLostError( + "Checkpoint writer thread does not match destination thread", + ); + } + const checkpoint: SandboxCheckpoint = { + id: destinationCheckpointId, + threadId: destinationThreadId, + parentCheckpointId: null, + createdAt: input.createdAt, + reason: "fork-root", + files: source.files, + conversation: source.conversation, + artifacts: source.artifacts, + }; + validateCheckpoint(checkpoint); + const d = destinationThreadId; + const result = (await redis.eval( + FORK, + [ + k.writer(d), + k.doc(sourceCheckpointId), + k.doc(destinationCheckpointId), + k.head(d), + k.index(d), + k.refs, + k.transcript(d), + k.threadRuns(d), + k.threadInterrupts(d), + k.threadGenerationRuns(d), + k.threadArtifacts(d), + ], + [ + ...writerArgs(writer), + JSON.stringify(checkpoint), + JSON.stringify(checkpoint.conversation), + member(destinationCheckpointId), + String(checkpoint.createdAt), + ttl, + ...blobKeysOf(checkpoint), + ], + )) as string; + if (result === "WRITER_LOST") throw writerLost(d); + if (result === "SOURCE_NOT_FOUND") { + throw new SandboxCheckpointError( + "SANDBOX_SNAPSHOT_FORK_SOURCE_NOT_FOUND", + "Source checkpoint was not found", + ); + } + if (result === "NOT_EMPTY") { + throw new SandboxCheckpointError( + "SANDBOX_SNAPSHOT_FORK_DESTINATION_NOT_EMPTY", + "Destination thread is not empty", + ); + } + return { checkpoint: copy(checkpoint) }; + }, + }; +} diff --git a/packages/tanstack-ai/src/sandbox/index.ts b/packages/tanstack-ai/src/sandbox/index.ts new file mode 100644 index 00000000..c08214b4 --- /dev/null +++ b/packages/tanstack-ai/src/sandbox/index.ts @@ -0,0 +1,7 @@ +// `@upstash/agentkit-tanstack-ai/sandbox` — `@tanstack/ai-sandbox` stores on Upstash Redis: the +// sandbox instance registry and the checkpoint (snapshot) history with forks. Needs +// `@tanstack/ai-sandbox`. +export { upstashSandboxInstanceStore } from "./instance-store.js"; +export type { UpstashSandboxInstanceStoreConfig } from "./instance-store.js"; +export { upstashSandboxCheckpointStore } from "./checkpoint-store.js"; +export type { UpstashSandboxCheckpointStoreConfig } from "./checkpoint-store.js"; diff --git a/packages/tanstack-ai/src/sandbox/instance-store.ts b/packages/tanstack-ai/src/sandbox/instance-store.ts new file mode 100644 index 00000000..581b441d --- /dev/null +++ b/packages/tanstack-ai/src/sandbox/instance-store.ts @@ -0,0 +1,57 @@ +import type { Redis } from "@upstash/redis"; +import { Redis as RedisClient } from "@upstash/redis"; +import type { SandboxInstanceRecord, SandboxInstanceStore } from "@tanstack/ai-sandbox"; +import { addTelemetry } from "../telemetry.js"; + +export interface UpstashSandboxInstanceStoreConfig { + /** Upstash Redis client. Defaults to `Redis.fromEnv()`. */ + redis?: Redis; + /** Base key prefix. Defaults to `agentkit:tanstack`. */ + prefix?: string; + /** + * Expire a record this many seconds after its last `upsert`. Omit to keep records until `delete`. + * Set it to at least your provider's sandbox lifetime: a record that expires early makes the next + * run create a new sandbox instead of resuming the old one. + */ + ttlSeconds?: number; + /** Set `false` to opt out of the telemetry header. */ + enableTelemetry?: boolean; +} + +/** + * `SandboxInstanceStore` (from `@tanstack/ai-sandbox`) on Upstash Redis: which provider sandbox (and + * latest snapshot) belongs to a sandbox key, so any server can resume a thread's sandbox instead of + * only the one that created it. One JSON document per key. Checked against + * `runSandboxInstanceStoreConformance` from `@tanstack/ai-sandbox/testkit`. + * + * @example + * ```ts + * import { provideSandboxInstanceStore } from "@tanstack/ai-sandbox"; + * import { upstashSandboxInstanceStore } from "@upstash/agentkit-tanstack-ai/sandbox"; + * + * provideSandboxInstanceStore(upstashSandboxInstanceStore()); + * ``` + */ +export function upstashSandboxInstanceStore( + config: UpstashSandboxInstanceStoreConfig = {}, +): SandboxInstanceStore { + const redis = config.redis ?? RedisClient.fromEnv(); + addTelemetry(redis, config.enableTelemetry); + const p = config.prefix ?? "agentkit:tanstack"; + const key = (k: string) => `${p}:sandbox:instance:${encodeURIComponent(k)}`; + + return { + async get(k) { + return (await redis.get(key(k))) ?? null; + }, + async upsert(record) { + // JSON drops `undefined` optional fields, so they read back absent, like the reference store. + const value = JSON.stringify(record); + if (config.ttlSeconds) await redis.set(key(record.key), value, { ex: config.ttlSeconds }); + else await redis.set(key(record.key), value); + }, + async delete(k) { + await redis.del(key(k)); + }, + }; +} diff --git a/packages/tanstack-ai/src/sandbox/sandbox.test.ts b/packages/tanstack-ai/src/sandbox/sandbox.test.ts new file mode 100644 index 00000000..c40ed096 --- /dev/null +++ b/packages/tanstack-ai/src/sandbox/sandbox.test.ts @@ -0,0 +1,119 @@ +import { afterAll, describe, expect, it } from "vitest"; +import { + runDurableRunFieldsConformance, + runSandboxCheckpointForkConformance, + runSandboxCheckpointStoreConformance, + runSandboxInstanceStoreConformance, +} from "@tanstack/ai-sandbox/testkit"; +import { upstashPersistence } from "../persistence/persistence.js"; +import { upstashSandboxCheckpointStore } from "./checkpoint-store.js"; +import { upstashSandboxInstanceStore } from "./instance-store.js"; +import { cleanupKeys, hasRedisCreds, testRedis, uniquePrefix } from "../testing/test-support.js"; + +// TanStack AI's own sandbox store conformance suites, run against a real Upstash Redis. Every case +// gets a fresh keyspace: the suites use fixed thread ids and assume an empty store. +describe.skipIf(!hasRedisCreds)("sandbox stores (live Redis)", () => { + const redis = testRedis(); + const prefix = uniquePrefix("tsb"); + let n = 0; + const next = () => `${prefix}:${++n}`; + + afterAll(async () => { + await cleanupKeys(redis, prefix); + }); + + runSandboxInstanceStoreConformance("upstashSandboxInstanceStore", () => + upstashSandboxInstanceStore({ redis, prefix: next() }), + ); + + runSandboxCheckpointStoreConformance("upstashSandboxCheckpointStore", (options) => + upstashSandboxCheckpointStore({ redis, prefix: next(), ...options }), + ); + + // The fork writes the transcript into `upstashPersistence()`'s message store on the same prefix. + runSandboxCheckpointForkConformance("upstashSandboxCheckpointStore", () => { + const p = next(); + const checkpoints = upstashSandboxCheckpointStore({ redis, prefix: p }); + const fork = checkpoints.forkFromCheckpoint.bind(checkpoints); + return { + persistence: upstashPersistence({ redis, prefix: p }), + checkpoints: { + ...checkpoints, + // The suite creates the fork promise, then reads the destination over the network, and only + // then attaches `expect(...).rejects`. Against a real Redis the rejection lands in between, + // and vitest reports it as unhandled. A no-op handler marks it handled; the suite still + // receives (and asserts on) the same rejected promise. + forkFromCheckpoint: (input) => { + const promise = fork(input); + promise.catch(() => {}); + return promise; + }, + }, + }; + }); + + // Sandbox durability stores its driver fields on the run record: check `upstashPersistence()`'s + // run store keeps (and clears) them, since sandbox runs reuse it. + runDurableRunFieldsConformance( + "upstashPersistence runs", + () => upstashPersistence({ redis, prefix: next() }).stores.runs, + ); + + it("refuses a fork into a thread that already has persisted runs", async () => { + const p = next(); + const persistence = upstashPersistence({ redis, prefix: p }); + const checkpoints = upstashSandboxCheckpointStore({ redis, prefix: p }); + const source = await checkpoints.acquireWriter("src"); + await checkpoints.append({ + checkpoint: { + id: "root", + threadId: "src", + parentCheckpointId: null, + createdAt: 1, + reason: "named", + files: [], + conversation: [{ role: "user", content: "hi" }], + artifacts: [], + }, + expectedHeadId: null, + writer: source, + }); + await persistence.stores.runs.createOrResume({ + runId: "r1", + threadId: "dst", + status: "running", + startedAt: 1, + } as never); + const writer = await checkpoints.acquireWriter("dst"); + await expect( + checkpoints.forkFromCheckpoint({ + sourceThreadId: "src", + sourceCheckpointId: "root", + destinationThreadId: "dst", + destinationCheckpointId: "fork", + createdAt: 2, + writer, + }), + ).rejects.toMatchObject({ code: "SANDBOX_SNAPSHOT_FORK_DESTINATION_NOT_EMPTY" }); + expect(await checkpoints.getHead("dst")).toBeNull(); + }); + + it("keeps numeric-looking checkpoint ids as strings", async () => { + const checkpoints = upstashSandboxCheckpointStore({ redis, prefix: next() }); + const writer = await checkpoints.acquireWriter("t"); + const base = { + threadId: "t", + reason: "automatic" as const, + files: [], + conversation: [], + artifacts: [], + }; + await checkpoints.append({ + checkpoint: { ...base, id: "007", parentCheckpointId: null, createdAt: 1 }, + expectedHeadId: null, + writer, + }); + expect(await checkpoints.getHead("t")).toBe("007"); + expect((await checkpoints.list("t")).map((c) => c.id)).toEqual(["007"]); + }); +}); diff --git a/packages/tanstack-ai/src/sandbox/validate.ts b/packages/tanstack-ai/src/sandbox/validate.ts new file mode 100644 index 00000000..7959cbf1 --- /dev/null +++ b/packages/tanstack-ai/src/sandbox/validate.ts @@ -0,0 +1,172 @@ +import { + SandboxCheckpointInvalidEntryError, + SandboxCheckpointInvalidIdError, + type SandboxCheckpoint, +} from "@tanstack/ai-sandbox"; + +// Input checks shared by the checkpoint store. They mirror `@tanstack/ai-sandbox`'s own in-memory +// store (which does not export them) so both stores refuse the same inputs with the same codes. + +const utf8 = new TextEncoder(); + +/** Order by UTF-8 bytes, the order the contract specifies for checkpoint ids and blob keys. */ +export function compareUtf8Bytes(a: string, b: string): number { + const x = utf8.encode(a); + const y = utf8.encode(b); + const n = Math.min(x.length, y.length); + for (let i = 0; i < n; i++) if (x[i] !== y[i]) return x[i]! - y[i]!; + return x.length - y.length; +} + +function hasUnpairedSurrogate(value: string): boolean { + for (let i = 0; i < value.length; i++) { + const code = value.charCodeAt(i); + if (code >= 0xd800 && code <= 0xdbff) { + const next = value.charCodeAt(i + 1); + if (Number.isNaN(next) || next < 0xdc00 || next > 0xdfff) return true; + i++; + } else if (code >= 0xdc00 && code <= 0xdfff) return true; + } + return false; +} + +export function assertValidIdentifier(value: unknown, label: string): asserts value is string { + if (typeof value !== "string" || value.length === 0 || hasUnpairedSurrogate(value)) { + throw new SandboxCheckpointInvalidIdError( + `${label} must be a non-empty well-formed Unicode string`, + ); + } +} + +const hasOwn = (value: object, key: string) => Object.prototype.hasOwnProperty.call(value, key); +const FILE_KEY = /^sandbox-files\/sha256\/[0-9a-f]{64}$/; +const ARTIFACT_KEY = /^sandbox-artifacts\/sha256\/[0-9a-f]{64}$/; + +function validateEntries(checkpoint: SandboxCheckpoint): void { + if (!Array.isArray(checkpoint.files)) { + throw new SandboxCheckpointInvalidEntryError("Checkpoint files must be an array"); + } + const paths = new Set(); + const kinds = new Map(); + for (const entry of checkpoint.files as unknown[]) { + if (entry === null || typeof entry !== "object") { + throw new SandboxCheckpointInvalidEntryError("Checkpoint entry must be an object"); + } + const c = entry as Record; + const path = c.path; + if ( + typeof path !== "string" || + path.length === 0 || + path.includes("\0") || + path.startsWith("/") || + path.startsWith("\\") || + /^[A-Za-z]:([\\/]|$)/.test(path) || + path.includes("\\") || + path.split("/").some((part) => part.length === 0 || part === "." || part === "..") + ) { + throw new SandboxCheckpointInvalidEntryError( + "Checkpoint entry path must be a normalized workspace-relative path", + ); + } + if (paths.has(path)) { + throw new SandboxCheckpointInvalidEntryError( + `Checkpoint contains duplicate entry path '${path}'`, + ); + } + for (let s = path.indexOf("/"); s !== -1; s = path.indexOf("/", s + 1)) { + const ancestor = path.slice(0, s); + if (kinds.get(ancestor) === "file") { + throw new SandboxCheckpointInvalidEntryError( + `Checkpoint entry '${path}' is beneath file '${ancestor}'`, + ); + } + } + if (c.kind === "file" && [...kinds.keys()].some((other) => other.startsWith(`${path}/`))) { + throw new SandboxCheckpointInvalidEntryError( + `Checkpoint file '${path}' is an ancestor of another entry`, + ); + } + paths.add(path); + if (c.kind === "file") { + if ( + typeof c.blobKey !== "string" || + hasUnpairedSurrogate(c.blobKey) || + !FILE_KEY.test(c.blobKey) + ) { + throw new SandboxCheckpointInvalidEntryError( + "File entries require a valid content-addressed blobKey", + ); + } + if ( + !hasOwn(c, "size") || + typeof c.size !== "number" || + !Number.isSafeInteger(c.size) || + c.size < 0 + ) { + throw new SandboxCheckpointInvalidEntryError( + "File entry size must be a non-negative safe integer", + ); + } + } else if (c.kind === "dir") { + if (hasOwn(c, "blobKey") || hasOwn(c, "size")) { + throw new SandboxCheckpointInvalidEntryError( + "Directory entries cannot contain file fields", + ); + } + } else { + throw new SandboxCheckpointInvalidEntryError("Checkpoint entry kind must be file or dir"); + } + kinds.set(path, c.kind); + } +} + +function validateArtifacts(checkpoint: SandboxCheckpoint): void { + if (!Array.isArray(checkpoint.artifacts)) { + throw new SandboxCheckpointInvalidEntryError("Checkpoint artifacts must be an array"); + } + for (const artifact of checkpoint.artifacts as unknown[]) { + if (artifact === null || typeof artifact !== "object") { + throw new SandboxCheckpointInvalidEntryError("Checkpoint artifact must be an object"); + } + const a = artifact as Record; + const ok = + typeof a.artifactId === "string" && + a.artifactId.length > 0 && + !hasUnpairedSurrogate(a.artifactId) && + typeof a.name === "string" && + a.name.length > 0 && + typeof a.mimeType === "string" && + a.mimeType.length > 0 && + typeof a.blobKey === "string" && + !hasUnpairedSurrogate(a.blobKey) && + ARTIFACT_KEY.test(a.blobKey) && + typeof a.size === "number" && + Number.isSafeInteger(a.size) && + a.size >= 0 && + typeof a.createdAt === "number" && + Number.isFinite(a.createdAt); + if (!ok) throw new SandboxCheckpointInvalidEntryError("Checkpoint artifact has invalid fields"); + } +} + +/** Every field check the contract makes on a checkpoint, ids included. */ +export function validateCheckpoint(checkpoint: SandboxCheckpoint): void { + assertValidIdentifier(checkpoint.id, "Checkpoint id"); + assertValidIdentifier(checkpoint.threadId, "Checkpoint thread id"); + if (checkpoint.parentCheckpointId != null) { + assertValidIdentifier(checkpoint.parentCheckpointId, "Parent checkpoint id"); + } + if (typeof checkpoint.createdAt !== "number" || !Number.isFinite(checkpoint.createdAt)) { + throw new SandboxCheckpointInvalidEntryError("Checkpoint createdAt must be a finite number"); + } + validateEntries(checkpoint); + validateArtifacts(checkpoint); +} + +/** The distinct blob keys a checkpoint references (files and artifacts). */ +export function blobKeysOf(checkpoint: SandboxCheckpoint): string[] { + const keys = new Set(); + for (const entry of checkpoint.files) if (entry.kind === "file") keys.add(entry.blobKey); + for (const artifact of checkpoint.artifacts) keys.add(artifact.blobKey); + return [...keys]; +} diff --git a/packages/tanstack-ai/src/sandbox/wiring.test.ts b/packages/tanstack-ai/src/sandbox/wiring.test.ts new file mode 100644 index 00000000..a43bb8ef --- /dev/null +++ b/packages/tanstack-ai/src/sandbox/wiring.test.ts @@ -0,0 +1,25 @@ +import { describe, expect, it } from "vitest"; +import { createSandboxSnapshots } from "@tanstack/ai-sandbox"; +import { Redis } from "@upstash/redis"; +import type { BlobBucketLike } from "../persistence/blob-store.js"; +import { upstashLocks } from "../locks/locks.js"; +import { upstashPersistence } from "../persistence/persistence.js"; +import { upstashSandboxCheckpointStore } from "./checkpoint-store.js"; +import { upstashSandboxInstanceStore } from "./instance-store.js"; + +// The README's "everything on Upstash" snapshot wiring must keep type-checking against +// `createSandboxSnapshots`: persistence (messages + artifacts + blobs), checkpoints, instances, locks. +describe("sandbox snapshot wiring", () => { + it("accepts every Upstash store where createSandboxSnapshots expects one", () => { + const redis = new Redis({ url: "https://example.upstash.io", token: "t" }); + // Never called here: the bucket only has to have the shape `upstashPersistence` accepts. + const bucket = {} as BlobBucketLike; + const snapshots = createSandboxSnapshots({ + persistence: upstashPersistence({ redis, bucket }), + checkpoints: upstashSandboxCheckpointStore({ redis }), + instances: upstashSandboxInstanceStore({ redis }), + locks: upstashLocks({ redis }), + }); + expect(typeof snapshots.fork).toBe("function"); + }); +}); diff --git a/packages/tanstack-ai/tsup.config.ts b/packages/tanstack-ai/tsup.config.ts index b98eee6a..63bca606 100644 --- a/packages/tanstack-ai/tsup.config.ts +++ b/packages/tanstack-ai/tsup.config.ts @@ -6,6 +6,8 @@ export default defineConfig({ persistence: "src/persistence/index.ts", memory: "src/memory/index.ts", byok: "src/byok/index.ts", + sandbox: "src/sandbox/index.ts", + mcp: "src/mcp/index.ts", }, format: ["esm"], dts: true, diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 761d111a..b1cfe094 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -466,12 +466,18 @@ importers: '@tanstack/ai': specifier: 1.4.0 version: 1.4.0(@opentelemetry/api@1.9.1)(zod@4.4.3) + '@tanstack/ai-mcp': + specifier: 1.2.0 + version: 1.2.0(@opentelemetry/api@1.9.1)(zod@4.4.3) '@tanstack/ai-memory': specifier: 1.1.0 version: 1.1.0(@tanstack/ai@1.4.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: 1.0.1 version: 1.0.1(@tanstack/ai@1.4.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-sandbox': + specifier: 1.2.1 + version: 1.2.1(@tanstack/ai-persistence@1.0.1(@tanstack/ai@1.4.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@1.4.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(vitest@2.1.9(@types/node@20.19.43)(lightningcss@1.32.0))(zod@4.4.3) '@upstash/blob': specifier: 0.0.7 version: 0.0.7(react@19.2.6) @@ -2491,6 +2497,15 @@ packages: '@tanstack/ai-event-client@1.0.1': resolution: {integrity: sha512-IURnn4oRKJbtbAuC3GcXRcuKHi8DrPmwuACqaTq6ApSGlyJyHhq0NTksKEmzql2Z12+njFjsKSXG9c7UbSHTxg==} + '@tanstack/ai-mcp@1.2.0': + resolution: {integrity: sha512-EIJfHjaAuZJjJ5FL841jHKPrmaZCpRlBLfToqTBvtTjA9k4eoDpBDB/HVeN+/kUIOWayVgLON1xZI8SqwCGzWg==} + hasBin: true + peerDependencies: + '@tanstack/ai-harness': ^0.3.0 + peerDependenciesMeta: + '@tanstack/ai-harness': + optional: true + '@tanstack/ai-memory@1.1.0': resolution: {integrity: sha512-Z8u8ExkOkb7tKYjsI++eqnJanF5fxqNrGskKEsEyiLtbYVBPAPXKzTGr1JOEntm/AhSmlASzDevYScVIBAtGHQ==} peerDependencies: @@ -2540,6 +2555,38 @@ packages: react-dom: optional: true + '@tanstack/ai-sandbox@1.2.1': + resolution: {integrity: sha512-fhQx9b5gThjAQhkHeQisp8fKwslVmxPTdjYDaMxteIcmSLRikheeYTHPD9RuqHx7EZ+aJGPpIm5or6dEDM2bdA==} + peerDependencies: + '@ngrok/ngrok': ^1.0.0 + '@tanstack/ai': ^1.3.0 + '@tanstack/ai-harness': ^0.3.0 + '@tanstack/ai-persistence': ^1.0.1 + vitest: ^4.1.10 + peerDependenciesMeta: + '@ngrok/ngrok': + optional: true + '@tanstack/ai-harness': + optional: true + '@tanstack/ai-persistence': + optional: true + vitest: + optional: true + + '@tanstack/ai-skills@1.0.2': + resolution: {integrity: sha512-snOdkl6rLUvfykR6Z7kK0bqS0mp0iXIZ646txNzB2O5vkK9yOh2FzyYAofehhjdUNNC+iU+3fyhMKFuyjPKNyA==} + engines: {node: '>=18'} + peerDependencies: + '@tanstack/ai': ^1.3.0 + '@tanstack/ai-harness': ^0.3.0 + vitest: ^4.1.10 + zod: ^3.0.0 || ^4.0.0 + peerDependenciesMeta: + '@tanstack/ai-harness': + optional: true + vitest: + optional: true + '@tanstack/ai-utils@1.1.0': resolution: {integrity: sha512-li47v/794azVLourgteS/+jjH/ZF+cIp/3RMVA1VybtKAFiMy3ZGxREU10OzfUghMzYLdfpK2l6ahvlWxNVHug==} @@ -7088,6 +7135,16 @@ snapshots: dependencies: '@tanstack/devtools-event-client': 0.4.4 + '@tanstack/ai-mcp@1.2.0(@opentelemetry/api@1.9.1)(zod@4.4.3)': + dependencies: + '@modelcontextprotocol/client': 2.3.1 + '@modelcontextprotocol/server': 2.3.1 + '@tanstack/ai': 1.4.0(@opentelemetry/api@1.9.1)(zod@4.4.3) + jose: 6.2.12 + transitivePeerDependencies: + - '@opentelemetry/api' + - zod + '@tanstack/ai-memory@1.1.0(@tanstack/ai@1.4.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': 1.4.0(@opentelemetry/api@1.9.1)(zod@4.4.3) @@ -7131,6 +7188,25 @@ snapshots: - octane - zod + '@tanstack/ai-sandbox@1.2.1(@tanstack/ai-persistence@1.0.1(@tanstack/ai@1.4.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@1.4.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(vitest@2.1.9(@types/node@20.19.43)(lightningcss@1.32.0))(zod@4.4.3)': + dependencies: + '@modelcontextprotocol/server': 2.3.1 + '@tanstack/ai': 1.4.0(@opentelemetry/api@1.9.1)(zod@4.4.3) + '@tanstack/ai-skills': 1.0.2(@tanstack/ai@1.4.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(vitest@2.1.9(@types/node@20.19.43)(lightningcss@1.32.0))(zod@4.4.3) + '@tanstack/ai-utils': 1.1.0 + optionalDependencies: + '@tanstack/ai-persistence': 1.0.1(@tanstack/ai@1.4.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(vitest@2.1.9(@types/node@20.19.43)(lightningcss@1.32.0)) + vitest: 2.1.9(@types/node@20.19.43)(lightningcss@1.32.0) + transitivePeerDependencies: + - zod + + '@tanstack/ai-skills@1.0.2(@tanstack/ai@1.4.0(@opentelemetry/api@1.9.1)(zod@4.4.3))(vitest@2.1.9(@types/node@20.19.43)(lightningcss@1.32.0))(zod@4.4.3)': + dependencies: + '@tanstack/ai': 1.4.0(@opentelemetry/api@1.9.1)(zod@4.4.3) + zod: 4.4.3 + optionalDependencies: + vitest: 2.1.9(@types/node@20.19.43)(lightningcss@1.32.0) + '@tanstack/ai-utils@1.1.0': {} '@tanstack/ai@1.4.0(@opentelemetry/api@1.9.1)(zod@4.4.3)': From 32e6e7b2ab1b0aa065874fb057848bc5d61df108 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 11 Nov 2026 11:31:44 +0000 Subject: [PATCH 2/7] feat(tanstack-ai): upstashMcpSessionStore for MCP Apps McpServerDescriptor carries functions and objects (needsApproval, toolFilter, toolName, a transport's fetch and authProvider) that cannot be stored in Redis. So set() stores the JSON part and records which code-only fields were dropped. get() takes them from a user-supplied resolve() callback. get() fails closed: if a recorded field is still missing, it returns null, so createMcpAppCallHandler falls back to its registry or refuses the call. A widget can never reach a server that lost its approval gate. Transport headers and stdio env can be encrypted at rest with AES-256-GCM, bound to the prefix, thread and server. The helpers move to src/secret-box.ts, which the BYOK credential store now shares. The TTL slides like the in-memory store's (1800 s by default). An end-to-end test drives createMcpAppCallHandler against a real MCP server over HTTP. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01CMK4etJMDDBFkjEyWDVCdy --- .changeset/tanstack-sandbox-mcp-stores.md | 4 + CLAUDE.md | 15 +- packages/tanstack-ai/README.md | 44 +++- packages/tanstack-ai/src/byok/credentials.ts | 38 +-- packages/tanstack-ai/src/mcp/index.ts | 8 +- .../tanstack-ai/src/mcp/session-store.test.ts | 221 ++++++++++++++++ packages/tanstack-ai/src/mcp/session-store.ts | 244 ++++++++++++++++++ packages/tanstack-ai/src/secret-box.ts | 79 ++++++ 8 files changed, 609 insertions(+), 44 deletions(-) create mode 100644 packages/tanstack-ai/src/mcp/session-store.test.ts create mode 100644 packages/tanstack-ai/src/mcp/session-store.ts create mode 100644 packages/tanstack-ai/src/secret-box.ts diff --git a/.changeset/tanstack-sandbox-mcp-stores.md b/.changeset/tanstack-sandbox-mcp-stores.md index d2303680..85853997 100644 --- a/.changeset/tanstack-sandbox-mcp-stores.md +++ b/.changeset/tanstack-sandbox-mcp-stores.md @@ -11,4 +11,8 @@ Add stores for TanStack AI's sandbox and MCP packages, all on Upstash Redis. - `@upstash/agentkit-tanstack-ai/mcp` (needs `@tanstack/ai-mcp` >= 1.0.0): - `upstashMcpTaskStore()` is the `taskStore` for `createMCPServer`. - `upstashMcpEventStore()` stores MCP Events webhook subscriptions, shared across instances. + - `upstashMcpSessionStore()` is the `McpSessionStore` for MCP Apps. + - The JSON part of each server descriptor is stored in Redis. Transport headers and env can be encrypted. + - Functions such as `needsApproval` come back through your `resolve` callback. + - It fails closed: it returns `null` rather than a descriptor that lost its approval gate. - `upstashPersistence({ bucket })` is now typed with a required `blobs` store, so you can pass it straight to `createSandboxSnapshots()`. diff --git a/CLAUDE.md b/CLAUDE.md index 13fb0e29..2f3a467c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1711,9 +1711,18 @@ find node_modules -path "*/zod/package.json" | while read f; do echo "$f $(node need_tools/replay loop), QStash-backed background runs writing to `upstashStream`, a QStash schedule for `reapDetachedRuns`, `@tanstack/ai-code-mode`'s `AgentStore` and `@tanstack/ai-code-mode-snippets`' `SnippetStorage` (both small keyed stores; `ai-code-mode` peers the pre-1.0 `@tanstack/ai-harness`). - **Deliberately not built:** `@tanstack/ai-mcp`'s `McpSessionStore` (MCP Apps) — its - `McpServerDescriptor` carries functions (`toolFilter`, `needsApproval`, `toolName`) and transport - objects, so it cannot round-trip through Redis. + **`McpSessionStore` (MCP Apps) is `upstashMcpSessionStore`:** + - **What can't be stored:** its `McpServerDescriptor` carries functions (`toolFilter`, `needsApproval`, + `toolName`) and transport objects (`fetch`, `authProvider`). + - **`set`:** stores only JSON-safe fields, checked by prototype and recursion in `isJsonSafe`, and + records the dropped paths in `m`. Transport `headers` / `env` are sealed with AES-GCM (`src/secret-box.ts`, + shared with BYOK; the AAD is prefix + thread + server). + - **`get`:** calls the user's `resolve({ threadId, serverId, descriptor, missing })`, then FAILS CLOSED, + returning `null` if any recorded path is still missing. Upstream's `createMcpAppCallHandler` then + falls back to its registry (`descriptor.x ?? fromRegistry?.x`, per field) or refuses. + - **Why fail closed:** a per-thread server that lost `needsApproval` would let a widget run a gated tool. + - **Test:** `src/mcp/session-store.test.ts` drives the real handler against a real `createMCPServer` + over node:http that requires the stored Authorization header. - **ai 7.0.169 + @tanstack/ai 1.4.0 + @tanstack/ai-openai 1.1.1 (2026-11-11):** devDependency/example pins only (no published-content change, no changeset); `@ai-sdk/*` ranges moved with `pnpm -r update "@ai-sdk/*"` (one `ai@7.0.169` in the lockfile) and `pnpm dedupe` dropped a stale second `@tanstack/ai@1.1.0` that `@tanstack/ai-client@1.1.0` (`^1.1.0`) kept in the lockfile. eve stays on 0.91.0: no stable release since, npm `latest` is `1.0.0-rc.9` (not adopted); published `@upstash/agentkit-eve-extension@0.19.0` + `@upstash/agentkit-eve@0.19.0` install on it, `eve build` clean and the mocked `eve eval` passes 4/4. A `skipLibCheck: false` consumer of `ai` errors inside `@ai-sdk/provider` (`TS7016 … module 'json-schema'`, no `@types/json-schema` dependency) on 7.0.163 and 7.0.169 alike: upstream, fixed by `npm i -D @types/json-schema`. The free throwaway Redis caps search indexes at 1: run index-using test files one at a time with `FLUSHDB` in between (`FLUSHDB` drops the indexes too; `SEARCH.DROPINDEX` is not available). - **ai 7.0.163 (2026-11-05):** pins bumped from 7.0.159 (no code change); `examples/ai-sdk-demo`'s `@ai-sdk/react` raised to `^4.0.167` because 4.0.160 depends on an *exact* `ai@7.0.156`, which split the `UIMessage` type from the demo's `ai@7.0.163` and failed `next build` (`SteeringMessageUIPart`). Re-checked 2026-11-05: published `@upstash/agentkit-eve-extension@0.19.0` / `@upstash/agentkit-eve@0.19.0` build clean on `eve@1.0.0-rc.5` (the 0.91.0 build stays); TanStack AI 1.x (ai 1.1.0, ai-persistence 1.0.1, ai-memory 1.1.0) still not adopted (same `threadId`/`blobKey` test-file errors). - **TanStack AI 1.x applied (2026-11-05):** `@tanstack/ai` 1.1.0, `ai-persistence` 1.0.1, `ai-memory` 1.1.0 (devDeps, exact) and the tanstack-ai-demo's `ai-client` 1.1.0 / `ai-react` 1.0.2 / `ai-openai` 1.1.0. `src` needed no change; only tests did: `chat()` calls pass `threadId` (`middleware.test.ts`, `search-tools.test.ts`) and `ArtifactRecord` fixtures carry `blobKey` (`persistence.test.ts`). The peer floors (`>=0.65.0` / `>=0.8.0` / `>=0.2.8`) are untouched: the source is unchanged and still works on 0.x, and published `@upstash/agentkit-tanstack-ai@0.2.0` imports on 1.x, so no changeset. Upstream `@tanstack/ai@1.1.0`'s `spawn.d.ts` still has the broken `import('../../..')` (TS2307 with `skipLibCheck: false`). The 1.0 contracts `SessionIndexStore` and `CredentialStore` are implemented in the follow-up below. diff --git a/packages/tanstack-ai/README.md b/packages/tanstack-ai/README.md index 2eac591a..e68410d4 100644 --- a/packages/tanstack-ai/README.md +++ b/packages/tanstack-ai/README.md @@ -13,7 +13,7 @@ instances, reloads and devices. | `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 | | `createSandboxSnapshots()` / `withSandbox()` stores | `upstashSandboxCheckpointStore()`, `upstashSandboxInstanceStore()` | Portable sandbox snapshots with history and forks, and any server can resume a thread's sandbox | -| `createMCPServer({ taskStore })`, MCP Events `{ store }` | `upstashMcpTaskStore()`, `upstashMcpEventStore()` | Long-running MCP tasks and webhook subscriptions that survive restarts and work across instances | +| `createMCPServer({ taskStore })`, MCP Events `{ store }`, `createMcpAppCallHandler({ store })` | `upstashMcpTaskStore()`, `upstashMcpEventStore()`, `upstashMcpSessionStore()` | Long-running MCP tasks, webhook subscriptions and MCP Apps sessions that survive restarts and work across instances | | 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 | @@ -295,8 +295,46 @@ const mcp = await createMCPClient({ deadlines are shared, so a webhook can land on any server. `update` is a compare-and-set, so concurrent cursor updates never lose a field. The subscription's webhook `secret` is stored as is. -Needs `@tanstack/ai-mcp` >= 1.0.0. Not provided: `McpSessionStore` (MCP Apps). Its server descriptors -carry functions (`toolFilter`, `needsApproval`, ...) that cannot be stored in Redis. +Needs `@tanstack/ai-mcp` >= 1.0.0. + +### MCP Apps sessions + +```ts +import { createMcpAppCallHandler } from "@tanstack/ai-mcp/apps"; +import { upstashMcpSessionStore } from "@upstash/agentkit-tanstack-ai/mcp"; + +const store = upstashMcpSessionStore({ + encryptionKey: process.env.MCP_SESSION_KEY, // encrypts transport headers / stdio env at rest + // The parts of a descriptor that cannot live in Redis come from your code, per server: + resolve: ({ threadId, serverId }) => ({ + needsApproval: (tool) => tool.annotations?.readOnlyHint !== true, + authProvider: oauthProviderFor(threadId, serverId), + }), +}); + +await store.set(threadId, servers); // in the chat route +const handleCall = createMcpAppCallHandler({ clients, store }); // widget calls, on any instance +``` + +`upstashMcpSessionStore()` is the `McpSessionStore` for MCP Apps. It records which MCP servers a +thread's widgets may call, so `createMcpAppCallHandler` can reconnect on any instance. + +A server descriptor carries a policy for widget calls (`toolFilter`, `needsApproval`, `toolName`) and +often a transport `fetch` or OAuth `authProvider`. None of those survive JSON, so `set` splits each +descriptor: + +- **The JSON part goes to Redis:** the transport (with `headers` / `env` encrypted when you pass + `encryptionKey`), `prefix`, `requestOptions`, and an array `toolFilter`. One hash per thread, one field + per server. +- **Everything else is dropped and recorded by name.** `resolve` supplies it again on `get`, from your + code. It receives `{ threadId, serverId, descriptor, missing }`. +- **It fails closed.** If `resolve` doesn't return a recorded field (say `needsApproval`), `get` returns + `null`. The handler then uses its own registry entry for that server, or refuses the call. A widget + never reaches a server whose approval gate was lost on the way through Redis. + +Servers you also pass to `createMcpAppCallHandler({ clients })` get `toolFilter` / `needsApproval` / +`toolName` / `requestOptions` from that registry, so `resolve` matters most for servers that exist only +per thread. The TTL slides like `inMemoryMcpSessionStore` (`ttlSeconds`, default 1800). Are you building an MCP _server_ with the official `@modelcontextprotocol/server` SDK instead of TanStack? [`@upstash/mcp-toolkit`](https://github.com/upstash/agentkit/tree/main/packages/mcp-toolkit) gives you diff --git a/packages/tanstack-ai/src/byok/credentials.ts b/packages/tanstack-ai/src/byok/credentials.ts index 369e06df..04316d8c 100644 --- a/packages/tanstack-ai/src/byok/credentials.ts +++ b/packages/tanstack-ai/src/byok/credentials.ts @@ -3,6 +3,7 @@ import { Redis as RedisClient } from "@upstash/redis"; import type { Credential, CredentialScope, CredentialStore } from "@tanstack/ai/byok"; import { addTelemetry } from "../telemetry.js"; import { KEY_LOCKING, assertId } from "../persistence/records.js"; +import { b64, toKeyBytes } from "../secret-box.js"; export interface UpstashCredentialsConfig { /** Upstash Redis client. Defaults to `Redis.fromEnv()`. */ @@ -65,43 +66,6 @@ interface Stored { ct?: string; } -const b64 = { - // Chunked: spreading a large array into String.fromCharCode overflows the call stack. - encode(bytes: Uint8Array): string { - let binary = ""; - for (let i = 0; i < bytes.length; i += 0x8000) { - binary += String.fromCharCode(...bytes.subarray(i, i + 0x8000)); - } - return globalThis.btoa(binary); - }, - decode: (text: string) => Uint8Array.from(globalThis.atob(text), (c) => c.charCodeAt(0)), -}; - -const KEY_HELP = - "@upstash/agentkit-tanstack-ai: `encryptionKey` must be 32 bytes: a Uint8Array, or a base64, base64url or hex string (e.g. `openssl rand -base64 32`)."; - -function toKeyBytes(key: string | Uint8Array): Uint8Array { - let bytes: Uint8Array; - if (typeof key !== "string") bytes = key; - else { - const text = key.trim(); - try { - bytes = /^[0-9a-f]{64}$/i.test(text) - ? Uint8Array.from(text.match(/../g)!, (h) => parseInt(h, 16)) - : b64.decode( - text - .replace(/-/g, "+") - .replace(/_/g, "/") - .padEnd(Math.ceil(text.length / 4) * 4, "="), - ); - } catch { - throw new Error(KEY_HELP); - } - } - if (bytes.length !== 32) throw new Error(KEY_HELP); - return bytes; -} - /** A segment no encoded id can equal: `encodeURIComponent` never emits a bare "%n". */ const NONE = "%none"; const seg = (v: string | undefined) => (v === undefined ? NONE : encodeURIComponent(v)); diff --git a/packages/tanstack-ai/src/mcp/index.ts b/packages/tanstack-ai/src/mcp/index.ts index 4bfe3ede..c764433f 100644 --- a/packages/tanstack-ai/src/mcp/index.ts +++ b/packages/tanstack-ai/src/mcp/index.ts @@ -1,6 +1,12 @@ // `@upstash/agentkit-tanstack-ai/mcp` — `@tanstack/ai-mcp` stores on Upstash Redis: MCP server task -// records and MCP Events webhook subscriptions. Needs `@tanstack/ai-mcp`. +// records, MCP Events webhook subscriptions, and MCP Apps sessions. Needs `@tanstack/ai-mcp`. export { upstashMcpTaskStore } from "./task-store.js"; export type { UpstashMcpTaskStoreConfig } from "./task-store.js"; export { upstashMcpEventStore } from "./event-store.js"; export type { UpstashMcpEventStoreConfig } from "./event-store.js"; +export { upstashMcpSessionStore } from "./session-store.js"; +export type { + McpSessionOverrides, + McpSessionResolveContext, + UpstashMcpSessionStoreConfig, +} from "./session-store.js"; diff --git a/packages/tanstack-ai/src/mcp/session-store.test.ts b/packages/tanstack-ai/src/mcp/session-store.test.ts new file mode 100644 index 00000000..fa6734d1 --- /dev/null +++ b/packages/tanstack-ai/src/mcp/session-store.test.ts @@ -0,0 +1,221 @@ +import { createServer, type Server } from "node:http"; +import type { AddressInfo } from "node:net"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { toolDefinition } from "@tanstack/ai"; +import { createMcpAppCallHandler } from "@tanstack/ai-mcp/apps"; +import type { McpServerDescriptor } from "@tanstack/ai-mcp/apps"; +import { createMCPServer } from "@tanstack/ai-mcp/server"; +import { z } from "zod"; +import { upstashMcpSessionStore } from "./session-store.js"; +import { cleanupKeys, hasRedisCreds, testRedis, uniquePrefix } from "../testing/test-support.js"; + +const KEY = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"; +const TOKEN = "Bearer s3cr3t-token"; +const isDelete = (tool: { name: string }) => tool.name === "delete_doc"; + +/** A real MCP server over HTTP that only answers requests carrying the stored `Authorization`. */ +async function startMcpServer(): Promise<{ url: string; close: () => Promise }> { + const read = toolDefinition({ + name: "read_doc", + description: "Read a doc", + inputSchema: z.object({ id: z.string() }), + }).server(async ({ id }) => `doc ${id}`); + const remove = toolDefinition({ + name: "delete_doc", + description: "Delete a doc", + inputSchema: z.object({ id: z.string() }), + }).server(async ({ id }) => `deleted ${id}`); + const mcp = createMCPServer({ + name: "docs", + version: "1.0.0", + tools: [read, remove], + sessions: "stateless", + }); + + const server: Server = createServer(async (req, res) => { + if (req.headers.authorization !== TOKEN) { + res.writeHead(401).end(); + return; + } + const chunks: Buffer[] = []; + for await (const chunk of req) chunks.push(chunk as Buffer); + const body = chunks.length > 0 ? Buffer.concat(chunks) : undefined; + const response = await mcp.fetch( + new Request(`http://localhost${req.url}`, { + method: req.method, + headers: req.headers as Record, + ...(body && req.method !== "GET" && req.method !== "HEAD" ? { body } : {}), + }), + ); + res.writeHead(response.status, Object.fromEntries(response.headers)); + if (response.body) for await (const chunk of response.body) res.write(chunk); + res.end(); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const { port } = server.address() as AddressInfo; + return { + url: `http://127.0.0.1:${port}/mcp`, + close: () => new Promise((resolve) => server.close(() => resolve())), + }; +} + +describe.skipIf(!hasRedisCreds)("upstashMcpSessionStore (live Redis)", () => { + const redis = testRedis(); + const prefix = uniquePrefix("tmcps"); + let n = 0; + const next = () => `${prefix}:${++n}`; + let mcp: Awaited>; + + beforeAll(async () => { + mcp = await startMcpServer(); + }); + afterAll(async () => { + await mcp.close(); + await cleanupKeys(redis, prefix); + }); + + const docs = (): McpServerDescriptor => ({ + transport: { type: "http", url: mcp.url, headers: { authorization: TOKEN } }, + prefix: "docs", + toolFilter: ["read_doc", "delete_doc"], + needsApproval: isDelete, + requestOptions: { timeout: 10_000 }, + }); + + it("round-trips the JSON part and returns null for unknown threads and servers", async () => { + const store = upstashMcpSessionStore({ redis, prefix: next() }); + expect(await store.get("t1", "docs")).toBeNull(); + await store.set("t1", { + docs: { transport: { type: "stdio", command: "node", args: ["server.js"] }, prefix: "docs" }, + }); + expect(await store.get("t1", "docs")).toEqual({ + transport: { type: "stdio", command: "node", args: ["server.js"] }, + prefix: "docs", + }); + expect(await store.get("t1", "other")).toBeNull(); + expect(await store.get("t2", "docs")).toBeNull(); + }); + + it("resolves an omitted serverId only when the thread has exactly one server", async () => { + const store = upstashMcpSessionStore({ redis, prefix: next() }); + const stdio = (command: string): McpServerDescriptor => ({ + transport: { type: "stdio", command }, + }); + await store.set("t", { a: stdio("a") }); + expect(await store.get("t", undefined)).toMatchObject({ transport: { command: "a" } }); + await store.set("t", { a: stdio("a"), b: stdio("b") }); + expect(await store.get("t", undefined)).toBeNull(); + await store.set("t", { b: stdio("b") }); + expect(await store.get("t", "a")).toBeNull(); // set replaces the thread's whole map + }); + + it("fails closed when a code-only field is not supplied again, and restores it via resolve", async () => { + const p = next(); + await upstashMcpSessionStore({ redis, prefix: p }).set("t", { docs: docs() }); + expect(await upstashMcpSessionStore({ redis, prefix: p }).get("t", "docs")).toBeNull(); + + const seen: string[][] = []; + const store = upstashMcpSessionStore({ + redis, + prefix: p, + resolve: ({ missing }) => { + seen.push([...missing]); + return { needsApproval: isDelete }; + }, + }); + const descriptor = await store.get("t", "docs"); + expect(seen).toEqual([["needsApproval"]]); + expect(descriptor?.needsApproval).toBe(isDelete); + // JSON-safe fields, the array toolFilter included, came back from Redis as they were. + expect(descriptor).toMatchObject({ + prefix: "docs", + toolFilter: ["read_doc", "delete_doc"], + requestOptions: { timeout: 10_000 }, + transport: { type: "http", url: mcp.url, headers: { authorization: TOKEN } }, + }); + }); + + it("encrypts headers at rest, bound to the thread and server", async () => { + const p = next(); + const store = upstashMcpSessionStore({ + redis, + prefix: p, + encryptionKey: KEY, + resolve: () => ({ needsApproval: isDelete }), + }); + await store.set("t", { docs: docs() }); + const raw = JSON.stringify(await redis.hgetall(`${p}:mcp:session:t`)); + expect(raw).not.toContain("s3cr3t-token"); + expect(raw).toContain(mcp.url); + expect((await store.get("t", "docs"))?.transport).toMatchObject({ + headers: { authorization: TOKEN }, + }); + + // Copying the row to another thread does not decrypt: the ciphertext is bound to where it was saved. + await redis.hset(`${p}:mcp:session:u`, { + docs: (await redis.hget(`${p}:mcp:session:t`, "docs"))!, + }); + await expect(store.get("u", "docs")).rejects.toThrow(); + // And a store without the key refuses to hand out a descriptor with its headers missing. + await expect(upstashMcpSessionStore({ redis, prefix: p }).get("t", "docs")).rejects.toThrow( + /encryptionKey/, + ); + }); + + it("slides the TTL on get", async () => { + const p = next(); + const store = upstashMcpSessionStore({ redis, prefix: p, ttlSeconds: 100 }); + await store.set("t", { a: { transport: { type: "stdio", command: "a" } } }); + await redis.expire(`${p}:mcp:session:t`, 5); + await store.get("t", "a"); + expect(await redis.ttl(`${p}:mcp:session:t`)).toBeGreaterThan(50); + }); + + describe("with createMcpAppCallHandler against a real MCP server", () => { + it("calls allowed tools and keeps the approval gate a widget must not bypass", async () => { + const store = upstashMcpSessionStore({ + redis, + prefix: next(), + encryptionKey: KEY, + resolve: () => ({ needsApproval: isDelete }), + }); + await store.set("thread-1", { docs: docs() }); + // No static clients: the thread's servers exist only in the store. + const handle = createMcpAppCallHandler({ clients: [], store }); + + const read = await handle({ + threadId: "thread-1", + serverId: "docs", + toolName: "read_doc", + args: { id: "7" }, + }); + expect(read.ok).toBe(true); + expect(JSON.stringify(read)).toContain("doc 7"); + + const remove = await handle({ + threadId: "thread-1", + serverId: "docs", + toolName: "delete_doc", + args: { id: "7" }, + }); + expect(remove).toEqual({ ok: false, error: "Tool needs approval: delete_doc" }); + }); + + it("refuses the call instead of dropping the gate when resolve does not supply it", async () => { + const p = next(); + await upstashMcpSessionStore({ redis, prefix: p }).set("thread-1", { docs: docs() }); + const handle = createMcpAppCallHandler({ + clients: [], + store: upstashMcpSessionStore({ redis, prefix: p }), + }); + const remove = await handle({ + threadId: "thread-1", + serverId: "docs", + toolName: "delete_doc", + args: { id: "7" }, + }); + expect(remove.ok).toBe(false); + expect(JSON.stringify(remove)).not.toContain("deleted"); + }); + }); +}); diff --git a/packages/tanstack-ai/src/mcp/session-store.ts b/packages/tanstack-ai/src/mcp/session-store.ts new file mode 100644 index 00000000..142421ca --- /dev/null +++ b/packages/tanstack-ai/src/mcp/session-store.ts @@ -0,0 +1,244 @@ +import type { Redis } from "@upstash/redis"; +import { Redis as RedisClient } from "@upstash/redis"; +import type { McpServerDescriptor, McpSessionStore } from "@tanstack/ai-mcp/apps"; +import { addTelemetry } from "../telemetry.js"; +import { secretBox, type Sealed } from "../secret-box.js"; + +/** The parts of a descriptor that live in code, not in Redis: `resolve` supplies them per server. */ +export interface McpSessionOverrides { + toolFilter?: McpServerDescriptor["toolFilter"]; + needsApproval?: McpServerDescriptor["needsApproval"]; + toolName?: McpServerDescriptor["toolName"]; + clientOptions?: McpServerDescriptor["clientOptions"]; + requestOptions?: McpServerDescriptor["requestOptions"]; + /** Merged into an `http` / `sse` transport. */ + fetch?: typeof fetch; + /** Merged into an `http` / `sse` transport (an OAuth client provider object). */ + authProvider?: unknown; +} + +export interface McpSessionResolveContext { + threadId: string; + /** The server's key in the map passed to `set` (its `prefix` for `createMcpAppCallHandler`). */ + serverId: string; + /** The stored descriptor: JSON fields and decrypted headers/env, without the code-only fields. */ + descriptor: McpServerDescriptor; + /** Code-only fields the descriptor had when it was saved, e.g. `["needsApproval", "transport.fetch"]`. */ + missing: readonly string[]; +} + +export interface UpstashMcpSessionStoreConfig { + /** Upstash Redis client. Defaults to `Redis.fromEnv()`. */ + redis?: Redis; + /** Base key prefix. Defaults to `agentkit:tanstack`. */ + prefix?: string; + /** + * Supply the fields Redis cannot hold (functions and objects such as `needsApproval`, + * `toolFilter` as a function, `toolName`, a transport's `fetch` / `authProvider`) for a stored + * server. Called on every `get`. What you return is merged over the stored descriptor. + * + * A descriptor saved with such a field comes back only if `resolve` supplies it again: otherwise + * `get` returns `null` (fail closed), so `createMcpAppCallHandler` falls back to the server's + * registry entry, or refuses the call. A widget never reaches a server whose approval gate or tool + * filter was dropped on the way through Redis. + */ + resolve?: ( + context: McpSessionResolveContext, + ) => McpSessionOverrides | undefined | Promise; + /** + * Encrypt the transport's `headers` (often `Authorization`) and a stdio transport's `env` at rest + * with AES-256-GCM, bound to the prefix, thread and server. A 32-byte key as a `Uint8Array`, or a + * base64 / base64url / hex string (`openssl rand -base64 32`). Strongly recommended when headers + * carry credentials: without it they are stored as plain JSON. + */ + encryptionKey?: string | Uint8Array; + /** + * Forget a thread's servers this many seconds after its last `set` or `get` (sliding, like + * `inMemoryMcpSessionStore`). @default 1800 + */ + ttlSeconds?: number; + /** Set `false` to opt out of the telemetry header. */ + enableTelemetry?: boolean; +} + +interface StoredServer { + /** The JSON-safe part of the descriptor (headers/env removed when encrypted). */ + d: McpServerDescriptor; + /** Sealed `{ headers?, env? }` of the transport, when an `encryptionKey` is set. */ + s?: Sealed; + /** Paths of the code-only fields the descriptor had, which `resolve` must supply again. */ + m: string[]; +} + +const DESCRIPTOR_FIELDS = [ + "prefix", + "clientOptions", + "toolFilter", + "needsApproval", + "toolName", + "requestOptions", +] as const; + +/** A value that round-trips through JSON unchanged: plain objects, arrays, strings, finite numbers, booleans, null. */ +function isJsonSafe(value: unknown): boolean { + if (value === null || typeof value === "string" || typeof value === "boolean") return true; + if (typeof value === "number") return Number.isFinite(value); + if (Array.isArray(value)) return value.every(isJsonSafe); + if (typeof value === "object") { + const proto = Object.getPrototypeOf(value); + if (proto !== Object.prototype && proto !== null) return false; + return Object.values(value as object).every((v) => v === undefined || isJsonSafe(v)); + } + return false; +} + +/** + * `McpSessionStore` (MCP Apps, `@tanstack/ai-mcp/apps`) on Upstash Redis: which MCP servers a thread + * may call from its app widgets, so `createMcpAppCallHandler` can reconnect on any instance. + * + * A descriptor is split on `set`: its JSON part goes to Redis (one hash per thread, one field per + * server), and anything that is not JSON (`needsApproval`, a function `toolFilter`, `toolName`, a + * transport's `fetch` / `authProvider`) is dropped and recorded by name. `resolve` puts those back + * on `get`, from your code, per server. If it does not, `get` returns `null` rather than a descriptor + * missing its approval gate or tool filter. Transport headers and stdio env can be encrypted at rest. + * + * For servers you configure in code and pass to `createMcpAppCallHandler({ clients })`, the handler + * already takes `toolFilter` / `needsApproval` / `toolName` / `requestOptions` from that registry when + * the stored descriptor has none; `resolve` matters for servers that exist only per thread. + * + * @example + * ```ts + * import { createMcpAppCallHandler } from "@tanstack/ai-mcp/apps"; + * import { upstashMcpSessionStore } from "@upstash/agentkit-tanstack-ai/mcp"; + * + * const store = upstashMcpSessionStore({ + * encryptionKey: process.env.MCP_SESSION_KEY, + * resolve: ({ serverId }) => ({ + * needsApproval: (tool) => tool.annotations?.readOnlyHint !== true, + * }), + * }); + * const handleCall = createMcpAppCallHandler({ clients, store }); + * ``` + */ +export function upstashMcpSessionStore(config: UpstashMcpSessionStoreConfig = {}): McpSessionStore { + const redis = config.redis ?? RedisClient.fromEnv(); + addTelemetry(redis, config.enableTelemetry); + const p = config.prefix ?? "agentkit:tanstack"; + const ttl = config.ttlSeconds ?? 1800; + const box = config.encryptionKey === undefined ? undefined : secretBox(config.encryptionKey); + const key = (threadId: string) => `${p}:mcp:session:${encodeURIComponent(threadId)}`; + const aad = (threadId: string, serverId: string) => + `${p}\u0000mcp-session\u0000${threadId}\u0000${serverId}`; + + const pack = async ( + threadId: string, + serverId: string, + descriptor: McpServerDescriptor, + ): Promise => { + const d: Record = {}; + const m: string[] = []; + for (const field of DESCRIPTOR_FIELDS) { + const value = descriptor[field]; + if (value === undefined) continue; + if (isJsonSafe(value)) d[field] = value; + else m.push(field); + } + let sealed: Sealed | undefined; + if (descriptor.transport !== undefined) { + const transport: Record = {}; + for (const [k, v] of Object.entries(descriptor.transport)) { + if (v === undefined) continue; + if (isJsonSafe(v)) transport[k] = v; + else m.push(`transport.${k}`); + } + if (box) { + const secret: Record = {}; + for (const k of ["headers", "env"]) { + if (transport[k] !== undefined) { + secret[k] = transport[k]; + delete transport[k]; + } + } + if (Object.keys(secret).length > 0) + sealed = await box.seal(secret, aad(threadId, serverId)); + } + d.transport = transport; + } + return { d: d as unknown as McpServerDescriptor, m, ...(sealed ? { s: sealed } : {}) }; + }; + + const unpack = async ( + threadId: string, + serverId: string, + stored: StoredServer, + ): Promise => { + const descriptor = { ...stored.d } as McpServerDescriptor & Record; + if (stored.s) { + if (!box) { + throw new Error( + `@upstash/agentkit-tanstack-ai: MCP session server ${JSON.stringify(serverId)} has encrypted headers; pass \`encryptionKey\`.`, + ); + } + const secret = await box.open>(stored.s, aad(threadId, serverId)); + descriptor.transport = { + ...descriptor.transport, + ...secret, + } as McpServerDescriptor["transport"]; + } + const overrides = config.resolve + ? await config.resolve({ threadId, serverId, descriptor, missing: stored.m }) + : undefined; + const { fetch: fetchOverride, authProvider, ...fields } = overrides ?? {}; + Object.assign(descriptor, fields); + if ((fetchOverride !== undefined || authProvider !== undefined) && descriptor.transport) { + descriptor.transport = { + ...descriptor.transport, + ...(fetchOverride !== undefined ? { fetch: fetchOverride } : {}), + ...(authProvider !== undefined ? { authProvider } : {}), + } as McpServerDescriptor["transport"]; + } + // Fail closed: every code-only field the descriptor was saved with must be back. + for (const path of stored.m) { + const [head, sub] = path.split(".") as [string, string | undefined]; + const holder = sub === undefined ? descriptor : (descriptor[head] as Record); + if (holder?.[sub ?? head] === undefined) return null; + } + return descriptor; + }; + + return { + async set(threadId, servers) { + const entries = await Promise.all( + Object.entries(servers).map( + async ([serverId, descriptor]) => + [serverId, JSON.stringify(await pack(threadId, serverId, descriptor))] as const, + ), + ); + const tx = redis.multi(); + tx.del(key(threadId)); + if (entries.length > 0) { + tx.hset(key(threadId), Object.fromEntries(entries)); + if (ttl > 0) tx.expire(key(threadId), ttl); + } + await tx.exec(); + }, + + async get(threadId, serverId) { + let id: string; + let stored: StoredServer; + if (serverId === undefined) { + // The widget omitted serverId: only a thread with exactly one server resolves. + const all = await redis.hgetall>(key(threadId)); + const entries = Object.entries(all ?? {}); + if (entries.length !== 1) return null; + [id, stored] = entries[0]!; + } else { + const one = await redis.hget(key(threadId), serverId); + if (!one) return null; + [id, stored] = [serverId, one]; + } + if (ttl > 0) await redis.expire(key(threadId), ttl); + return unpack(threadId, id, stored); + }, + }; +} diff --git a/packages/tanstack-ai/src/secret-box.ts b/packages/tanstack-ai/src/secret-box.ts new file mode 100644 index 00000000..075e6b43 --- /dev/null +++ b/packages/tanstack-ai/src/secret-box.ts @@ -0,0 +1,79 @@ +// AES-256-GCM helpers shared by the stores that keep secrets at rest (BYOK credentials, MCP session +// headers). Web Crypto only, so they run on Node, edge and serverless runtimes alike. + +export const b64 = { + // Chunked: spreading a large array into String.fromCharCode overflows the call stack. + encode(bytes: Uint8Array): string { + let binary = ""; + for (let i = 0; i < bytes.length; i += 0x8000) { + binary += String.fromCharCode(...bytes.subarray(i, i + 0x8000)); + } + return globalThis.btoa(binary); + }, + decode: (text: string) => Uint8Array.from(globalThis.atob(text), (c) => c.charCodeAt(0)), +}; + +const KEY_HELP = + "@upstash/agentkit-tanstack-ai: `encryptionKey` must be 32 bytes: a Uint8Array, or a base64, base64url or hex string (e.g. `openssl rand -base64 32`)."; + +/** A 32-byte key from a `Uint8Array`, or a base64 / base64url / hex string. */ +export function toKeyBytes(key: string | Uint8Array): Uint8Array { + let bytes: Uint8Array; + if (typeof key !== "string") bytes = key; + else { + const text = key.trim(); + try { + bytes = /^[0-9a-f]{64}$/i.test(text) + ? Uint8Array.from(text.match(/../g)!, (h) => parseInt(h, 16)) + : b64.decode( + text + .replace(/-/g, "+") + .replace(/_/g, "/") + .padEnd(Math.ceil(text.length / 4) * 4, "="), + ); + } catch { + throw new Error(KEY_HELP); + } + } + if (bytes.length !== 32) throw new Error(KEY_HELP); + return bytes; +} + +export interface Sealed { + /** Base64 IV. */ + iv: string; + /** Base64 ciphertext (with the GCM tag). */ + ct: string; +} + +const utf8 = new TextEncoder(); + +/** Seal and open JSON values with one key; `aad` binds a ciphertext to where it is stored. */ +export function secretBox(key: string | Uint8Array) { + const keyBytes = toKeyBytes(key); + let cryptoKey: ReturnType | undefined; + const getKey = () => + (cryptoKey ??= crypto.subtle.importKey("raw", keyBytes as never, "AES-GCM", false, [ + "encrypt", + "decrypt", + ])); + return { + async seal(value: unknown, aad: string): Promise { + const iv = crypto.getRandomValues(new Uint8Array(12)); + const ct = await crypto.subtle.encrypt( + { name: "AES-GCM", iv, additionalData: utf8.encode(aad) }, + await getKey(), + utf8.encode(JSON.stringify(value)), + ); + return { iv: b64.encode(iv), ct: b64.encode(new Uint8Array(ct)) }; + }, + async open(sealed: Sealed, aad: string): Promise { + const plain = await crypto.subtle.decrypt( + { name: "AES-GCM", iv: b64.decode(sealed.iv), additionalData: utf8.encode(aad) }, + await getKey(), + b64.decode(sealed.ct), + ); + return JSON.parse(new TextDecoder().decode(plain)) as T; + }, + }; +} From 8b3aff76afb39830ad656f1327f9791e7f7a2470 Mon Sep 17 00:00:00 2001 From: CahidArda Date: Fri, 13 Nov 2026 14:00:02 +0300 Subject: [PATCH 3/7] fix(tanstack-ai): don't drop a cross-instance log write that lands before the first poll subscribe() took its baseline from its first poll, which runs after subscribe returns. A write from another machine in between became the baseline and was never reported. Start each watch from this store's own last write to the log instead. --- .changeset/log-subscribe-race.md | 5 +++ .../tanstack-ai/src/persistence/log-store.ts | 33 +++++++++++-------- 2 files changed, 25 insertions(+), 13 deletions(-) create mode 100644 .changeset/log-subscribe-race.md diff --git a/.changeset/log-subscribe-race.md b/.changeset/log-subscribe-race.md new file mode 100644 index 00000000..37ff82af --- /dev/null +++ b/.changeset/log-subscribe-race.md @@ -0,0 +1,5 @@ +--- +"@upstash/agentkit-tanstack-ai": patch +--- + +Fix a race in the `upstashPersistence()` log store's `subscribe`: a write from another machine that landed before the subscriber's first poll was treated as the starting point and never reported, so the listener could miss it. A watch now starts from the store's own last write to that log. diff --git a/packages/tanstack-ai/src/persistence/log-store.ts b/packages/tanstack-ai/src/persistence/log-store.ts index 6728934c..5883d688 100644 --- a/packages/tanstack-ai/src/persistence/log-store.ts +++ b/packages/tanstack-ai/src/persistence/log-store.ts @@ -237,9 +237,8 @@ interface Watch { listeners: Set<() => void>; /** The highest change counter the listeners were told about: each counter value is reported once. */ reported: number; - /** The highest counter a poll has read. The first read is the baseline, not a change. */ + /** The highest counter a poll has read. */ observed: number; - baselined: boolean; /** A poll is in flight: the next tick skips, so slow requests never pile up. */ busy: boolean; timer?: ReturnType; @@ -290,8 +289,18 @@ export function redisLogStore( } } }; + /** + * The latest change counter of each log this store wrote, so a later subscribe knows which writes + * it already saw. Bounded: the oldest entries go first, and a forgotten log only costs one extra + * listener call on its next subscribe. + */ + const written = new Map(); + const MAX_WRITTEN = 10_000; /** A write through this store: remember its counter so the poll does not report it twice, then call listeners. */ const changed = (logId: string, tick: unknown) => { + written.delete(logId); + written.set(logId, Number(tick)); + if (written.size > MAX_WRITTEN) written.delete(written.keys().next().value!); const watch = watches.get(logId); if (!watch) return; // A poll may have seen this write's counter before the write's own reply arrived, and already @@ -369,28 +378,26 @@ export function redisLogStore( if (!watch) { const created: Watch = { listeners: new Set(), - reported: 0, - observed: 0, - baselined: false, + reported: written.get(logId) ?? 0, + observed: written.get(logId) ?? 0, busy: false, }; watch = created; watches.set(logId, created); const tickKey = keys(logId)[6]!; - // The first successful read of the counter is the baseline; later reads above it are writes - // from other machines. A failed read (a REST blip) just tries again on the next tick. + // A counter above the last one known means a write this store has not reported. The watch + // starts from this store's own last write to the log, not from its first read: subscribe() + // returns before that read, so a write from another machine can land in between and must + // still be reported. A log written only by other machines before the subscribe gets one + // extra call; listeners re-read the log, so that is harmless, while a missed call would + // leave a reader waiting. A failed read (a REST blip) retries on the next tick. const check = async () => { if (created.busy) return; created.busy = true; try { const tick = Number((await redis.get(tickKey)) ?? 0); if (watches.get(logId) !== created) return; - if (!created.baselined) { - // Not marked as reported: if this read already includes a write made through this store, - // that write's own reply still has to call the listeners. - created.baselined = true; - created.observed = Math.max(created.observed, tick); - } else if (tick > created.observed) { + if (tick > created.observed) { created.observed = tick; if (tick > created.reported) { created.reported = tick; From aa0c60499a90c3f0e8017c94dd5dcf8c7a8a97d0 Mon Sep 17 00:00:00 2001 From: CahidArda Date: Fri, 13 Nov 2026 14:02:14 +0300 Subject: [PATCH 4/7] ci: pin @upstash/qstash-cli to 2.44.2 in the mcp-toolkit-demo e2e 2.44.3 exits with code 1 and no output, even for --help, so the e2e failed before it started. --- examples/mcp-toolkit-demo/package.json | 2 +- examples/mcp-toolkit-demo/scripts/e2e.mjs | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/examples/mcp-toolkit-demo/package.json b/examples/mcp-toolkit-demo/package.json index 31ca0b16..afd3c8d0 100644 --- a/examples/mcp-toolkit-demo/package.json +++ b/examples/mcp-toolkit-demo/package.json @@ -6,7 +6,7 @@ "dev": "next dev", "build": "next build", "start": "next start", - "qstash": "npx @upstash/qstash-cli dev", + "qstash": "npx @upstash/qstash-cli@2.44.2 dev", "smoke": "node scripts/smoke.mjs", "e2e": "node scripts/e2e.mjs" }, diff --git a/examples/mcp-toolkit-demo/scripts/e2e.mjs b/examples/mcp-toolkit-demo/scripts/e2e.mjs index 290dd782..972215f6 100644 --- a/examples/mcp-toolkit-demo/scripts/e2e.mjs +++ b/examples/mcp-toolkit-demo/scripts/e2e.mjs @@ -68,7 +68,9 @@ async function waitFor(what, ready, seconds = 90) { } // 1. The QStash dev server. Its token and signing keys are read from what it prints. -const qstash = start("QStash dev server", "npx", ["-y", "@upstash/qstash-cli", "dev", "-port", QSTASH_PORT]); +// Pinned: @upstash/qstash-cli 2.44.3 exits with code 1 and no output, even for `--help`. +const QSTASH_CLI = "@upstash/qstash-cli@2.44.2"; +const qstash = start("QStash dev server", "npx", ["-y", QSTASH_CLI, "dev", "-port", QSTASH_PORT]); const credentials = await waitFor("the QStash dev server", async () => { const read = (name) => new RegExp(`${name}=(\\S+)`).exec(qstash.output)?.[1]; const [token, current, next] = ["QSTASH_TOKEN", "QSTASH_CURRENT_SIGNING_KEY", "QSTASH_NEXT_SIGNING_KEY"].map(read); From 66318551eca59b9da162d5fd66038895c43274b8 Mon Sep 17 00:00:00 2001 From: CahidArda Date: Fri, 13 Nov 2026 14:11:07 +0300 Subject: [PATCH 5/7] ci: unpin @upstash/qstash-cli in the mcp-toolkit-demo e2e latest points to a working release again. --- examples/mcp-toolkit-demo/package.json | 2 +- examples/mcp-toolkit-demo/scripts/e2e.mjs | 4 +--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/examples/mcp-toolkit-demo/package.json b/examples/mcp-toolkit-demo/package.json index afd3c8d0..31ca0b16 100644 --- a/examples/mcp-toolkit-demo/package.json +++ b/examples/mcp-toolkit-demo/package.json @@ -6,7 +6,7 @@ "dev": "next dev", "build": "next build", "start": "next start", - "qstash": "npx @upstash/qstash-cli@2.44.2 dev", + "qstash": "npx @upstash/qstash-cli dev", "smoke": "node scripts/smoke.mjs", "e2e": "node scripts/e2e.mjs" }, diff --git a/examples/mcp-toolkit-demo/scripts/e2e.mjs b/examples/mcp-toolkit-demo/scripts/e2e.mjs index 972215f6..290dd782 100644 --- a/examples/mcp-toolkit-demo/scripts/e2e.mjs +++ b/examples/mcp-toolkit-demo/scripts/e2e.mjs @@ -68,9 +68,7 @@ async function waitFor(what, ready, seconds = 90) { } // 1. The QStash dev server. Its token and signing keys are read from what it prints. -// Pinned: @upstash/qstash-cli 2.44.3 exits with code 1 and no output, even for `--help`. -const QSTASH_CLI = "@upstash/qstash-cli@2.44.2"; -const qstash = start("QStash dev server", "npx", ["-y", QSTASH_CLI, "dev", "-port", QSTASH_PORT]); +const qstash = start("QStash dev server", "npx", ["-y", "@upstash/qstash-cli", "dev", "-port", QSTASH_PORT]); const credentials = await waitFor("the QStash dev server", async () => { const read = (name) => new RegExp(`${name}=(\\S+)`).exec(qstash.output)?.[1]; const [token, current, next] = ["QSTASH_TOKEN", "QSTASH_CURRENT_SIGNING_KEY", "QSTASH_NEXT_SIGNING_KEY"].map(read); From cca3b119a68ece57aa2ace79debc61d2ca16da68 Mon Sep 17 00:00:00 2001 From: CahidArda Date: Fri, 13 Nov 2026 15:57:59 +0300 Subject: [PATCH 6/7] fix(tanstack-ai): harden and simplify the sandbox and MCP stores - MCP sessions: with encryptionKey, seal each stored entry whole, so a Redis writer cannot change the server url or strip the fail-closed field list - MCP events: revisions come from one counter, so a stale update cannot match a deleted and re-created subscription; drop the key-to-id hash - fork: write the transcript first, so a failing JSON.SET leaves nothing half-written; reject a non-integer messagesTtlSeconds - sandbox validation: keep only the checks this store relies on - BYOK: use the shared secretBox (same stored format) --- .changeset/tanstack-sandbox-mcp-stores.md | 2 +- CLAUDE.md | 12 +- packages/tanstack-ai/README.md | 7 +- packages/tanstack-ai/src/byok/credentials.ts | 38 ++---- packages/tanstack-ai/src/mcp/event-store.ts | 30 ++-- .../tanstack-ai/src/mcp/session-store.test.ts | 23 +--- packages/tanstack-ai/src/mcp/session-store.ts | 48 +++---- .../src/sandbox/checkpoint-store.ts | 13 +- packages/tanstack-ai/src/sandbox/validate.ts | 128 ++++-------------- packages/tanstack-ai/src/secret-box.ts | 4 +- 10 files changed, 91 insertions(+), 214 deletions(-) diff --git a/.changeset/tanstack-sandbox-mcp-stores.md b/.changeset/tanstack-sandbox-mcp-stores.md index 85853997..cad90d6a 100644 --- a/.changeset/tanstack-sandbox-mcp-stores.md +++ b/.changeset/tanstack-sandbox-mcp-stores.md @@ -12,7 +12,7 @@ Add stores for TanStack AI's sandbox and MCP packages, all on Upstash Redis. - `upstashMcpTaskStore()` is the `taskStore` for `createMCPServer`. - `upstashMcpEventStore()` stores MCP Events webhook subscriptions, shared across instances. - `upstashMcpSessionStore()` is the `McpSessionStore` for MCP Apps. - - The JSON part of each server descriptor is stored in Redis. Transport headers and env can be encrypted. + - The JSON part of each server descriptor is stored in Redis. With `encryptionKey`, each stored entry is encrypted. - Functions such as `needsApproval` come back through your `resolve` callback. - It fails closed: it returns `null` rather than a descriptor that lost its approval gate. - `upstashPersistence({ bucket })` is now typed with a required `blobs` store, so you can pass it straight to `createSandboxSnapshots()`. diff --git a/CLAUDE.md b/CLAUDE.md index c74cdf1a..446f7b73 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1715,8 +1715,9 @@ find node_modules -path "*/zod/package.json" | while read f; do echo "$f $(node - **What can't be stored:** its `McpServerDescriptor` carries functions (`toolFilter`, `needsApproval`, `toolName`) and transport objects (`fetch`, `authProvider`). - **`set`:** stores only JSON-safe fields, checked by prototype and recursion in `isJsonSafe`, and - records the dropped paths in `m`. Transport `headers` / `env` are sealed with AES-GCM (`src/secret-box.ts`, - shared with BYOK; the AAD is prefix + thread + server). + records the dropped paths in `m`. With a key, the whole `{d, m}` is sealed with AES-GCM + (`src/secret-box.ts`, shared with BYOK; the AAD is prefix + thread + server), so a Redis writer + cannot change the url or strip `m`. - **`get`:** calls the user's `resolve({ threadId, serverId, descriptor, missing })`, then FAILS CLOSED, returning `null` if any recorded path is still missing. Upstream's `createMcpAppCallHandler` then falls back to its registry (`descriptor.x ?? fromRegistry?.x`, per field) or refuses. @@ -1741,9 +1742,10 @@ find node_modules -path "*/zod/package.json" | while read f; do echo "$f $(node (3) `runSandboxCheckpointForkConformance` creates the fork promise before reading destination state over the network, so a live-Redis rejection shows up as a vitest "unhandled error"; the test wraps `forkFromCheckpoint` with a no-op `.catch` (assertions unchanged). (4) Heads, index members, id-index - values and set members are `"~" + id`, because `@upstash/redis` JSON-parses `"007"` to `7`. (5) The - validation helpers mirror `@tanstack/ai-sandbox`'s unexported `validateEntries`/`validateArtifacts` - (`src/sandbox/validate.ts`) — re-diff on sandbox bumps. (6) `upstashPersistence` gained an overload: with + values and set members are `"~" + id`, because `@upstash/redis` JSON-parses `"007"` to `7`. (5) + `src/sandbox/validate.ts` checks only ids, `createdAt` and blob keys (they become hash fields); + TanStack re-checks paths, sizes and hashes on restore. FORK writes the transcript first: a Lua + script is not rolled back on error. (6) `upstashPersistence` gained an overload: with `bucket` the return type has a required `blobs`, which `createSandboxSnapshots` needs (`src/sandbox/wiring.test.ts` pins it). Also runs `runDurableRunFieldsConformance` on the run store. - **TanStack AI 1.0 contracts implemented (2026-11-05, `@upstash/agentkit-tanstack-ai` minor):** `stores.log` (`persistence/log-store.ts`, see below), `stores.sessions` (`SessionIndexStore`, `persistence/sessions.ts`) and `./byok` (`upstashCredentials`, `byok/credentials.ts`, `CredentialStore`). **Peer floors raised to `@tanstack/ai >=1.0.0` and `@tanstack/ai-persistence >=1.0.0`** (verified by packing 0.71.0/0.10.1: neither has `SessionIndexStore` / `@tanstack/ai/byok`; `ai-memory` floor untouched). Sessions: doc `{r,e}` per `(tenant, thread)` + one ZSET per tenant (score `updatedAt`, member `~` so the client never JSON-parses a numeric-looking id); `list` walks the ZSET in batches of whole score groups and re-reads a tie group in full so the `threadId` UTF-8 order is exact, filtering after loading; `ZRANGE BYSCORE REV` takes `(max, min)` (cast through `number`, the SDK types forbid `(`-exclusive strings). Revisions come from an `INCR` counter, never reused (the credential contract requires that; sessions may reuse). Passes the 19 `sessions` cases of `runPersistenceConformance` (they run only when the store is present). Credentials: one hash per owner (`%none` segment = undefined tenant/user, distinct from `""`), value `{r,t,x?,c | iv,ct}` with type/expiry in the clear so `list` never decrypts; optional AES-256-GCM `encryptionKey`, AAD = owner key + id (a copied ciphertext does not decrypt). `storedProviderKeys` lives in `@tanstack/ai/byok/server`, not `/byok`. **`stores.log` (`LogStore`):** per log a meta hash (head, first, epoch, owner, harness, exp, rel, wake, aff, del, laneH, oseq), a records ZSET (score seq, member `:`), snapshot string (`seq` must stay the first JSON key: `truncate` reads it with a Lua pattern), offers hash + info hash + pending ZSET, and a tick counter; shared `logDue` ZSET (when each log needs an owner) and `logHosts` hash. Every op is one Lua script over the same nine declared keys; the store clock is Redis `TIME` inside the script; numbers written from Lua go through `%.0f` (plain `tostring` prints epoch ms in scientific notation); ids returned from Lua are `~`-prefixed because the client JSON-parses numeric-looking strings (`"007"` -> 7). `subscribe` = immediate local callback + a `reported` and an `observed` tick per log so each counter value is reported once (a baseline read that already contains a local write must not mark it reported: that swallowed the notification), plus a 250 ms tick poll for other machines (non-overlapping; the first successful read is the baseline, so a failed first read is retried instead of silencing the subscription). `changed()` skips a counter the poll already reported (deterministic repro in tests: delay the append reply with a Proxy), which fixed a double notification (a poll can see the tick before the write reply arrives; it failed the `calls a listener after each append` conformance case in CI). The baseline read is still async, so a write by another machine landing in the first round trip after `subscribe` can go unreported (listeners are wake-ups, hosts reconcile with `read`). `append` rejects more than 10,000 records (a 50,000-record batch outran the Lua time limit and left part of the writes). Known follow-ups: `listLogsWithoutOwner` re-checks entries skipped by `harnesses` / a live pin on every call; consumed offers, old hosts and `deleteLog` tombstones are never cleaned up; `read` / `listOfferedInputs` without a limit return everything. Conformance needs the `twin` option (a second store set on the same keyspace) for the 4 `log across instances` cases; all 56 log cases and 19 sessions cases pass. Not implemented: `ownership` TTL-based cleanup of finished logs (logs and their keys live until `deleteLog`). diff --git a/packages/tanstack-ai/README.md b/packages/tanstack-ai/README.md index e68410d4..e9dc173d 100644 --- a/packages/tanstack-ai/README.md +++ b/packages/tanstack-ai/README.md @@ -304,7 +304,7 @@ import { createMcpAppCallHandler } from "@tanstack/ai-mcp/apps"; import { upstashMcpSessionStore } from "@upstash/agentkit-tanstack-ai/mcp"; const store = upstashMcpSessionStore({ - encryptionKey: process.env.MCP_SESSION_KEY, // encrypts transport headers / stdio env at rest + encryptionKey: process.env.MCP_SESSION_KEY, // encrypts each stored entry at rest // The parts of a descriptor that cannot live in Redis come from your code, per server: resolve: ({ threadId, serverId }) => ({ needsApproval: (tool) => tool.annotations?.readOnlyHint !== true, @@ -323,9 +323,8 @@ A server descriptor carries a policy for widget calls (`toolFilter`, `needsAppro often a transport `fetch` or OAuth `authProvider`. None of those survive JSON, so `set` splits each descriptor: -- **The JSON part goes to Redis:** the transport (with `headers` / `env` encrypted when you pass - `encryptionKey`), `prefix`, `requestOptions`, and an array `toolFilter`. One hash per thread, one field - per server. +- **The JSON part goes to Redis:** the transport, `prefix`, `requestOptions`, and an array `toolFilter`. + One hash per thread, one field per server, encrypted whole when you pass `encryptionKey`. - **Everything else is dropped and recorded by name.** `resolve` supplies it again on `get`, from your code. It receives `{ threadId, serverId, descriptor, missing }`. - **It fails closed.** If `resolve` doesn't return a recorded field (say `needsApproval`), `get` returns diff --git a/packages/tanstack-ai/src/byok/credentials.ts b/packages/tanstack-ai/src/byok/credentials.ts index 04316d8c..97cded05 100644 --- a/packages/tanstack-ai/src/byok/credentials.ts +++ b/packages/tanstack-ai/src/byok/credentials.ts @@ -3,7 +3,7 @@ import { Redis as RedisClient } from "@upstash/redis"; import type { Credential, CredentialScope, CredentialStore } from "@tanstack/ai/byok"; import { addTelemetry } from "../telemetry.js"; import { KEY_LOCKING, assertId } from "../persistence/records.js"; -import { b64, toKeyBytes } from "../secret-box.js"; +import { secretBox } from "../secret-box.js"; export interface UpstashCredentialsConfig { /** Upstash Redis client. Defaults to `Redis.fromEnv()`. */ @@ -98,17 +98,10 @@ export function upstashCredentials(config: UpstashCredentialsConfig = {}): Crede if (config.requireEncryption && config.encryptionKey === undefined) { throw new Error("@upstash/agentkit-tanstack-ai: `requireEncryption` needs an `encryptionKey`."); } - const keyBytes = - config.encryptionKey === undefined ? undefined : toKeyBytes(config.encryptionKey); - let cryptoKey: ReturnType | undefined; - const aesKey = () => - (cryptoKey ??= crypto.subtle.importKey("raw", keyBytes! as never, "AES-GCM", false, [ - "encrypt", - "decrypt", - ])); + const box = config.encryptionKey === undefined ? undefined : secretBox(config.encryptionKey); // Bind the ciphertext to its owner, id and the clear-text type/expiry that `list` shows. const aad = (scope: CredentialScope, id: string, type: string, expiresAt: number | undefined) => - new TextEncoder().encode(JSON.stringify([owner(scope), id, type, expiresAt ?? null])); + JSON.stringify([owner(scope), id, type, expiresAt ?? null]); /** The JSON text after `{"r":N,` for a credential. */ async function pack(scope: CredentialScope, id: string, credential: Credential): Promise { @@ -118,15 +111,8 @@ export function upstashCredentials(config: UpstashCredentialsConfig = {}): Crede ? { x: credential.expiresAt } : {}), }; - if (keyBytes) { - const iv = crypto.getRandomValues(new Uint8Array(12)); - const ct = await crypto.subtle.encrypt( - { name: "AES-GCM", iv, additionalData: aad(scope, id, stored.t, stored.x) }, - await aesKey(), - new TextEncoder().encode(JSON.stringify(credential)), - ); - stored.iv = b64.encode(iv); - stored.ct = b64.encode(new Uint8Array(ct)); + if (box) { + Object.assign(stored, await box.seal(credential, aad(scope, id, stored.t, stored.x))); } else { stored.c = credential; } @@ -147,21 +133,15 @@ export function upstashCredentials(config: UpstashCredentialsConfig = {}): Crede `@upstash/agentkit-tanstack-ai: credential ${JSON.stringify(id)} is corrupt.`, ); } - if (!keyBytes) { + if (!box) { throw new Error( `@upstash/agentkit-tanstack-ai: credential ${JSON.stringify(id)} is encrypted; pass \`encryptionKey\`.`, ); } - const plain = await crypto.subtle.decrypt( - { - name: "AES-GCM", - iv: b64.decode(stored.iv) as never, - additionalData: aad(scope, id, stored.t, stored.x), - }, - await aesKey(), - b64.decode(stored.ct) as never, + return box.open( + { iv: stored.iv, ct: stored.ct }, + aad(scope, id, stored.t, stored.x), ); - return JSON.parse(new TextDecoder().decode(plain)) as Credential; } const checkId = (id: string) => { diff --git a/packages/tanstack-ai/src/mcp/event-store.ts b/packages/tanstack-ai/src/mcp/event-store.ts index 99f57d1f..8a03af2a 100644 --- a/packages/tanstack-ai/src/mcp/event-store.ts +++ b/packages/tanstack-ai/src/mcp/event-store.ts @@ -17,30 +17,30 @@ export interface UpstashMcpEventStoreConfig { /** * Write a subscription and keep the id index in step. - * KEYS: document, id→key hash, key→id hash, key set. + * KEYS: document hash (fields d, r, i), id→key hash, key set, revision counter. * ARGV: key, document JSON, expected revision ("" = unconditional), id ("" = none). - * Returns the new revision, or -1 if the revision moved (a concurrent `update`). + * Returns the new revision, or -1 if the revision moved. Revisions come from one counter, so a + * delete and re-create never brings an old revision back. */ const WRITE = `${KEY_LOCKING}if ARGV[3] ~= "" and tostring(redis.call("HGET", KEYS[1], "r")) ~= ARGV[3] then return -1 end -local prev = redis.call("HGET", KEYS[3], ARGV[1]) +local prev = redis.call("HGET", KEYS[1], "i") if prev and prev ~= ARGV[4] then redis.call("HDEL", KEYS[2], prev) end -local rev = redis.call("HINCRBY", KEYS[1], "r", 1) -redis.call("HSET", KEYS[1], "d", ARGV[2]) +local rev = redis.call("INCR", KEYS[4]) +redis.call("HSET", KEYS[1], "d", ARGV[2], "r", rev) if ARGV[4] ~= "" then + redis.call("HSET", KEYS[1], "i", ARGV[4]) redis.call("HSET", KEYS[2], ARGV[4], "~" .. ARGV[1]) - redis.call("HSET", KEYS[3], ARGV[1], ARGV[4]) else - redis.call("HDEL", KEYS[3], ARGV[1]) + redis.call("HDEL", KEYS[1], "i") end -redis.call("SADD", KEYS[4], "~" .. ARGV[1]) +redis.call("SADD", KEYS[3], "~" .. ARGV[1]) return rev`; -/** KEYS: document, id→key hash, key→id hash, key set. ARGV: key. */ -const DELETE = `${KEY_LOCKING}local prev = redis.call("HGET", KEYS[3], ARGV[1]) +/** KEYS: document hash, id→key hash, key set. ARGV: key. */ +const DELETE = `${KEY_LOCKING}local prev = redis.call("HGET", KEYS[1], "i") if prev then redis.call("HDEL", KEYS[2], prev) end -redis.call("HDEL", KEYS[3], ARGV[1]) redis.call("DEL", KEYS[1]) -redis.call("SREM", KEYS[4], "~" .. ARGV[1]) +redis.call("SREM", KEYS[3], "~" .. ARGV[1]) return 1`; interface Stored { @@ -79,10 +79,10 @@ export function upstashMcpEventStore(config: UpstashMcpEventStoreConfig = {}): M const k = { doc: (key: string) => `${p}:mcp:eventSub:${encodeURIComponent(key)}`, ids: `${p}:mcp:eventSubIds`, - keyIds: `${p}:mcp:eventSubKeyIds`, all: `${p}:mcp:eventSubs`, + rev: `${p}:mcp:eventSubRev`, }; - const keys = (key: string) => [k.doc(key), k.ids, k.keyIds, k.all]; + const keys = (key: string) => [k.doc(key), k.ids, k.all, k.rev]; const load = async (key: string): Promise => { const hash = await redis.hgetall<{ d: MCPEventSubscription; r: number }>(k.doc(key)); @@ -128,7 +128,7 @@ export function upstashMcpEventStore(config: UpstashMcpEventStoreConfig = {}): M }, async delete(key) { - await redis.eval(DELETE, keys(key), [key]); + await redis.eval(DELETE, [k.doc(key), k.ids, k.all], [key]); }, async *list() { diff --git a/packages/tanstack-ai/src/mcp/session-store.test.ts b/packages/tanstack-ai/src/mcp/session-store.test.ts index fa6734d1..e76bab42 100644 --- a/packages/tanstack-ai/src/mcp/session-store.test.ts +++ b/packages/tanstack-ai/src/mcp/session-store.test.ts @@ -135,7 +135,7 @@ describe.skipIf(!hasRedisCreds)("upstashMcpSessionStore (live Redis)", () => { }); }); - it("encrypts headers at rest, bound to the thread and server", async () => { + it("encrypts each entry at rest, bound to the thread and server", async () => { const p = next(); const store = upstashMcpSessionStore({ redis, @@ -146,7 +146,7 @@ describe.skipIf(!hasRedisCreds)("upstashMcpSessionStore (live Redis)", () => { await store.set("t", { docs: docs() }); const raw = JSON.stringify(await redis.hgetall(`${p}:mcp:session:t`)); expect(raw).not.toContain("s3cr3t-token"); - expect(raw).toContain(mcp.url); + expect(raw).not.toContain(mcp.url); expect((await store.get("t", "docs"))?.transport).toMatchObject({ headers: { authorization: TOKEN }, }); @@ -156,7 +156,7 @@ describe.skipIf(!hasRedisCreds)("upstashMcpSessionStore (live Redis)", () => { docs: (await redis.hget(`${p}:mcp:session:t`, "docs"))!, }); await expect(store.get("u", "docs")).rejects.toThrow(); - // And a store without the key refuses to hand out a descriptor with its headers missing. + // And a store without the key refuses instead of handing out a partial descriptor. await expect(upstashMcpSessionStore({ redis, prefix: p }).get("t", "docs")).rejects.toThrow( /encryptionKey/, ); @@ -200,22 +200,5 @@ describe.skipIf(!hasRedisCreds)("upstashMcpSessionStore (live Redis)", () => { }); expect(remove).toEqual({ ok: false, error: "Tool needs approval: delete_doc" }); }); - - it("refuses the call instead of dropping the gate when resolve does not supply it", async () => { - const p = next(); - await upstashMcpSessionStore({ redis, prefix: p }).set("thread-1", { docs: docs() }); - const handle = createMcpAppCallHandler({ - clients: [], - store: upstashMcpSessionStore({ redis, prefix: p }), - }); - const remove = await handle({ - threadId: "thread-1", - serverId: "docs", - toolName: "delete_doc", - args: { id: "7" }, - }); - expect(remove.ok).toBe(false); - expect(JSON.stringify(remove)).not.toContain("deleted"); - }); }); }); diff --git a/packages/tanstack-ai/src/mcp/session-store.ts b/packages/tanstack-ai/src/mcp/session-store.ts index 142421ca..6153ba10 100644 --- a/packages/tanstack-ai/src/mcp/session-store.ts +++ b/packages/tanstack-ai/src/mcp/session-store.ts @@ -61,15 +61,15 @@ export interface UpstashMcpSessionStoreConfig { enableTelemetry?: boolean; } -interface StoredServer { - /** The JSON-safe part of the descriptor (headers/env removed when encrypted). */ +/** The JSON-safe part of a descriptor plus the paths of the code-only fields `resolve` must supply. */ +interface PlainServer { d: McpServerDescriptor; - /** Sealed `{ headers?, env? }` of the transport, when an `encryptionKey` is set. */ - s?: Sealed; - /** Paths of the code-only fields the descriptor had, which `resolve` must supply again. */ m: string[]; } +/** Stored as is, or the whole entry sealed when an `encryptionKey` is set (so nothing can be altered). */ +type StoredServer = PlainServer | { s: Sealed }; + const DESCRIPTOR_FIELDS = [ "prefix", "clientOptions", @@ -100,7 +100,7 @@ function isJsonSafe(value: unknown): boolean { * server), and anything that is not JSON (`needsApproval`, a function `toolFilter`, `toolName`, a * transport's `fetch` / `authProvider`) is dropped and recorded by name. `resolve` puts those back * on `get`, from your code, per server. If it does not, `get` returns `null` rather than a descriptor - * missing its approval gate or tool filter. Transport headers and stdio env can be encrypted at rest. + * missing its approval gate or tool filter. With `encryptionKey`, each stored entry is encrypted whole. * * For servers you configure in code and pass to `createMcpAppCallHandler({ clients })`, the handler * already takes `toolFilter` / `needsApproval` / `toolName` / `requestOptions` from that registry when @@ -128,7 +128,7 @@ export function upstashMcpSessionStore(config: UpstashMcpSessionStoreConfig = {} const box = config.encryptionKey === undefined ? undefined : secretBox(config.encryptionKey); const key = (threadId: string) => `${p}:mcp:session:${encodeURIComponent(threadId)}`; const aad = (threadId: string, serverId: string) => - `${p}\u0000mcp-session\u0000${threadId}\u0000${serverId}`; + `${p}\u0000mcp-session\u0000${encodeURIComponent(threadId)}\u0000${serverId}`; const pack = async ( threadId: string, @@ -143,7 +143,6 @@ export function upstashMcpSessionStore(config: UpstashMcpSessionStoreConfig = {} if (isJsonSafe(value)) d[field] = value; else m.push(field); } - let sealed: Sealed | undefined; if (descriptor.transport !== undefined) { const transport: Record = {}; for (const [k, v] of Object.entries(descriptor.transport)) { @@ -151,20 +150,10 @@ export function upstashMcpSessionStore(config: UpstashMcpSessionStoreConfig = {} if (isJsonSafe(v)) transport[k] = v; else m.push(`transport.${k}`); } - if (box) { - const secret: Record = {}; - for (const k of ["headers", "env"]) { - if (transport[k] !== undefined) { - secret[k] = transport[k]; - delete transport[k]; - } - } - if (Object.keys(secret).length > 0) - sealed = await box.seal(secret, aad(threadId, serverId)); - } d.transport = transport; } - return { d: d as unknown as McpServerDescriptor, m, ...(sealed ? { s: sealed } : {}) }; + const plain: PlainServer = { d: d as unknown as McpServerDescriptor, m }; + return box ? { s: await box.seal(plain, aad(threadId, serverId)) } : plain; }; const unpack = async ( @@ -172,21 +161,20 @@ export function upstashMcpSessionStore(config: UpstashMcpSessionStoreConfig = {} serverId: string, stored: StoredServer, ): Promise => { - const descriptor = { ...stored.d } as McpServerDescriptor & Record; - if (stored.s) { + let plain: PlainServer; + if ("s" in stored) { if (!box) { throw new Error( - `@upstash/agentkit-tanstack-ai: MCP session server ${JSON.stringify(serverId)} has encrypted headers; pass \`encryptionKey\`.`, + `@upstash/agentkit-tanstack-ai: MCP session server ${JSON.stringify(serverId)} is encrypted; pass \`encryptionKey\`.`, ); } - const secret = await box.open>(stored.s, aad(threadId, serverId)); - descriptor.transport = { - ...descriptor.transport, - ...secret, - } as McpServerDescriptor["transport"]; + plain = await box.open(stored.s, aad(threadId, serverId)); + } else { + plain = stored; } + const descriptor = { ...plain.d } as McpServerDescriptor & Record; const overrides = config.resolve - ? await config.resolve({ threadId, serverId, descriptor, missing: stored.m }) + ? await config.resolve({ threadId, serverId, descriptor, missing: plain.m }) : undefined; const { fetch: fetchOverride, authProvider, ...fields } = overrides ?? {}; Object.assign(descriptor, fields); @@ -198,7 +186,7 @@ export function upstashMcpSessionStore(config: UpstashMcpSessionStoreConfig = {} } as McpServerDescriptor["transport"]; } // Fail closed: every code-only field the descriptor was saved with must be back. - for (const path of stored.m) { + for (const path of plain.m) { const [head, sub] = path.split(".") as [string, string | undefined]; const holder = sub === undefined ? descriptor : (descriptor[head] as Record); if (holder?.[sub ?? head] === undefined) return null; diff --git a/packages/tanstack-ai/src/sandbox/checkpoint-store.ts b/packages/tanstack-ai/src/sandbox/checkpoint-store.ts index 909b0647..c6d49648 100644 --- a/packages/tanstack-ai/src/sandbox/checkpoint-store.ts +++ b/packages/tanstack-ai/src/sandbox/checkpoint-store.ts @@ -144,12 +144,13 @@ end for i = 8, 11 do if redis.call("ZCARD", KEYS[i]) > 0 then return "NOT_EMPTY" end end -redis.call("SET", KEYS[3], ARGV[4]) -redis.call("SET", KEYS[4], ARGV[6]) -redis.call("ZADD", KEYS[5], ARGV[7], ARGV[6]) +-- The transcript first: a script is not rolled back on error, so a failing JSON.SET must leave nothing. redis.call("JSON.SET", KEYS[7], "$", ARGV[5]) if ARGV[8] ~= "0" then redis.call("EXPIRE", KEYS[7], ARGV[8]) end for i = 9, #ARGV do redis.call("HINCRBY", KEYS[6], ARGV[i], 1) end +redis.call("SET", KEYS[3], ARGV[4]) +redis.call("SET", KEYS[4], ARGV[6]) +redis.call("ZADD", KEYS[5], ARGV[7], ARGV[6]) return "OK"`; /** A checkpoint is JSON by contract; copying through JSON also detaches it from the caller's object. */ @@ -198,6 +199,12 @@ export function upstashSandboxCheckpointStore( // The lease key outlives its lease by a day: expiry is decided by `expiresAt`, the PX only // garbage-collects leases nobody released. const px = String(Math.ceil(leaseDurationMs) + 86_400_000); + if ( + config.messagesTtlSeconds !== undefined && + !(Number.isSafeInteger(config.messagesTtlSeconds) && config.messagesTtlSeconds > 0) + ) { + throw new RangeError("messagesTtlSeconds must be a positive integer"); + } const ttl = String(config.messagesTtlSeconds ?? 0); const enc = encodeURIComponent; diff --git a/packages/tanstack-ai/src/sandbox/validate.ts b/packages/tanstack-ai/src/sandbox/validate.ts index 7959cbf1..7360e9b9 100644 --- a/packages/tanstack-ai/src/sandbox/validate.ts +++ b/packages/tanstack-ai/src/sandbox/validate.ts @@ -4,9 +4,6 @@ import { type SandboxCheckpoint, } from "@tanstack/ai-sandbox"; -// Input checks shared by the checkpoint store. They mirror `@tanstack/ai-sandbox`'s own in-memory -// store (which does not export them) so both stores refuse the same inputs with the same codes. - const utf8 = new TextEncoder(); /** Order by UTF-8 bytes, the order the contract specifies for checkpoint ids and blob keys. */ @@ -38,126 +35,47 @@ export function assertValidIdentifier(value: unknown, label: string): asserts va } } -const hasOwn = (value: object, key: string) => Object.prototype.hasOwnProperty.call(value, key); const FILE_KEY = /^sandbox-files\/sha256\/[0-9a-f]{64}$/; const ARTIFACT_KEY = /^sandbox-artifacts\/sha256\/[0-9a-f]{64}$/; +const invalid = (message: string) => new SandboxCheckpointInvalidEntryError(message); + +// Only what this store relies on: blob keys become Redis hash fields, so they must be well formed. +// TanStack re-checks paths, sizes and file hashes when it restores a snapshot. function validateEntries(checkpoint: SandboxCheckpoint): void { - if (!Array.isArray(checkpoint.files)) { - throw new SandboxCheckpointInvalidEntryError("Checkpoint files must be an array"); - } - const paths = new Set(); - const kinds = new Map(); + if (!Array.isArray(checkpoint.files)) throw invalid("Checkpoint files must be an array"); for (const entry of checkpoint.files as unknown[]) { - if (entry === null || typeof entry !== "object") { - throw new SandboxCheckpointInvalidEntryError("Checkpoint entry must be an object"); - } - const c = entry as Record; - const path = c.path; - if ( - typeof path !== "string" || - path.length === 0 || - path.includes("\0") || - path.startsWith("/") || - path.startsWith("\\") || - /^[A-Za-z]:([\\/]|$)/.test(path) || - path.includes("\\") || - path.split("/").some((part) => part.length === 0 || part === "." || part === "..") - ) { - throw new SandboxCheckpointInvalidEntryError( - "Checkpoint entry path must be a normalized workspace-relative path", - ); - } - if (paths.has(path)) { - throw new SandboxCheckpointInvalidEntryError( - `Checkpoint contains duplicate entry path '${path}'`, - ); - } - for (let s = path.indexOf("/"); s !== -1; s = path.indexOf("/", s + 1)) { - const ancestor = path.slice(0, s); - if (kinds.get(ancestor) === "file") { - throw new SandboxCheckpointInvalidEntryError( - `Checkpoint entry '${path}' is beneath file '${ancestor}'`, - ); + const e = entry as Record | null; + if (e === null || typeof e !== "object") throw invalid("Checkpoint entry must be an object"); + if (e.kind === "file") { + if (typeof e.blobKey !== "string" || !FILE_KEY.test(e.blobKey)) { + throw invalid("File entries require a valid content-addressed blobKey"); } + } else if (e.kind !== "dir") { + throw invalid("Checkpoint entry kind must be file or dir"); } - if (c.kind === "file" && [...kinds.keys()].some((other) => other.startsWith(`${path}/`))) { - throw new SandboxCheckpointInvalidEntryError( - `Checkpoint file '${path}' is an ancestor of another entry`, - ); - } - paths.add(path); - if (c.kind === "file") { - if ( - typeof c.blobKey !== "string" || - hasUnpairedSurrogate(c.blobKey) || - !FILE_KEY.test(c.blobKey) - ) { - throw new SandboxCheckpointInvalidEntryError( - "File entries require a valid content-addressed blobKey", - ); - } - if ( - !hasOwn(c, "size") || - typeof c.size !== "number" || - !Number.isSafeInteger(c.size) || - c.size < 0 - ) { - throw new SandboxCheckpointInvalidEntryError( - "File entry size must be a non-negative safe integer", - ); - } - } else if (c.kind === "dir") { - if (hasOwn(c, "blobKey") || hasOwn(c, "size")) { - throw new SandboxCheckpointInvalidEntryError( - "Directory entries cannot contain file fields", - ); - } - } else { - throw new SandboxCheckpointInvalidEntryError("Checkpoint entry kind must be file or dir"); - } - kinds.set(path, c.kind); } } function validateArtifacts(checkpoint: SandboxCheckpoint): void { - if (!Array.isArray(checkpoint.artifacts)) { - throw new SandboxCheckpointInvalidEntryError("Checkpoint artifacts must be an array"); - } + if (!Array.isArray(checkpoint.artifacts)) throw invalid("Checkpoint artifacts must be an array"); for (const artifact of checkpoint.artifacts as unknown[]) { - if (artifact === null || typeof artifact !== "object") { - throw new SandboxCheckpointInvalidEntryError("Checkpoint artifact must be an object"); + const a = artifact as Record | null; + if ( + a === null || + typeof a !== "object" || + typeof a.blobKey !== "string" || + !ARTIFACT_KEY.test(a.blobKey) + ) { + throw invalid("Checkpoint artifact has invalid fields"); } - const a = artifact as Record; - const ok = - typeof a.artifactId === "string" && - a.artifactId.length > 0 && - !hasUnpairedSurrogate(a.artifactId) && - typeof a.name === "string" && - a.name.length > 0 && - typeof a.mimeType === "string" && - a.mimeType.length > 0 && - typeof a.blobKey === "string" && - !hasUnpairedSurrogate(a.blobKey) && - ARTIFACT_KEY.test(a.blobKey) && - typeof a.size === "number" && - Number.isSafeInteger(a.size) && - a.size >= 0 && - typeof a.createdAt === "number" && - Number.isFinite(a.createdAt); - if (!ok) throw new SandboxCheckpointInvalidEntryError("Checkpoint artifact has invalid fields"); } } -/** Every field check the contract makes on a checkpoint, ids included. */ +/** Field checks on a checkpoint; callers check its ids first. */ export function validateCheckpoint(checkpoint: SandboxCheckpoint): void { - assertValidIdentifier(checkpoint.id, "Checkpoint id"); - assertValidIdentifier(checkpoint.threadId, "Checkpoint thread id"); - if (checkpoint.parentCheckpointId != null) { - assertValidIdentifier(checkpoint.parentCheckpointId, "Parent checkpoint id"); - } if (typeof checkpoint.createdAt !== "number" || !Number.isFinite(checkpoint.createdAt)) { - throw new SandboxCheckpointInvalidEntryError("Checkpoint createdAt must be a finite number"); + throw invalid("Checkpoint createdAt must be a finite number"); } validateEntries(checkpoint); validateArtifacts(checkpoint); diff --git a/packages/tanstack-ai/src/secret-box.ts b/packages/tanstack-ai/src/secret-box.ts index 075e6b43..7c08093f 100644 --- a/packages/tanstack-ai/src/secret-box.ts +++ b/packages/tanstack-ai/src/secret-box.ts @@ -1,7 +1,7 @@ // AES-256-GCM helpers shared by the stores that keep secrets at rest (BYOK credentials, MCP session // headers). Web Crypto only, so they run on Node, edge and serverless runtimes alike. -export const b64 = { +const b64 = { // Chunked: spreading a large array into String.fromCharCode overflows the call stack. encode(bytes: Uint8Array): string { let binary = ""; @@ -17,7 +17,7 @@ const KEY_HELP = "@upstash/agentkit-tanstack-ai: `encryptionKey` must be 32 bytes: a Uint8Array, or a base64, base64url or hex string (e.g. `openssl rand -base64 32`)."; /** A 32-byte key from a `Uint8Array`, or a base64 / base64url / hex string. */ -export function toKeyBytes(key: string | Uint8Array): Uint8Array { +function toKeyBytes(key: string | Uint8Array): Uint8Array { let bytes: Uint8Array; if (typeof key !== "string") bytes = key; else { From 53f61954036c3526aca582895e5dbea8d3ddbe4a Mon Sep 17 00:00:00 2001 From: CahidArda Date: Fri, 13 Nov 2026 16:20:04 +0300 Subject: [PATCH 7/7] test(tanstack-ai): cover the log subscribe race; fix stale docs - regression test: another store's write before the first poll is reported - drop references to the removed baseline read and header-only encryption - move the MCP session store notes in CLAUDE.md under the stores entry --- CLAUDE.md | 31 ++++++++++--------- packages/tanstack-ai/README.md | 2 +- packages/tanstack-ai/src/mcp/event-store.ts | 2 +- packages/tanstack-ai/src/mcp/session-store.ts | 9 +++--- .../src/persistence/persistence.test.ts | 26 +++++++++++++--- 5 files changed, 43 insertions(+), 27 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 446f7b73..428cd7cb 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1711,19 +1711,6 @@ find node_modules -path "*/zod/package.json" | while read f; do echo "$f $(node need_tools/replay loop), QStash-backed background runs writing to `upstashStream`, a QStash schedule for `reapDetachedRuns`, `@tanstack/ai-code-mode`'s `AgentStore` and `@tanstack/ai-code-mode-snippets`' `SnippetStorage` (both small keyed stores; `ai-code-mode` peers the pre-1.0 `@tanstack/ai-harness`). - **`McpSessionStore` (MCP Apps) is `upstashMcpSessionStore`:** - - **What can't be stored:** its `McpServerDescriptor` carries functions (`toolFilter`, `needsApproval`, - `toolName`) and transport objects (`fetch`, `authProvider`). - - **`set`:** stores only JSON-safe fields, checked by prototype and recursion in `isJsonSafe`, and - records the dropped paths in `m`. With a key, the whole `{d, m}` is sealed with AES-GCM - (`src/secret-box.ts`, shared with BYOK; the AAD is prefix + thread + server), so a Redis writer - cannot change the url or strip `m`. - - **`get`:** calls the user's `resolve({ threadId, serverId, descriptor, missing })`, then FAILS CLOSED, - returning `null` if any recorded path is still missing. Upstream's `createMcpAppCallHandler` then - falls back to its registry (`descriptor.x ?? fromRegistry?.x`, per field) or refuses. - - **Why fail closed:** a per-thread server that lost `needsApproval` would let a widget run a gated tool. - - **Test:** `src/mcp/session-store.test.ts` drives the real handler against a real `createMCPServer` - over node:http that requires the stored Authorization header. - **ai 7.0.172 (2026-11-13):** pins moved 7.0.170 → 7.0.172 with `@ai-sdk/*` refreshed (`openai ^4.0.117`, `provider ^4.0.38`, `react ^4.0.176`; one `ai@7.0.172` in the lockfile); no code change. Behaviour note from 7.0.171: `toUIMessageStream` now sends a failed tool call's real error text to the client (was "An error occurred."; opt back in with `redactToolErrors: true`), so any message an agentkit tool throws is user-visible. eve still on 0.91.0 (npm `latest` is `1.0.0-rc.13`, not adopted); published `@upstash/agentkit-eve-extension@0.19.0` + `@upstash/agentkit-eve@0.19.0` install on rc.13 with one eve/ai/zod, `eve build` clean, mocked `eve eval` 4/4. - **ai 7.0.170 + @tanstack/ai 1.4.0 + @tanstack/ai-openai 1.1.1 (2026-11-12):** devDependency/example pins only (no published-content change, no changeset); `@ai-sdk/*` ranges moved with `pnpm -r update "@ai-sdk/*"` (one `ai@7.0.170` in the lockfile) and `pnpm dedupe` dropped a stale second `@tanstack/ai@1.1.0` that `@tanstack/ai-client@1.1.0` (`^1.1.0`) kept in the lockfile. eve stays on 0.91.0: no stable release since, npm `latest` is `1.0.0-rc.10` (not adopted; eve 1.0 adoption is PR #104); published `@upstash/agentkit-eve-extension@0.19.0` + `@upstash/agentkit-eve@0.19.0` install on it, `eve build` clean and the mocked `eve eval` passes 4/4. A `skipLibCheck: false` consumer of `ai` errors inside `@ai-sdk/provider` (`TS7016 … module 'json-schema'`, no `@types/json-schema` dependency) on 7.0.163 through 7.0.170 alike: upstream, fixed by `npm i -D @types/json-schema`. The free throwaway Redis caps search indexes at 1: run index-using test files one at a time with `FLUSHDB` in between (`FLUSHDB` drops the indexes too; `SEARCH.DROPINDEX` is not available). - **ai 7.0.163 (2026-11-05):** pins bumped from 7.0.159 (no code change); `examples/ai-sdk-demo`'s `@ai-sdk/react` raised to `^4.0.167` because 4.0.160 depends on an *exact* `ai@7.0.156`, which split the `UIMessage` type from the demo's `ai@7.0.163` and failed `next build` (`SteeringMessageUIPart`). Re-checked 2026-11-05: published `@upstash/agentkit-eve-extension@0.19.0` / `@upstash/agentkit-eve@0.19.0` build clean on `eve@1.0.0-rc.5` (the 0.91.0 build stays); TanStack AI 1.x (ai 1.1.0, ai-persistence 1.0.1, ai-memory 1.1.0) still not adopted (same `threadId`/`blobKey` test-file errors). @@ -1731,7 +1718,8 @@ find node_modules -path "*/zod/package.json" | while read f; do echo "$f $(node - **Sandbox + MCP stores (2026-11-11, `@upstash/agentkit-tanstack-ai` minor):** `./sandbox` = `upstashSandboxCheckpointStore` (`SandboxCheckpointStore` + `forkFromCheckpoint`) and `upstashSandboxInstanceStore`; `./mcp` = `upstashMcpTaskStore` (`TaskStore` for `createMCPServer`) and - `upstashMcpEventStore` (`MCPEventStore`, MCP Events direct mode). Optional peers `@tanstack/ai-sandbox` + `upstashMcpEventStore` (`MCPEventStore`, MCP Events direct mode) and `upstashMcpSessionStore` (below). + Optional peers `@tanstack/ai-sandbox` / `@tanstack/ai-mcp` `>=1.0.0` (checkpoint/instance `.d.ts` and `MCPEventSubscription` are identical from 1.0.0 to 1.2.x). Gotchas: (1) writer leases are `"||"` strings compared against the CALLER's clock in Lua (ARGV), not a Redis TTL — the conformance suite injects a fake `now`; the PX is @@ -1748,5 +1736,18 @@ find node_modules -path "*/zod/package.json" | while read f; do echo "$f $(node script is not rolled back on error. (6) `upstashPersistence` gained an overload: with `bucket` the return type has a required `blobs`, which `createSandboxSnapshots` needs (`src/sandbox/wiring.test.ts` pins it). Also runs `runDurableRunFieldsConformance` on the run store. -- **TanStack AI 1.0 contracts implemented (2026-11-05, `@upstash/agentkit-tanstack-ai` minor):** `stores.log` (`persistence/log-store.ts`, see below), `stores.sessions` (`SessionIndexStore`, `persistence/sessions.ts`) and `./byok` (`upstashCredentials`, `byok/credentials.ts`, `CredentialStore`). **Peer floors raised to `@tanstack/ai >=1.0.0` and `@tanstack/ai-persistence >=1.0.0`** (verified by packing 0.71.0/0.10.1: neither has `SessionIndexStore` / `@tanstack/ai/byok`; `ai-memory` floor untouched). Sessions: doc `{r,e}` per `(tenant, thread)` + one ZSET per tenant (score `updatedAt`, member `~` so the client never JSON-parses a numeric-looking id); `list` walks the ZSET in batches of whole score groups and re-reads a tie group in full so the `threadId` UTF-8 order is exact, filtering after loading; `ZRANGE BYSCORE REV` takes `(max, min)` (cast through `number`, the SDK types forbid `(`-exclusive strings). Revisions come from an `INCR` counter, never reused (the credential contract requires that; sessions may reuse). Passes the 19 `sessions` cases of `runPersistenceConformance` (they run only when the store is present). Credentials: one hash per owner (`%none` segment = undefined tenant/user, distinct from `""`), value `{r,t,x?,c | iv,ct}` with type/expiry in the clear so `list` never decrypts; optional AES-256-GCM `encryptionKey`, AAD = owner key + id (a copied ciphertext does not decrypt). `storedProviderKeys` lives in `@tanstack/ai/byok/server`, not `/byok`. **`stores.log` (`LogStore`):** per log a meta hash (head, first, epoch, owner, harness, exp, rel, wake, aff, del, laneH, oseq), a records ZSET (score seq, member `:`), snapshot string (`seq` must stay the first JSON key: `truncate` reads it with a Lua pattern), offers hash + info hash + pending ZSET, and a tick counter; shared `logDue` ZSET (when each log needs an owner) and `logHosts` hash. Every op is one Lua script over the same nine declared keys; the store clock is Redis `TIME` inside the script; numbers written from Lua go through `%.0f` (plain `tostring` prints epoch ms in scientific notation); ids returned from Lua are `~`-prefixed because the client JSON-parses numeric-looking strings (`"007"` -> 7). `subscribe` = immediate local callback + a `reported` and an `observed` tick per log so each counter value is reported once (a baseline read that already contains a local write must not mark it reported: that swallowed the notification), plus a 250 ms tick poll for other machines (non-overlapping; the first successful read is the baseline, so a failed first read is retried instead of silencing the subscription). `changed()` skips a counter the poll already reported (deterministic repro in tests: delay the append reply with a Proxy), which fixed a double notification (a poll can see the tick before the write reply arrives; it failed the `calls a listener after each append` conformance case in CI). The baseline read is still async, so a write by another machine landing in the first round trip after `subscribe` can go unreported (listeners are wake-ups, hosts reconcile with `read`). `append` rejects more than 10,000 records (a 50,000-record batch outran the Lua time limit and left part of the writes). Known follow-ups: `listLogsWithoutOwner` re-checks entries skipped by `harnesses` / a live pin on every call; consumed offers, old hosts and `deleteLog` tombstones are never cleaned up; `read` / `listOfferedInputs` without a limit return everything. Conformance needs the `twin` option (a second store set on the same keyspace) for the 4 `log across instances` cases; all 56 log cases and 19 sessions cases pass. Not implemented: `ownership` TTL-based cleanup of finished logs (logs and their keys live until `deleteLog`). + **`upstashMcpSessionStore` (`McpSessionStore`, MCP Apps):** + - **What can't be stored:** its `McpServerDescriptor` carries functions (`toolFilter`, `needsApproval`, + `toolName`) and transport objects (`fetch`, `authProvider`). + - **`set`:** stores only JSON-safe fields, checked by prototype and recursion in `isJsonSafe`, and + records the dropped paths in `m`. With a key, the whole `{d, m}` is sealed with AES-GCM + (`src/secret-box.ts`, shared with BYOK; the AAD is prefix + thread + server), so a Redis writer + cannot change the url or strip `m`. + - **`get`:** calls the user's `resolve({ threadId, serverId, descriptor, missing })`, then FAILS CLOSED, + returning `null` if any recorded path is still missing. Upstream's `createMcpAppCallHandler` then + falls back to its registry (`descriptor.x ?? fromRegistry?.x`, per field) or refuses. + - **Why fail closed:** a per-thread server that lost `needsApproval` would let a widget run a gated tool. + - **Test:** `src/mcp/session-store.test.ts` drives the real handler against a real `createMCPServer` + over node:http that requires the stored Authorization header. +- **TanStack AI 1.0 contracts implemented (2026-11-05, `@upstash/agentkit-tanstack-ai` minor):** `stores.log` (`persistence/log-store.ts`, see below), `stores.sessions` (`SessionIndexStore`, `persistence/sessions.ts`) and `./byok` (`upstashCredentials`, `byok/credentials.ts`, `CredentialStore`). **Peer floors raised to `@tanstack/ai >=1.0.0` and `@tanstack/ai-persistence >=1.0.0`** (verified by packing 0.71.0/0.10.1: neither has `SessionIndexStore` / `@tanstack/ai/byok`; `ai-memory` floor untouched). Sessions: doc `{r,e}` per `(tenant, thread)` + one ZSET per tenant (score `updatedAt`, member `~` so the client never JSON-parses a numeric-looking id); `list` walks the ZSET in batches of whole score groups and re-reads a tie group in full so the `threadId` UTF-8 order is exact, filtering after loading; `ZRANGE BYSCORE REV` takes `(max, min)` (cast through `number`, the SDK types forbid `(`-exclusive strings). Revisions come from an `INCR` counter, never reused (the credential contract requires that; sessions may reuse). Passes the 19 `sessions` cases of `runPersistenceConformance` (they run only when the store is present). Credentials: one hash per owner (`%none` segment = undefined tenant/user, distinct from `""`), value `{r,t,x?,c | iv,ct}` with type/expiry in the clear so `list` never decrypts; optional AES-256-GCM `encryptionKey`, AAD = owner key + id (a copied ciphertext does not decrypt). `storedProviderKeys` lives in `@tanstack/ai/byok/server`, not `/byok`. **`stores.log` (`LogStore`):** per log a meta hash (head, first, epoch, owner, harness, exp, rel, wake, aff, del, laneH, oseq), a records ZSET (score seq, member `:`), snapshot string (`seq` must stay the first JSON key: `truncate` reads it with a Lua pattern), offers hash + info hash + pending ZSET, and a tick counter; shared `logDue` ZSET (when each log needs an owner) and `logHosts` hash. Every op is one Lua script over the same nine declared keys; the store clock is Redis `TIME` inside the script; numbers written from Lua go through `%.0f` (plain `tostring` prints epoch ms in scientific notation); ids returned from Lua are `~`-prefixed because the client JSON-parses numeric-looking strings (`"007"` -> 7). `subscribe` = immediate local callback + a `reported` and an `observed` tick per log so each counter value is reported once (a baseline read that already contains a local write must not mark it reported: that swallowed the notification), plus a 250 ms tick poll for other machines (non-overlapping; a watch starts from this store's own last write to the log, the `written` map, not from its first read, so another machine's write that lands before the first poll is still reported; a failed read is retried on the next tick). `changed()` skips a counter the poll already reported (deterministic repro in tests: delay the append reply with a Proxy), which fixed a double notification (a poll can see the tick before the write reply arrives; it failed the `calls a listener after each append` conformance case in CI). `append` rejects more than 10,000 records (a 50,000-record batch outran the Lua time limit and left part of the writes). Known follow-ups: `listLogsWithoutOwner` re-checks entries skipped by `harnesses` / a live pin on every call; consumed offers, old hosts and `deleteLog` tombstones are never cleaned up; `read` / `listOfferedInputs` without a limit return everything. Conformance needs the `twin` option (a second store set on the same keyspace) for the 4 `log across instances` cases; all 56 log cases and 19 sessions cases pass. Not implemented: `ownership` TTL-based cleanup of finished logs (logs and their keys live until `deleteLog`). - **eve 0.91.0 + ai 7.0.159 (2026-10-31):** `eve ^0.91.0` / `ai 7.0.159` everywhere; the eve-extension dist is rebuilt on 0.91.0 (manifest `builtWithEve: 0.91.0`) and its `eve` peer floor is `>=0.91.0 || >=1.0.0-0` — verified by packing the 0.91.0-built dist into consumers: eve 0.90.0 refuses it (`was built with eve 0.91.0. Upgrade eve to 0.91.0 or later`), 0.91.0 and `1.0.0-rc.1` build clean. Both eve peers carry `|| >=1.0.0-0` because semver never matches a prerelease against `>=0.x`: `npm install eve@latest` (`1.0.0-rc.1`) failed ERESOLVE against the published `@upstash/agentkit-eve@0.18.0` (`peer eve@">=0.87.0"`). **Not adopted:** `eve@1.0.0-rc.1` (a dist built on it stamps `builtWithEve: 1.0.0-rc.1` and would refuse every 0.x eve, so build on stable) and **TanStack AI 1.0** (`@tanstack/ai` 1.0.1, `-persistence` 1.0.1, `-memory` 1.0.0): `chat()` now requires `threadId` and `ArtifactRecord` requires `blobKey`, so `packages/tanstack-ai` test files fail `tsc` and the middleware tests throw `chat() needs a threadId` (source + build are fine); pins stay on 0.71.0 / 0.10.1 / 0.2.16. diff --git a/packages/tanstack-ai/README.md b/packages/tanstack-ai/README.md index e9dc173d..bf041e6e 100644 --- a/packages/tanstack-ai/README.md +++ b/packages/tanstack-ai/README.md @@ -323,7 +323,7 @@ A server descriptor carries a policy for widget calls (`toolFilter`, `needsAppro often a transport `fetch` or OAuth `authProvider`. None of those survive JSON, so `set` splits each descriptor: -- **The JSON part goes to Redis:** the transport, `prefix`, `requestOptions`, and an array `toolFilter`. +- **The JSON part goes to Redis:** the transport, `prefix`, `requestOptions`, `clientOptions`, and an array `toolFilter`. One hash per thread, one field per server, encrypted whole when you pass `encryptionKey`. - **Everything else is dropped and recorded by name.** `resolve` supplies it again on `get`, from your code. It receives `{ threadId, serverId, descriptor, missing }`. diff --git a/packages/tanstack-ai/src/mcp/event-store.ts b/packages/tanstack-ai/src/mcp/event-store.ts index 8a03af2a..a590cc6e 100644 --- a/packages/tanstack-ai/src/mcp/event-store.ts +++ b/packages/tanstack-ai/src/mcp/event-store.ts @@ -48,7 +48,7 @@ interface Stored { r: number; } -/** Subscriptions per Redis round trip while `list` walks the store. */ +/** `SSCAN` page size while `list` walks the store. */ const LIST_BATCH = 100; /** diff --git a/packages/tanstack-ai/src/mcp/session-store.ts b/packages/tanstack-ai/src/mcp/session-store.ts index 6153ba10..91724b6e 100644 --- a/packages/tanstack-ai/src/mcp/session-store.ts +++ b/packages/tanstack-ai/src/mcp/session-store.ts @@ -21,7 +21,7 @@ export interface McpSessionResolveContext { threadId: string; /** The server's key in the map passed to `set` (its `prefix` for `createMcpAppCallHandler`). */ serverId: string; - /** The stored descriptor: JSON fields and decrypted headers/env, without the code-only fields. */ + /** The stored descriptor (decrypted), without the code-only fields. */ descriptor: McpServerDescriptor; /** Code-only fields the descriptor had when it was saved, e.g. `["needsApproval", "transport.fetch"]`. */ missing: readonly string[]; @@ -46,10 +46,9 @@ export interface UpstashMcpSessionStoreConfig { context: McpSessionResolveContext, ) => McpSessionOverrides | undefined | Promise; /** - * Encrypt the transport's `headers` (often `Authorization`) and a stdio transport's `env` at rest - * with AES-256-GCM, bound to the prefix, thread and server. A 32-byte key as a `Uint8Array`, or a - * base64 / base64url / hex string (`openssl rand -base64 32`). Strongly recommended when headers - * carry credentials: without it they are stored as plain JSON. + * Encrypt each stored entry at rest with AES-256-GCM, bound to the prefix, thread and server. A + * 32-byte key as a `Uint8Array`, or a base64 / base64url / hex string (`openssl rand -base64 32`). + * Strongly recommended when headers carry credentials: without it entries are stored as plain JSON. */ encryptionKey?: string | Uint8Array; /** diff --git a/packages/tanstack-ai/src/persistence/persistence.test.ts b/packages/tanstack-ai/src/persistence/persistence.test.ts index 5fba79ac..b55de9c9 100644 --- a/packages/tanstack-ai/src/persistence/persistence.test.ts +++ b/packages/tanstack-ai/src/persistence/persistence.test.ts @@ -378,16 +378,16 @@ describe.skipIf(!hasRedisCreds)("key-locking scripts (live Redis)", () => { }).stores; let calls = 0; const stop = log.subscribe({ logId: "once", listener: () => void (calls += 1) }); - await new Promise((r) => setTimeout(r, 60)); // the baseline read is done + await new Promise((r) => setTimeout(r, 60)); // the first poll is done await log.append({ logId: "once", seq: 1, epoch: 0, records: [{ type: "one" }] }); await new Promise((r) => setTimeout(r, 100)); stop(); expect(calls).toBe(1); }); - it("log: a baseline read that already includes a local write does not swallow its notification", async () => { - // The baseline read is issued 100 ms late, after the append ran, and the append's reply is held - // back, so the baseline sees the new counter before the append returns. + it("log: a first poll that already includes a local write does not swallow its notification", async () => { + // The first poll is issued 100 ms late, after the append ran, and the append's reply is held + // back, so the poll sees the new counter before the append returns. const { log } = upstashPersistence({ redis: slowed({ get: { before: 100 }, eval: { after: 300 } }), prefix: `${prefix}:absorb`, @@ -401,6 +401,22 @@ describe.skipIf(!hasRedisCreds)("key-locking scripts (live Redis)", () => { expect(calls).toBe(1); }); + it("log: reports another machine's write that lands before the first poll", async () => { + // The subscriber's first poll is held back 200 ms, so the other store's append lands first. + const logPrefix = `${prefix}:early`; + const subscriber = upstashPersistence({ + redis: slowed({ get: { before: 200 } }), + prefix: logPrefix, + logPollIntervalMs: 50, + }).stores.log; + const writer = upstashPersistence({ redis, prefix: logPrefix }).stores.log; + let calls = 0; + const stop = subscriber.subscribe({ logId: "early", listener: () => void (calls += 1) }); + await writer.append({ logId: "early", seq: 1, epoch: 0, records: [{ type: "x" }] }); + await expect.poll(() => calls, { timeout: 2000 }).toBe(1); + stop(); + }); + it("log: a failed first poll does not silence a subscription", async () => { const logPrefix = `${prefix}:blip`; let failed = false; @@ -423,7 +439,7 @@ describe.skipIf(!hasRedisCreds)("key-locking scripts (live Redis)", () => { const b = upstashPersistence({ redis, prefix: logPrefix }).stores.log; let calls = 0; const stop = a.subscribe({ logId: "blip", listener: () => void (calls += 1) }); - await new Promise((r) => setTimeout(r, 120)); // the retry baselines + await new Promise((r) => setTimeout(r, 120)); // the retry has run await b.append({ logId: "blip", seq: 1, epoch: 0, records: [{ type: "x" }] }); await expect.poll(() => calls, { timeout: 2000 }).toBe(1); stop();