diff --git a/.changeset/ai-compaction.md b/.changeset/ai-compaction.md new file mode 100644 index 0000000000..80a443e750 --- /dev/null +++ b/.changeset/ai-compaction.md @@ -0,0 +1,13 @@ +--- +'@tanstack/ai-compaction': minor +--- + +Add `@tanstack/ai-compaction` — context-window compaction as a `chat()` +middleware. `withCompaction({ maxTokens, strategy })` runs a pluggable +`CompactionStrategy` before each model call, so compaction is incremental and +rolling. Three strategies ship built in: `evictOldest` (drop old messages, the +default), `summarizeOldest` (replace them with an LLM summary), and +`clearToolResults` (stub old tool output, keep the messages). Combine them with +`composeStrategies`, which escalates through strategies until the transcript is +back under budget. Strategies preserve tool-call/result pairing and never touch +the system prompt. diff --git a/.changeset/compaction-persistence-integration.md b/.changeset/compaction-persistence-integration.md new file mode 100644 index 0000000000..3721134e51 --- /dev/null +++ b/.changeset/compaction-persistence-integration.md @@ -0,0 +1,8 @@ +--- +'@tanstack/ai': minor +'@tanstack/ai-compaction': minor +'@tanstack/ai-persistence': patch +--- + +Keep canonical chat history separate from compacted provider context. Reuse +validated compaction checkpoints through an optional persistence metadata store. diff --git a/docs/advanced/compaction.md b/docs/advanced/compaction.md new file mode 100644 index 0000000000..25508b5c05 --- /dev/null +++ b/docs/advanced/compaction.md @@ -0,0 +1,220 @@ +--- +title: Compaction +id: compaction +order: 3 +description: "Keep long chats under the context limit with @tanstack/ai-compaction. withCompaction runs a pluggable strategy before each model call: evict, summarize, or clear old tool output." +keywords: + - tanstack ai + - compaction + - context window + - middleware + - token limit + - summarize history +--- + +A long chat or a multi-step agent loop keeps adding messages. At some point the transcript passes the model's context limit and the call fails. You want the conversation to keep working without hitting that wall. + +`withCompaction` shrinks provider context before each model call. When the context passes `maxTokens`, a **strategy** rewrites what the model sees. The canonical transcript does not change. Add this [`ChatMiddleware`](./middleware) to the `middleware` array of any `chat()` call. + +## Install + +```bash +pnpm add @tanstack/ai-compaction +``` + +## Quick start + +The default strategy drops the oldest messages once the transcript passes `maxTokens` and keeps the recent ones. + +```typescript +import { chat, toServerSentEventsResponse } from "@tanstack/ai"; +import { openaiText } from "@tanstack/ai-openai"; +import { withCompaction } from "@tanstack/ai-compaction"; + +export async function POST(request: Request) { + const { messages } = await request.json(); + + const stream = chat({ + adapter: openaiText("gpt-5.5"), + messages, + middleware: [withCompaction({ maxTokens: 100_000 })], + }); + + return toServerSentEventsResponse(stream); +} +``` + +## Pick a strategy + +Pass `strategy` to change how the history shrinks. Three are built in. + +| Strategy | What it does | Cost | +|----------|--------------|------| +| `evictOldest` (default) | Drop the oldest messages, leave a marker | No extra model call | +| `summarizeOldest` | Replace the oldest messages with an LLM summary | One summarize call | +| `clearToolResults` | Stub the content of old tool results, keep the messages | No extra model call | + +### evictOldest + +Cheapest. Keeps the recent tail, drops the older head, and leaves a short marker in its place. This is the default, so you only name it to tune `keepRecentTokens`. + +```typescript +import { withCompaction, evictOldest } from "@tanstack/ai-compaction"; + +withCompaction({ + maxTokens: 100_000, + strategy: evictOldest({ keepRecentTokens: 40_000 }), +}); +``` + +### summarizeOldest + +Keeps the gist of old turns instead of dropping them, at the cost of one summarization call. Pass a `summarize` callback. It gets the messages about to be dropped and returns the summary text. Wire it to `summarize()` or any model call. + +```typescript +import { chat, summarize, toServerSentEventsResponse } from "@tanstack/ai"; +import { openaiText, openaiSummarize } from "@tanstack/ai-openai"; +import { withCompaction, summarizeOldest } from "@tanstack/ai-compaction"; +import type { ModelMessage } from "@tanstack/ai"; + +async function summarizeHistory(messages: Array): Promise { + const text = messages + .map((m) => `${m.role}: ${typeof m.content === "string" ? m.content : ""}`) + .join("\n"); + + const { summary } = await summarize({ + adapter: openaiSummarize("gpt-5.5"), + text, + }); + return summary; +} + +export async function POST(request: Request) { + const { messages } = await request.json(); + + const stream = chat({ + adapter: openaiText("gpt-5.5"), + messages, + middleware: [ + withCompaction({ + maxTokens: 100_000, + strategy: summarizeOldest({ summarize: summarizeHistory }), + }), + ], + }); + + return toServerSentEventsResponse(stream); +} +``` + +### clearToolResults + +Best for agent loops. Tool output (file reads, command output) is usually most of the tokens. This strategy replaces the content of old tool results with a stub and keeps every message and its tool-call pairing in place. The conversation shape does not change. + +```typescript +import { withCompaction, clearToolResults } from "@tanstack/ai-compaction"; + +withCompaction({ + maxTokens: 100_000, + // Keep the 5 most recent tool results in full, stub the older ones. + strategy: clearToolResults({ keepRecentToolResults: 5 }), +}); +``` + +### Write your own + +A strategy is a function. It gets the messages and the budget, and returns the rewritten messages, or `null` to change nothing. It runs only when the estimate is over `maxTokens`. + +```typescript +import { withCompaction } from "@tanstack/ai-compaction"; +import type { CompactionStrategy } from "@tanstack/ai-compaction"; + +// Keep only the last message. +const keepLastOnly: CompactionStrategy = (messages) => { + if (messages.length <= 1) return null; + return messages.slice(-1); +}; + +withCompaction({ + maxTokens: 100_000, + strategy: keepLastOnly, + strategyKey: "keep-last-v1", +}); +``` + +Set `strategyKey` when you combine a custom strategy with persistence. Change +the key when the strategy can produce different output. This prevents an old +checkpoint from using stale behavior. + +## Combine strategies + +`composeStrategies` runs several strategies in order and **escalates**: it stops as soon as the result is back under `maxTokens`. Put the cheap, targeted strategy first and a broad fallback last. Here it clears old tool output first, and only drops old messages if that was not enough. + +```typescript +import { + withCompaction, + composeStrategies, + clearToolResults, + evictOldest, +} from "@tanstack/ai-compaction"; + +withCompaction({ + maxTokens: 100_000, + strategy: composeStrategies(clearToolResults(), evictOldest()), +}); +``` + +## Options + +### withCompaction + +| Option | Type | Default | Description | +|--------|------|---------|-------------| +| `maxTokens` | `number` | - | **Required.** Compact when the estimated tokens across `messages` pass this. | +| `strategy` | `CompactionStrategy` | `evictOldest()` | How to shrink the messages. | +| `estimateTokens` | `(message: ModelMessage) => number` | characters / 4 | Per-message token estimate. Pass a real tokenizer if you need exact counts. | +| `strategyKey` | `string` | built-in strategy identity | Stable checkpoint identity. Set it for custom strategies, custom estimators, `summarizeOldest`, or a custom eviction marker. | +| `onCompact` | `(info: CompactionInfo) => void` | - | Runs after each compaction. `info` is `{ before, after, messagesBefore, messagesAfter }` (token and message counts). | + +### Strategy options + +| Strategy | Options | +|----------|---------| +| `evictOldest` | `keepRecentTokens` (default `maxTokens / 2`), `marker` | +| `summarizeOldest` | `summarize` (**required**), `keepRecentTokens`, `summaryRole` | +| `clearToolResults` | `keepRecentToolResults` (default `3`), `stub` | + +The token count is a rough `characters / 4` estimate. It is good enough to trigger on, not exact. Pass `estimateTokens` for provider-accurate counts. + +## What it keeps safe + +- **The system prompt is never dropped.** `chat()` keeps it separate from `messages`, so compaction only touches the conversation. +- **Tool calls stay paired with their results.** The built-in strategies never leave an orphaned tool result, so the request stays valid. +- **It runs before every model call.** Compaction is incremental: as the chat keeps growing it compacts again, and a later `summarizeOldest` pass folds an earlier summary into the new one. +- **The canonical transcript stays complete.** Compaction writes provider-only context. Persistence and other middleware still read `ctx.messages`. + +## Compaction and persistence + +Compaction and server-side [`withPersistence`](../persistence/chat-persistence) +use two message views: + +- `messages` is the complete canonical transcript. Persistence saves this view. +- `providerMessages` is temporary model context. Compaction rewrites this view. + +Middleware order does not change this split. Dropped, summarized, and stubbed +content remains in the message store. + +If the persistence adapter has a `metadata` store, compaction also saves a small +checkpoint. The next request validates the canonical prefix, restores the last +compacted result, and adds only new messages. A changed prefix or strategy key +invalidates the checkpoint. + +The default strategy, standard `evictOldest`, `clearToolResults`, and safe +compositions get a strategy key automatically. Set `strategyKey` for +`summarizeOldest`, custom strategies, custom estimators, or custom marker +functions. Without a metadata store or safe key, compaction stays stateless. + +## Next steps + +- [Middleware](./middleware): the full hook reference and how middleware composes +- [Built-in Middleware](./built-in-middleware): ready-made middleware that ships in `@tanstack/ai` diff --git a/docs/advanced/middleware.md b/docs/advanced/middleware.md index fd9d064436..70b720e005 100644 --- a/docs/advanced/middleware.md +++ b/docs/advanced/middleware.md @@ -157,7 +157,8 @@ const dynamicTemperature: ChatMiddleware = { | Field | Type | Description | |-------|------|-------------| -| `messages` | `ModelMessage[]` | Conversation history | +| `messages` | `ModelMessage[]` | Canonical conversation history. Persistence and `ctx.messages` use this field. | +| `providerMessages` | `ModelMessage[]` | Temporary context sent to the provider. Defaults to `messages`. | | `systemPrompts` | `string[]` | System prompts | | `tools` | `Tool[]` | Available tools | | `metadata` | `Record` | Request metadata | @@ -165,6 +166,10 @@ const dynamicTemperature: ChatMiddleware = { When multiple middleware define `onConfig`, the config is **piped** through them in order — each receives the merged config from the previous middleware. +Return `providerMessages` when a transform must affect only the model call. For +compatibility, returning `messages` also updates provider input unless the same +result sets `providerMessages` explicitly. + ### onStructuredOutputConfig Called once at the start of the final structured-output adapter call — only when `chat()` was invoked with `outputSchema` **and** `supportsCombinedToolsAndSchema()` does not return `true` for the current model/options. Pipes through middleware in order, like `onConfig`, but with access to the **JSON Schema** being sent to the provider. Use this hook when you need to transform the schema (e.g., inject `$defs`, strip vendor-incompatible keywords) or apply structured-output-specific behavior (e.g., suppress system prompts on the final call). @@ -195,7 +200,8 @@ const injectDefs: ChatMiddleware = { | Field | Type | Description | |-------|------|-------------| -| `messages` | `ModelMessage[]` | Conversation history sent to the final call | +| `messages` | `ModelMessage[]` | Canonical conversation history | +| `providerMessages` | `ModelMessage[]` | Temporary context sent to the final call | | `systemPrompts` | `SystemPrompt[]` | System prompts on the final call | | `metadata` | `Record` | Request metadata | | `modelOptions` | `Record` | Provider-native options — this is where sampling params (`temperature`, `top_p` / `topP`, the provider's `max*Tokens` key) now live, alongside every other model-specific knob. See [Moving Sampling Options into modelOptions](../migration/sampling-options-to-model-options). | @@ -1140,6 +1146,7 @@ import type { ## Next Steps - [Built-in Middleware](./built-in-middleware) — `toolCacheMiddleware`, `contentGuardMiddleware`, `otelMiddleware` +- [Compaction](./compaction): keep long chats under the context limit with `withCompaction` - [OpenTelemetry](./otel) — emit traces and metrics via `otelMiddleware`- [Tools](../tools/tools) — Learn about the isomorphic tool system - [Agentic Cycle](../chat/agentic-cycle) — Understand the multi-step agent loop - [Streaming](../chat/streaming) — How streaming works in TanStack AI diff --git a/docs/config.json b/docs/config.json index 18ee148510..ea12c9d6c1 100644 --- a/docs/config.json +++ b/docs/config.json @@ -278,7 +278,7 @@ "label": "Chat Persistence", "to": "persistence/chat-persistence", "addedAt": "2026-08-04", - "updatedAt": "2026-08-19" + "updatedAt": "2026-08-26" }, { "label": "Client Persistence", @@ -348,7 +348,7 @@ "label": "Store Reference", "to": "persistence/store-reference", "addedAt": "2026-08-04", - "updatedAt": "2026-08-14" + "updatedAt": "2026-08-26" }, { "label": "How Persistence Works", @@ -537,7 +537,7 @@ "label": "Middleware", "to": "advanced/middleware", "addedAt": "2026-04-15", - "updatedAt": "2026-08-19" + "updatedAt": "2026-08-26" }, { "label": "Built-in Middleware", @@ -545,6 +545,12 @@ "addedAt": "2026-06-03", "updatedAt": "2026-07-21" }, + { + "label": "Compaction", + "to": "advanced/compaction", + "addedAt": "2026-08-24", + "updatedAt": "2026-08-26" + }, { "label": "Locks", "to": "advanced/locks", diff --git a/docs/persistence/chat-persistence.md b/docs/persistence/chat-persistence.md index 4496155c20..088b754c16 100644 --- a/docs/persistence/chat-persistence.md +++ b/docs/persistence/chat-persistence.md @@ -91,6 +91,18 @@ generation hooks. [How persistence works](./internals) has the rest. middleware loads the stored transcript and the run picks up from there, so the client does not have to re-send history. +## Compaction keeps the transcript complete + +Do you add [`withCompaction`](../advanced/compaction) to the same `chat()`? The +saved thread remains canonical. Compaction changes only the provider context, +not `ctx.messages`. The message store keeps dropped content, summaries do not +replace old turns, and cleared tool output remains available for reloads. + +If your adapter provides `stores.metadata`, `withPersistence` exposes it to +other middleware. Compaction uses it automatically for validated checkpoints. +See +[Compaction and persistence](../advanced/compaction#compaction-and-persistence). + ## What gets persisted, and when `withPersistence` writes at **four** moments so a reload never loses a turn: diff --git a/docs/persistence/store-reference.md b/docs/persistence/store-reference.md index 31604ea730..6143eb2f72 100644 --- a/docs/persistence/store-reference.md +++ b/docs/persistence/store-reference.md @@ -269,6 +269,12 @@ composite identity. A stored `null` is indistinguishable from absence at the typ level, so wrap a value you must persist as `null` (e.g. `{ value: null }`), or reject nullish values outright the way the SQLite store above does. +`withPersistence` also provides this store through the core +`MetadataCapability`. Middleware can use it for derived state without depending +on `@tanstack/ai-persistence`. For example, `withCompaction` stores validated +context checkpoints here. Do not place the canonical transcript in metadata; +the `messages` store owns it. + ## GenerationRunStore The generation counterpart to `RunStore`. Keyed by its own `runId`, with diff --git a/packages/ai-compaction/README.md b/packages/ai-compaction/README.md new file mode 100644 index 0000000000..b8f092f15b --- /dev/null +++ b/packages/ai-compaction/README.md @@ -0,0 +1,130 @@ +# @tanstack/ai-compaction + +Context-window compaction as a `chat()` middleware. When the working message set +grows past `maxTokens`, `withCompaction` runs a pluggable **strategy** that +rewrites provider context. It runs before every model call, so compaction is +incremental and rolling. The canonical transcript and system prompt stay +unchanged. + +```bash +npm install @tanstack/ai-compaction +``` + +## Quick start + +The default strategy (`evictOldest`) drops the oldest messages and keeps the +recent ones. + +```ts +import { chat } from '@tanstack/ai' +import { withCompaction } from '@tanstack/ai-compaction' + +chat({ + adapter, + messages, + middleware: [withCompaction({ maxTokens: 100_000 })], +}) +``` + +## Strategies + +Pass `strategy` to change how the history shrinks. Three are built in. + +| Strategy | What it does | Cost | +| ----------------------- | ------------------------------------------------------- | ------------------- | +| `evictOldest` (default) | Drop the oldest messages, leave a marker | No extra model call | +| `summarizeOldest` | Replace the oldest messages with an LLM summary | One summarize call | +| `clearToolResults` | Stub the content of old tool results, keep the messages | No extra model call | + +```ts +import { + withCompaction, + evictOldest, + summarizeOldest, + clearToolResults, +} from '@tanstack/ai-compaction' + +// Tune how much recent history to keep. +withCompaction({ + maxTokens: 100_000, + strategy: evictOldest({ keepRecentTokens: 40_000 }), +}) + +// Summarize instead of dropping. `summarize` gets the messages being removed. +withCompaction({ + maxTokens: 100_000, + strategy: summarizeOldest({ summarize: (msgs) => summarizeToText(msgs) }), +}) + +// Best for agent loops: stub old tool output, keep the messages in place. +withCompaction({ + maxTokens: 100_000, + strategy: clearToolResults({ keepRecentToolResults: 5 }), +}) +``` + +### Combine them + +`composeStrategies` runs strategies in order and escalates: it stops once the +result is back under `maxTokens`. Put the cheap one first. + +```ts +import { + withCompaction, + composeStrategies, + clearToolResults, + evictOldest, +} from '@tanstack/ai-compaction' + +// Clear old tool output first; only drop old messages if that isn't enough. +withCompaction({ + maxTokens: 100_000, + strategy: composeStrategies(clearToolResults(), evictOldest()), +}) +``` + +### Write your own + +A strategy gets the messages and the budget, and returns the rewritten messages +(or `null` to change nothing). It runs only when the estimate is over +`maxTokens`. + +```ts +import type { CompactionStrategy } from '@tanstack/ai-compaction' + +const keepLastOnly: CompactionStrategy = (messages) => + messages.length <= 1 ? null : messages.slice(-1) + +withCompaction({ + maxTokens: 100_000, + strategy: keepLastOnly, + strategyKey: 'keep-last-v1', +}) +``` + +## Options + +### `withCompaction` + +| Option | Default | What it does | +| ---------------- | --------------- | ---------------------------------------------------------------------------- | +| `maxTokens` | (required) | Compact when estimated tokens exceed this. | +| `strategy` | `evictOldest()` | How to shrink the messages. | +| `estimateTokens` | chars / 4 | Per-message token estimate. Swap in a real tokenizer for accuracy. | +| `strategyKey` | built-in key | Stable checkpoint identity. Set it for custom strategies or estimators. | +| `onCompact` | — | Observe each compaction (`before`/`after`/`messagesBefore`/`messagesAfter`). | + +### Strategy options + +| Strategy | Options | +| ------------------ | --------------------------------------------------------- | +| `evictOldest` | `keepRecentTokens` (default `maxTokens / 2`), `marker` | +| `summarizeOldest` | `summarize` (required), `keepRecentTokens`, `summaryRole` | +| `clearToolResults` | `keepRecentToolResults` (default `3`), `stub` | + +The token estimate is a rough `chars / 4` heuristic, good enough to trigger on, +not exact. Pass `estimateTokens` if you need provider-accurate counts. + +When `withPersistence` provides a metadata store, compaction saves a validated +checkpoint automatically. The next request reuses the compacted prefix and adds +new canonical messages. Without metadata, compaction remains stateless. diff --git a/packages/ai-compaction/package.json b/packages/ai-compaction/package.json new file mode 100644 index 0000000000..3c41c7b4fe --- /dev/null +++ b/packages/ai-compaction/package.json @@ -0,0 +1,50 @@ +{ + "name": "@tanstack/ai-compaction", + "version": "0.0.1", + "description": "Context-window compaction middleware for TanStack AI chat()", + "author": "", + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/TanStack/ai.git", + "directory": "packages/ai-compaction" + }, + "type": "module", + "module": "./dist/esm/index.js", + "types": "./dist/esm/index.d.ts", + "exports": { + ".": { + "types": "./dist/esm/index.d.ts", + "import": "./dist/esm/index.js" + } + }, + "sideEffects": false, + "files": [ + "dist", + "src" + ], + "scripts": { + "build": "vite build", + "clean": "premove ./build ./dist", + "lint:fix": "oxlint src --type-aware --fix", + "test:build": "publint --strict", + "test:oxlint": "oxlint src --type-aware", + "test:lib": "vitest --passWithNoTests", + "test:lib:dev": "pnpm test:lib --watch", + "test:types": "tsc" + }, + "keywords": [ + "ai", + "tanstack", + "compaction", + "context", + "middleware" + ], + "peerDependencies": { + "@tanstack/ai": "workspace:^" + }, + "devDependencies": { + "@tanstack/ai": "workspace:*", + "@vitest/coverage-v8": "4.1.10" + } +} diff --git a/packages/ai-compaction/src/index.test.ts b/packages/ai-compaction/src/index.test.ts new file mode 100644 index 0000000000..c24f94565e --- /dev/null +++ b/packages/ai-compaction/src/index.test.ts @@ -0,0 +1,345 @@ +import { describe, expect, it, vi } from 'vitest' +import type { + ChatMiddlewareConfig, + ChatMiddlewareContext, + MetadataStore, + ModelMessage, + ToolCall, +} from '@tanstack/ai' +import { provideMetadata } from '@tanstack/ai' +import { + clearToolResults, + composeStrategies, + estimateMessageTokens, + evictOldest, + summarizeOldest, + withCompaction, +} from './index' + +// Minimal onConfig driver. The middleware ignores ctx, so a bare stub is fine. +// oxlint-disable-next-line eslint-js/no-restricted-syntax -- test stub; onConfig never reads ctx +const CTX = {} as unknown as ChatMiddlewareContext +function runOnConfig( + mw: ReturnType, + messages: Array, + ctx = CTX, +) { + const config: ChatMiddlewareConfig = { + messages, + systemPrompts: [], + tools: [], + } + return mw.onConfig?.(ctx, config) +} + +function checkpointContext( + store: MetadataStore, + options: { aborted?: boolean } = {}, +): ChatMiddlewareContext { + // oxlint-disable-next-line eslint-js/no-restricted-syntax -- focused hook stub; only capability identity, threadId, and signal are read + const ctx = { + threadId: 'thread-1', + signal: options.aborted ? AbortSignal.abort() : undefined, + capabilities: { markProvided: () => undefined }, + } as unknown as ChatMiddlewareContext + provideMetadata(ctx, store) + return ctx +} + +const text = (role: ModelMessage['role'], content: string): ModelMessage => ({ + role, + content, +}) +// ~40 tokens each at chars/4. +const big = (role: ModelMessage['role']) => text(role, 'x'.repeat(160)) + +const call: ToolCall = { + id: 't1', + type: 'function', + function: { name: 'f', arguments: '{}' }, +} + +describe('withCompaction', () => { + it('passes through when under the token budget', async () => { + const mw = withCompaction({ maxTokens: 1000 }) + const result = await runOnConfig(mw, [ + text('user', 'hi'), + text('assistant', 'hello'), + ]) + expect(result).toBeUndefined() + }) + + it('defaults to evictOldest', async () => { + const mw = withCompaction({ maxTokens: 100 }) + const msgs = [big('user'), big('assistant'), big('user'), big('assistant')] + const result = await runOnConfig(mw, msgs) + const out = result?.providerMessages ?? [] + expect(out[0]?.content).toContain('omitted') + expect(out[out.length - 1]).toBe(msgs[msgs.length - 1]) + }) + + it('reports before/after token and message counts via onCompact', async () => { + const onCompact = vi.fn() + const mw = withCompaction({ maxTokens: 100, onCompact }) + await runOnConfig(mw, [ + big('user'), + big('assistant'), + big('user'), + big('assistant'), + ]) + expect(onCompact).toHaveBeenCalledOnce() + const info = onCompact.mock.calls[0]?.[0] + expect(info.after).toBeLessThan(info.before) + expect(info.messagesAfter).toBeLessThan(info.messagesBefore) + }) + + it('reuses a persisted checkpoint for an unchanged canonical prefix', async () => { + const values = new Map() + const store: MetadataStore = { + get: async (namespace, key) => values.get(`${namespace}:${key}`) ?? null, + set: async (namespace, key, value) => { + values.set(`${namespace}:${key}`, value) + }, + delete: async (namespace, key) => { + values.delete(`${namespace}:${key}`) + }, + } + const summarize = vi.fn(async () => 'the gist') + const messages = [ + big('user'), + big('assistant'), + big('user'), + big('assistant'), + ] + const options = { + maxTokens: 100, + strategy: summarizeOldest({ summarize, keepRecentTokens: 50 }), + strategyKey: 'summary-v1', + } + + const first = await runOnConfig( + withCompaction(options), + messages, + checkpointContext(store), + ) + const appended = [...messages, text('user', 'new')] + const second = await runOnConfig( + withCompaction(options), + appended, + checkpointContext(store), + ) + + expect(summarize).toHaveBeenCalledOnce() + expect(first?.providerMessages?.[0]?.content).toContain('the gist') + expect(second?.providerMessages?.[0]?.content).toContain('the gist') + expect(second?.providerMessages?.at(-1)?.content).toBe('new') + expect(appended).toHaveLength(5) + }) + + it('rejects a checkpoint when the canonical prefix changes', async () => { + const values = new Map() + const store: MetadataStore = { + get: async (namespace, key) => values.get(`${namespace}:${key}`) ?? null, + set: async (namespace, key, value) => { + values.set(`${namespace}:${key}`, value) + }, + delete: async () => undefined, + } + const summarize = vi.fn(async () => 'the gist') + const options = { + maxTokens: 100, + strategy: summarizeOldest({ summarize, keepRecentTokens: 50 }), + strategyKey: 'summary-v1', + } + const messages = [ + big('user'), + big('assistant'), + big('user'), + big('assistant'), + ] + + await runOnConfig( + withCompaction(options), + messages, + checkpointContext(store), + ) + await runOnConfig( + withCompaction(options), + [text('user', 'changed'.repeat(30)), ...messages.slice(1)], + checkpointContext(store), + ) + + expect(summarize).toHaveBeenCalledTimes(2) + }) + + it('does not write a checkpoint after cancellation', async () => { + const set = vi.fn() + const store: MetadataStore = { + get: async () => null, + set, + delete: async () => undefined, + } + + await runOnConfig( + withCompaction({ maxTokens: 100 }), + [big('user'), big('assistant'), big('user'), big('assistant')], + checkpointContext(store, { aborted: true }), + ) + + expect(set).not.toHaveBeenCalled() + }) +}) + +describe('evictOldest', () => { + it('keeps the recent tail and drops the head', async () => { + const mw = withCompaction({ + maxTokens: 100, + strategy: evictOldest({ keepRecentTokens: 50 }), + }) + const msgs = [big('user'), big('assistant'), big('user'), big('assistant')] + const out = (await runOnConfig(mw, msgs))?.providerMessages ?? [] + expect(out[0]?.content).toContain('omitted') + expect(out[out.length - 1]).toBe(msgs[msgs.length - 1]) + }) + + it('never lets the tail start with an orphaned tool result', async () => { + const assistantCall: ModelMessage = { + role: 'assistant', + content: 'x'.repeat(160), + toolCalls: [call], + } + const toolResult: ModelMessage = { + role: 'tool', + content: 'x'.repeat(160), + toolCallId: 't1', + } + const msgs = [big('user'), assistantCall, toolResult, big('user')] + const mw = withCompaction({ + maxTokens: 100, + strategy: evictOldest({ keepRecentTokens: 45 }), + }) + const out = (await runOnConfig(mw, msgs))?.providerMessages ?? [] + expect(out.slice(1).some((m) => m.role === 'tool')).toBe(false) + }) +}) + +describe('summarizeOldest', () => { + it('replaces the head with a summary', async () => { + const summarize = vi.fn(async () => 'the gist') + const mw = withCompaction({ + maxTokens: 100, + strategy: summarizeOldest({ summarize, keepRecentTokens: 50 }), + }) + const result = await runOnConfig(mw, [ + big('user'), + big('assistant'), + big('user'), + big('assistant'), + ]) + expect(summarize).toHaveBeenCalledOnce() + expect(result?.providerMessages?.[0]?.content).toBe( + 'Summary of earlier conversation:\nthe gist', + ) + }) +}) + +describe('clearToolResults', () => { + const toolMsg = (id: string): ModelMessage => ({ + role: 'tool', + content: 'x'.repeat(400), + toolCallId: id, + }) + + it('stubs old tool results but keeps recent ones and message count', async () => { + const msgs: Array = [ + text('user', 'go'), + toolMsg('a'), + toolMsg('b'), + toolMsg('c'), + toolMsg('d'), + ] + const mw = withCompaction({ + maxTokens: 100, + strategy: clearToolResults({ keepRecentToolResults: 2 }), + }) + const out = (await runOnConfig(mw, msgs))?.providerMessages ?? [] + // Same number of messages — structure is untouched. + expect(out.length).toBe(msgs.length) + // Oldest two tool results are stubbed. + expect(out[1]?.content).toBe('[tool output cleared to save context]') + expect(out[2]?.content).toBe('[tool output cleared to save context]') + // Two most recent tool results are untouched. + expect(out[3]?.content).toBe('x'.repeat(400)) + expect(out[4]?.content).toBe('x'.repeat(400)) + }) + + it('no-ops when there are not enough tool results to clear', async () => { + const msgs: Array = [big('user'), toolMsg('a'), big('user')] + const mw = withCompaction({ + maxTokens: 50, + strategy: clearToolResults({ keepRecentToolResults: 3 }), + }) + expect(await runOnConfig(mw, msgs)).toBeUndefined() + }) +}) + +describe('composeStrategies', () => { + const assistantCall = (id: string): ModelMessage => ({ + role: 'assistant', + content: '', + toolCalls: [ + { id, type: 'function', function: { name: 'f', arguments: '{}' } }, + ], + }) + const toolMsg = (id: string): ModelMessage => ({ + role: 'tool', + content: 'x'.repeat(800), // ~200 tokens + toolCallId: id, + }) + const history = (): Array => [ + text('user', 'HEAD_MARKER'), + assistantCall('a'), + toolMsg('a'), + assistantCall('b'), + toolMsg('b'), + text('user', 'last'), + ] + + it('stops after the first strategy once back under budget', async () => { + const mw = withCompaction({ + maxTokens: 260, + strategy: composeStrategies( + clearToolResults({ keepRecentToolResults: 1 }), + evictOldest({ keepRecentTokens: 50 }), + ), + }) + const msgs = history() + const out = (await runOnConfig(mw, msgs))?.providerMessages ?? [] + // Clearing one tool result was enough, so evict never ran: + // the head message and full message count survive. + expect(out.length).toBe(msgs.length) + expect(out.some((m) => m.content === 'HEAD_MARKER')).toBe(true) + expect(out[2]?.content).toBe('[tool output cleared to save context]') + }) + + it('escalates to the next strategy when the first is not enough', async () => { + const mw = withCompaction({ + maxTokens: 60, + strategy: composeStrategies( + clearToolResults({ keepRecentToolResults: 1 }), + evictOldest({ keepRecentTokens: 30 }), + ), + }) + const msgs = history() + const out = (await runOnConfig(mw, msgs))?.providerMessages ?? [] + // Clearing was not enough, so evict ran too: the head is dropped. + expect(out.some((m) => m.content === 'HEAD_MARKER')).toBe(false) + expect(out[0]?.content).toContain('omitted') + }) +}) + +describe('estimateMessageTokens', () => { + it('counts content and tool calls', () => { + expect(estimateMessageTokens(text('user', 'x'.repeat(40)))).toBe(10) + }) +}) diff --git a/packages/ai-compaction/src/index.ts b/packages/ai-compaction/src/index.ts new file mode 100644 index 0000000000..b35388b2ae --- /dev/null +++ b/packages/ai-compaction/src/index.ts @@ -0,0 +1,387 @@ +/** + * `@tanstack/ai-compaction` — context-window compaction as a `chat()` + * middleware. `withCompaction({ maxTokens, strategy })` runs before each model + * call: when the working message set grows past `maxTokens`, the chosen + * `CompactionStrategy` rewrites the messages. Because it runs every call, + * compaction is incremental and rolling. + * + * Strategies are pluggable, mirroring `AgentLoopStrategy`. Three are built in: + * {@link evictOldest}, {@link summarizeOldest}, and {@link clearToolResults}. + * Write your own by passing any {@link CompactionStrategy}. + * + * The system prompt is never touched — `chat()` keeps it separate from + * `messages`. + */ +import { MetadataCapability, getMetadata } from '@tanstack/ai' +import type { ChatMiddleware, ModelMessage } from '@tanstack/ai' + +const strategyKeys = new WeakMap() +const CHECKPOINT_NAMESPACE = '@tanstack/ai-compaction' + +interface CompactionCheckpoint { + schemaVersion: 1 + sourceMessageCount: number + sourceHash: string + strategyKey: string + compactedMessages: Array +} + +function identifyStrategy( + strategy: CompactionStrategy, + key: string | undefined, +): CompactionStrategy { + if (key) strategyKeys.set(strategy, key) + return strategy +} + +async function hashMessages( + messages: ReadonlyArray, +): Promise { + const bytes = new TextEncoder().encode(JSON.stringify(messages)) + const digest = await globalThis.crypto.subtle.digest('SHA-256', bytes) + return Array.from(new Uint8Array(digest), (byte) => + byte.toString(16).padStart(2, '0'), + ).join('') +} + +function isModelMessage(value: unknown): value is ModelMessage { + return ( + typeof value === 'object' && + value !== null && + 'role' in value && + (value.role === 'user' || + value.role === 'assistant' || + value.role === 'tool') && + 'content' in value + ) +} + +function isCompactionCheckpoint(value: unknown): value is CompactionCheckpoint { + return ( + typeof value === 'object' && + value !== null && + 'schemaVersion' in value && + value.schemaVersion === 1 && + 'sourceMessageCount' in value && + typeof value.sourceMessageCount === 'number' && + Number.isInteger(value.sourceMessageCount) && + value.sourceMessageCount >= 0 && + 'sourceHash' in value && + typeof value.sourceHash === 'string' && + 'strategyKey' in value && + typeof value.strategyKey === 'string' && + 'compactedMessages' in value && + Array.isArray(value.compactedMessages) && + value.compactedMessages.every(isModelMessage) + ) +} + +/** Rough token estimate for one message. Default: characters / 4. */ +export function estimateMessageTokens(message: ModelMessage): number { + let text = + typeof message.content === 'string' + ? message.content + : JSON.stringify(message.content ?? '') + if (message.toolCalls?.length) text += JSON.stringify(message.toolCalls) + return Math.ceil(text.length / 4) +} + +/** What a {@link CompactionStrategy} receives alongside the messages. */ +export interface CompactionContext { + /** The `maxTokens` budget from `withCompaction`. */ + maxTokens: number + /** The shared token estimator (default {@link estimateMessageTokens}). */ + estimate: (message: ModelMessage) => number +} + +/** + * Shrinks a message list. Called only when the estimate is over budget. + * Return the rewritten messages, or `null` to leave them unchanged. + */ +export type CompactionStrategy = ( + messages: ReadonlyArray, + ctx: CompactionContext, +) => Array | null | Promise | null> + +/** Reported to `onCompact` after each compaction event. */ +export interface CompactionInfo { + /** Estimated tokens before compaction. */ + before: number + /** Estimated tokens after compaction. */ + after: number + /** Message count before compaction. */ + messagesBefore: number + /** Message count after compaction (unchanged for {@link clearToolResults}). */ + messagesAfter: number +} + +export interface CompactionOptions { + /** Compact when estimated tokens across `messages` exceed this. */ + maxTokens: number + /** How to shrink the messages. Default: {@link evictOldest}. */ + strategy?: CompactionStrategy + /** Per-message token estimator. Default: {@link estimateMessageTokens}. */ + estimateTokens?: (message: ModelMessage) => number + /** + * Stable identity for persisted checkpoints. Set this for custom strategies + * or estimators, and change it when their output can change. + */ + strategyKey?: string + /** Observe each compaction (logging, metrics). */ + onCompact?: (info: CompactionInfo) => void +} + +const sum = ( + messages: ReadonlyArray, + estimate: (m: ModelMessage) => number, +) => messages.reduce((total, m) => total + estimate(m), 0) + +/** + * Find the split point that keeps the most recent messages up to + * `keepRecentTokens`, then moves the cut forward past any leading tool result + * so the kept tail never starts with an orphan (its tool call would be dropped). + * Returns the index where the tail begins (head is `messages[0..cut)`). + */ +function splitAtRecent( + messages: ReadonlyArray, + estimate: (m: ModelMessage) => number, + keepRecentTokens: number, +): number { + let kept = 0 + let cut = messages.length + while (cut > 0) { + const prev = messages[cut - 1] + if (!prev) break + const size = estimate(prev) + if (kept + size > keepRecentTokens) break + kept += size + cut-- + } + // Always keep at least the last message. + if (cut >= messages.length) cut = messages.length - 1 + while (cut < messages.length && messages[cut]?.role === 'tool') cut++ + return cut +} + +/** + * Drop the oldest messages and replace them with a short marker. Cheapest + * strategy — no extra model call. This is the default. + */ +export function evictOldest( + options: { + /** Tokens of recent messages to keep verbatim. Default `floor(maxTokens/2)`. */ + keepRecentTokens?: number + /** Build the marker that replaces the dropped head. */ + marker?: (droppedCount: number) => string + } = {}, +): CompactionStrategy { + const strategy: CompactionStrategy = (messages, ctx) => { + const keep = options.keepRecentTokens ?? Math.floor(ctx.maxTokens / 2) + const cut = splitAtRecent(messages, ctx.estimate, keep) + // ponytail: can't shrink past the recent window; raise keepRecentTokens or + // lower maxTokens if compaction never fires. + if (cut <= 0) return null + const marker = + options.marker?.(cut) ?? + `[${cut} earlier message(s) omitted to save context.]` + return [{ role: 'user', content: marker }, ...messages.slice(cut)] + } + return identifyStrategy( + strategy, + options.marker + ? undefined + : `evict-oldest:${options.keepRecentTokens ?? 'half'}`, + ) +} + +/** + * Drop the oldest messages and replace them with an LLM summary. Keeps the gist + * of old turns at the cost of one summarization call. Wire `summarize` to + * `summarize()` or any model call. + */ +export function summarizeOldest(options: { + summarize: (messages: Array) => Promise + /** Tokens of recent messages to keep verbatim. Default `floor(maxTokens/2)`. */ + keepRecentTokens?: number + /** Role of the injected summary message. Default `'user'`. */ + summaryRole?: 'user' | 'assistant' +}): CompactionStrategy { + return async (messages, ctx) => { + const keep = options.keepRecentTokens ?? Math.floor(ctx.maxTokens / 2) + const cut = splitAtRecent(messages, ctx.estimate, keep) + if (cut <= 0) return null + const summary = await options.summarize(messages.slice(0, cut)) + return [ + { + role: options.summaryRole ?? 'user', + content: `Summary of earlier conversation:\n${summary}`, + }, + ...messages.slice(cut), + ] + } +} + +/** + * Replace the content of old tool-result messages with a stub, keeping every + * message and its tool-call pairing in place. Best for agent loops where tool + * output (file reads, command output) dominates the token count — it clears the + * bulk without disturbing the conversation shape. No extra model call. + */ +export function clearToolResults( + options: { + /** Number of most-recent tool results to keep verbatim. Default `3`. */ + keepRecentToolResults?: number + /** Text that replaces a cleared tool result. */ + stub?: string + } = {}, +): CompactionStrategy { + const keepN = options.keepRecentToolResults ?? 3 + const stub = options.stub ?? '[tool output cleared to save context]' + const strategy: CompactionStrategy = (messages) => { + const toolIndexes: Array = [] + messages.forEach((m, i) => { + if (m.role === 'tool') toolIndexes.push(i) + }) + if (toolIndexes.length <= keepN) return null + const clearBefore = toolIndexes[toolIndexes.length - keepN] ?? 0 + let changed = false + const next = messages.map((m, i) => { + if (m.role === 'tool' && i < clearBefore && m.content !== stub) { + changed = true + return { ...m, content: stub } + } + return m + }) + return changed ? next : null + } + return identifyStrategy(strategy, `clear-tool-results:${keepN}:${stub}`) +} + +/** + * Run several strategies in order, escalating: stop as soon as the running + * estimate is back under `maxTokens`. Put the cheap, targeted strategy first + * (for example {@link clearToolResults}) and a broad fallback last (for example + * {@link evictOldest}) — the fallback only runs when clearing was not enough. + * A strategy that returns `null` (no change) is skipped and the next one runs. + * + * @example + * ```ts + * withCompaction({ + * maxTokens: 100_000, + * strategy: composeStrategies(clearToolResults(), evictOldest()), + * }) + * ``` + */ +export function composeStrategies( + ...strategies: Array +): CompactionStrategy { + const strategy: CompactionStrategy = async (messages, ctx) => { + let current: ReadonlyArray = messages + let result: Array | null = null + for (const itemStrategy of strategies) { + if (sum(current, ctx.estimate) <= ctx.maxTokens) break + const out = await itemStrategy(current, ctx) + if (out) { + current = out + result = out + } + } + return result + } + const keys = strategies.map((item) => strategyKeys.get(item)) + return identifyStrategy( + strategy, + keys.every((key) => key !== undefined) ? keys.join('|') : undefined, + ) +} + +/** + * Context-compaction middleware. Add to `chat({ middleware: [...] })`. + * + * @example + * ```ts + * chat({ + * adapter, + * messages, + * middleware: [withCompaction({ maxTokens: 100_000 })], // evictOldest by default + * }) + * ``` + */ +export function withCompaction(options: CompactionOptions): ChatMiddleware { + const estimate = options.estimateTokens ?? estimateMessageTokens + const strategy = options.strategy ?? evictOldest() + const strategyKey = + options.strategyKey ?? + (options.estimateTokens ? undefined : strategyKeys.get(strategy)) + const checkpointStrategyKey = strategyKey + ? `${strategyKey}:maxTokens=${options.maxTokens}` + : undefined + + return { + name: 'compaction', + optionalRequires: [MetadataCapability], + async onConfig(ctx, config) { + const { messages } = config + const inputMessages = config.providerMessages ?? messages + const metadata = getMetadata(ctx, { optional: true }) + let workingMessages = inputMessages + let reusedCheckpoint = false + + if (metadata && checkpointStrategyKey && inputMessages === messages) { + const stored = await metadata.get(CHECKPOINT_NAMESPACE, ctx.threadId) + if ( + isCompactionCheckpoint(stored) && + stored.strategyKey === checkpointStrategyKey && + stored.sourceMessageCount <= messages.length && + stored.sourceHash === + (await hashMessages(messages.slice(0, stored.sourceMessageCount))) + ) { + workingMessages = [ + ...stored.compactedMessages, + ...messages.slice(stored.sourceMessageCount), + ] + reusedCheckpoint = true + } + } + + const before = sum(workingMessages, estimate) + if (before <= options.maxTokens) { + return reusedCheckpoint + ? { providerMessages: workingMessages } + : undefined + } + + const next = await strategy(workingMessages, { + maxTokens: options.maxTokens, + estimate, + }) + if (!next || next === workingMessages) { + return reusedCheckpoint + ? { providerMessages: workingMessages } + : undefined + } + + options.onCompact?.({ + before, + after: sum(next, estimate), + messagesBefore: workingMessages.length, + messagesAfter: next.length, + }) + + if (metadata && checkpointStrategyKey && inputMessages === messages) { + const checkpoint: CompactionCheckpoint = { + schemaVersion: 1, + sourceMessageCount: messages.length, + sourceHash: await hashMessages(messages), + strategyKey: checkpointStrategyKey, + compactedMessages: next, + } + if (!ctx.signal?.aborted) { + await metadata.set(CHECKPOINT_NAMESPACE, ctx.threadId, checkpoint) + } + } + + return { providerMessages: next } + }, + } +} diff --git a/packages/ai-compaction/tsconfig.json b/packages/ai-compaction/tsconfig.json new file mode 100644 index 0000000000..29112eff9f --- /dev/null +++ b/packages/ai-compaction/tsconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "outDir": "dist" + }, + "include": ["vite.config.ts", "./src", "./tests"], + "exclude": ["node_modules", "dist"] +} diff --git a/packages/ai-compaction/vite.config.ts b/packages/ai-compaction/vite.config.ts new file mode 100644 index 0000000000..1f3542380f --- /dev/null +++ b/packages/ai-compaction/vite.config.ts @@ -0,0 +1,35 @@ +import { defineConfig, mergeConfig } from 'vitest/config' +import { tanstackViteConfig } from '@tanstack/vite-config' +import packageJson from './package.json' + +const config = defineConfig({ + test: { + name: packageJson.name, + dir: './', + watch: false, + globals: true, + environment: 'node', + include: ['src/**/*.test.ts', 'tests/**/*.test.ts'], + coverage: { + provider: 'v8', + reporter: ['text', 'json', 'html', 'lcov'], + exclude: [ + 'node_modules/', + 'dist/', + 'tests/', + '**/*.test.ts', + '**/*.config.ts', + ], + include: ['src/**/*.ts'], + }, + }, +}) + +export default mergeConfig( + config, + tanstackViteConfig({ + entry: ['./src/index.ts'], + srcDir: './src', + cjs: false, + }), +) diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index 7e8eb184e2..2ccb09635a 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -3,6 +3,8 @@ import { fromSpecTokenUsage, getDetachableRun, InterruptResumeValidationError, + MetadataCapability, + provideMetadata, readInterruptBinding, validateInterruptResumeBatch, wasCancelRequested, @@ -1961,6 +1963,7 @@ export function withPersistence( const provides = [ PersistenceCapability, PersistenceCompletionCapability, + ...(persistence.stores.metadata ? [MetadataCapability] : []), ...(wantsInterrupts ? [InterruptsCapability] : []), ] @@ -1969,6 +1972,9 @@ export function withPersistence( provides, setup(ctx: ChatMiddlewareContext) { providePersistence(ctx, persistence) + if (persistence.stores.metadata) { + provideMetadata(ctx, persistence.stores.metadata) + } let resolveCompletion: () => void = () => undefined let rejectCompletion: (error: unknown) => void = () => undefined diff --git a/packages/ai-persistence/src/types.ts b/packages/ai-persistence/src/types.ts index d16e992395..084042c08c 100644 --- a/packages/ai-persistence/src/types.ts +++ b/packages/ai-persistence/src/types.ts @@ -1,5 +1,6 @@ import type { ModelMessage, + MetadataStore, PersistedArtifactRef, RunStatus, RunStore, @@ -11,7 +12,7 @@ import type { // `@tanstack/ai` or `@tanstack/ai-persistence`. See {@link Scope} security notes: // pair a client-visible `threadId` with a server-trusted `userId`/`tenantId` // before authorizing load/save (e.g. via `reconstructChat({ authorize })`). -export type { Scope } +export type { MetadataStore, Scope } // =========================================================================== // Store contracts @@ -292,37 +293,6 @@ export interface InterruptStore { listPendingByRun: (runId: string) => Promise> } -/** - * Namespaced key/value store for arbitrary JSON metadata (app-owned). - * - * The first argument is an **app-defined namespace string**, not the shared - * {@link Scope} identity type from `@tanstack/ai`. Composite identity is - * `(namespace, key)` as two independent fields (SQL backends use a composite - * primary key; the in-memory store uses nested maps). Do not encode both into a - * single delimited string — `${namespace}:${key}` collides when either part - * contains `:`. - * - * The same `key` under different namespaces is independent. - */ -export interface MetadataStore { - /** - * Return the stored value for `(namespace, key)`, or `null` if absent. - * - * CAVEAT: the return type is `unknown | null`, where `| null` collapses into - * `unknown` — a stored value of `null` is therefore **indistinguishable from - * absence** at the type level. Callers that must persist a real `null` - * distinctly from "not set" should wrap it (e.g. store `{ value: null }`). - */ - get: (namespace: string, key: string) => Promise - /** Insert or overwrite the value for `(namespace, key)`. */ - set: (namespace: string, key: string, value: unknown) => Promise - /** - * Remove `(namespace, key)`. A no-op if absent. Does not affect other - * namespaces. - */ - delete: (namespace: string, key: string) => Promise -} - // =========================================================================== // Store typers // =========================================================================== diff --git a/packages/ai-persistence/tests/metadata-capability.test.ts b/packages/ai-persistence/tests/metadata-capability.test.ts new file mode 100644 index 0000000000..c01fca8344 --- /dev/null +++ b/packages/ai-persistence/tests/metadata-capability.test.ts @@ -0,0 +1,114 @@ +import { describe, expect, it } from 'vitest' +import { + EventType, + MetadataCapability, + chat, + defineChatMiddleware, + getMetadata, +} from '@tanstack/ai' +import type { AnyTextAdapter, MetadataStore, StreamChunk } from '@tanstack/ai' +import { memoryPersistence } from '../src/memory' +import { withPersistence } from '../src/middleware' +import { defineAIPersistence, defineMessageStore } from '../src/types' + +function mockAdapter() { + return { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': { + providerOptions: undefined, + inputModalities: undefined, + messageMetadataByModality: undefined, + toolCapabilities: undefined, + toolCallMetadata: undefined, + systemPromptMetadata: undefined, + }, + chatStream: () => + (async function* () { + yield { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: 1, + } as const + yield { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop', + timestamp: 1, + } as const + })(), + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } satisfies AnyTextAdapter +} + +async function collect(stream: AsyncIterable) { + for await (const _chunk of stream) { + // Drain the stream so terminal middleware hooks run. + } +} + +describe('metadata capability', () => { + it('provides the persistence metadata store before onConfig in either order', async () => { + const persistence = memoryPersistence() + let metadata: MetadataStore | undefined + const consumer = defineChatMiddleware({ + name: 'metadata-consumer', + optionalRequires: [MetadataCapability], + onConfig(ctx) { + metadata = getMetadata(ctx, { optional: true }) + }, + }) + + await collect( + chat({ + adapter: mockAdapter(), + messages: [{ role: 'user', content: 'hello' }], + middleware: [consumer, withPersistence(persistence)], + }), + ) + + expect(metadata).toBe(persistence.stores.metadata) + }) + + it('leaves the capability absent when persistence has no metadata store', async () => { + const threads = new Map>() + const persistence = defineAIPersistence({ + stores: { + messages: defineMessageStore({ + loadThread: async (threadId) => threads.get(threadId) ?? [], + saveThread: async (threadId, messages) => { + threads.set( + threadId, + messages.filter( + (message): message is { role: 'user'; content: string } => + message.role === 'user' && + typeof message.content === 'string', + ), + ) + }, + }), + }, + }) + let metadata: MetadataStore | undefined + const consumer = defineChatMiddleware({ + name: 'metadata-consumer', + optionalRequires: [MetadataCapability], + onConfig(ctx) { + metadata = getMetadata(ctx, { optional: true }) + }, + }) + + await collect( + chat({ + adapter: mockAdapter(), + messages: [{ role: 'user', content: 'hello' }], + middleware: [withPersistence(persistence), consumer], + }), + ) + + expect(metadata).toBeUndefined() + }) +}) diff --git a/packages/ai-persistence/tests/with-persistence.test.ts b/packages/ai-persistence/tests/with-persistence.test.ts index 2e0db2451a..3aa1a54087 100644 --- a/packages/ai-persistence/tests/with-persistence.test.ts +++ b/packages/ai-persistence/tests/with-persistence.test.ts @@ -3,6 +3,7 @@ import { EventType, chat } from '@tanstack/ai' import type { AdapterYieldChunk, AnyTextAdapter, + ChatMiddleware, ModelMessage, StreamChunk, Tool, @@ -140,6 +141,70 @@ describe('withPersistence (state-only)', () => { ]) }) + it('saves canonical history when middleware compacts provider messages', async () => { + const persistence = memoryPersistence() + const { adapter, calls } = mockAdapter([ + [ev.runStarted(), ev.text('hello'), ev.runFinished()], + ]) + + const dropOldest: ChatMiddleware = { + name: 'drop-oldest', + onConfig(ctx, config) { + if (ctx.phase !== 'beforeModel' || config.messages.length <= 1) return + return { providerMessages: config.messages.slice(1) } + }, + } + + await collect( + chat({ + adapter, + messages: [ + { role: 'user', content: 'DROP_ME_FIRST' }, + { role: 'user', content: 'KEEP_ME_LAST' }, + ], + runId: 'r1', + threadId: 't1', + middleware: [dropOldest, withPersistence(persistence)], + }) as AsyncIterable, + ) + + const thread = await persistence.stores.messages!.loadThread('t1') + expect(thread).toEqual([ + { role: 'user', content: 'DROP_ME_FIRST' }, + { role: 'user', content: 'KEEP_ME_LAST' }, + expect.objectContaining({ role: 'assistant', content: 'hello' }), + ]) + expect(calls[0]).toEqual( + expect.objectContaining({ + messages: [{ role: 'user', content: 'KEEP_ME_LAST' }], + }), + ) + }) + + it('does not add ids to caller messages while saving', async () => { + const persistence = memoryPersistence() + const { adapter } = mockAdapter([ + [ev.runStarted(), ev.text('hello'), ev.runFinished()], + ]) + const userMessage: ModelMessage = { role: 'user', content: 'hello' } + + await collect( + chat({ + adapter, + messages: [userMessage], + runId: 'r1', + threadId: 't1', + middleware: [withPersistence(persistence)], + }) as AsyncIterable, + ) + + expect(userMessage).toEqual({ role: 'user', content: 'hello' }) + expect(await persistence.stores.messages!.loadThread('t1')).toEqual([ + { role: 'user', content: 'hello' }, + expect.objectContaining({ role: 'assistant', content: 'hello' }), + ]) + }) + it('persists cumulative usage across model calls', async () => { const persistence = memoryPersistence() const { adapter } = mockAdapter([ diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 91f84d3490..17132ecddb 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -789,6 +789,7 @@ class TextEngine< private readonly effectiveSignal?: AbortSignal private messages: Array + private providerMessages: Array private iterationCount = 0 /** Cumulative tool calls counted in this run (emitted + pending resume). */ private toolCallCount = 0 @@ -936,6 +937,7 @@ class TextEngine< // Convert messages to ModelMessage format (handles both UIMessage and ModelMessage input) // This ensures consistent internal format regardless of what the client sends this.messages = convertMessagesToModelMessages(config.params.messages) + this.providerMessages = this.messages // Initialize lazy tool manager after messages are converted (needs message history for scanning) assertUniqueToolNames(config.params.tools || []) @@ -1498,7 +1500,7 @@ class TextEngine< for await (const raw of this.adapter.chatStream({ model: this.params.model, - messages: this.messages, + messages: this.providerMessages, tools: toolsWithJsonSchemas, metadata, request: this.effectiveRequest, @@ -3471,7 +3473,7 @@ class TextEngine< const structuredCallOptions = { chatOptions: { model: this.params.model, - messages: this.messages, + messages: this.providerMessages, metadata: postOnConfig.metadata, modelOptions: postOnConfig.modelOptions, systemPrompts: postOnConfig.systemPrompts, @@ -3950,6 +3952,7 @@ class TextEngine< private buildMiddlewareConfig(): ChatMiddlewareConfig { return { messages: this.messages, + providerMessages: this.messages, systemPrompts: [...this.systemPrompts], tools: [...this.tools], resume: this.params.resume, @@ -4368,6 +4371,7 @@ class TextEngine< private applyMiddlewareConfig(config: ChatMiddlewareConfig): void { this.applyResumeToolState(config.resumeToolState) this.messages = config.messages + this.providerMessages = config.providerMessages ?? config.messages this.systemPrompts = config.systemPrompts assertUniqueToolNames(config.tools) this.tools = config.tools diff --git a/packages/ai/src/activities/chat/middleware/compose.ts b/packages/ai/src/activities/chat/middleware/compose.ts index 3e811cfbaa..5493e20485 100644 --- a/packages/ai/src/activities/chat/middleware/compose.ts +++ b/packages/ai/src/activities/chat/middleware/compose.ts @@ -166,7 +166,13 @@ export class MiddlewareRunner< const result = await mw.onConfig(ctx, current) const hasTransform = result !== undefined && result !== null if (hasTransform) { - current = { ...current, ...result } + current = { + ...current, + ...result, + ...('messages' in result && !('providerMessages' in result) + ? { providerMessages: result.messages } + : {}), + } if (!skip) { this.logger.config( `middleware=${mw.name ?? 'unnamed'} keys=${Object.keys(result).join(',')}`, @@ -221,7 +227,13 @@ export class MiddlewareRunner< const result = await mw.onStructuredOutputConfig(ctx, current) const hasTransform = result !== undefined && result !== null if (hasTransform) { - current = { ...current, ...result } + current = { + ...current, + ...result, + ...('messages' in result && !('providerMessages' in result) + ? { providerMessages: result.messages } + : {}), + } if (!skip) { this.logger.config( `middleware=${mw.name ?? 'unnamed'} keys=${Object.keys(result).join(',')}`, diff --git a/packages/ai/src/activities/chat/middleware/index.ts b/packages/ai/src/activities/chat/middleware/index.ts index 1f913cfa77..53da815b24 100644 --- a/packages/ai/src/activities/chat/middleware/index.ts +++ b/packages/ai/src/activities/chat/middleware/index.ts @@ -67,6 +67,9 @@ export { } from './locks' export type { LockStore } from './locks' +export { MetadataCapability, getMetadata, provideMetadata } from './metadata' +export type { MetadataStore } from './metadata' + export { isRunStatus, isTerminalRunStatus, diff --git a/packages/ai/src/activities/chat/middleware/metadata.ts b/packages/ai/src/activities/chat/middleware/metadata.ts new file mode 100644 index 0000000000..08d71987f1 --- /dev/null +++ b/packages/ai/src/activities/chat/middleware/metadata.ts @@ -0,0 +1,20 @@ +import { createCapability } from './capabilities' + +/** + * Namespaced key/value store for app and middleware metadata. + * + * `(namespace, key)` is the composite identity. Keep both values separate; + * joining them with a delimiter can create collisions. + */ +export interface MetadataStore { + /** Return the value for `(namespace, key)`, or `null` when it is absent. */ + get: (namespace: string, key: string) => Promise + /** Insert or replace the value for `(namespace, key)`. */ + set: (namespace: string, key: string, value: unknown) => Promise + /** Delete `(namespace, key)`. Do nothing when it is absent. */ + delete: (namespace: string, key: string) => Promise +} + +export const MetadataCapability = createCapability()('metadata') + +export const [getMetadata, provideMetadata] = MetadataCapability diff --git a/packages/ai/src/activities/chat/middleware/types.ts b/packages/ai/src/activities/chat/middleware/types.ts index 6cd5de10fb..f17f04d642 100644 --- a/packages/ai/src/activities/chat/middleware/types.ts +++ b/packages/ai/src/activities/chat/middleware/types.ts @@ -305,7 +305,10 @@ export interface ChatMiddlewareContext { * that middleware is allowed to modify. */ export interface ChatMiddlewareConfig { + /** Canonical conversation history. Middleware and persistence read this. */ messages: Array + /** Provider-only context. Defaults to `messages` when it is not set. */ + providerMessages?: Array | undefined systemPrompts: Array tools: Array resume?: Array | undefined diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 39b77827cf..8a300522bd 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -281,6 +281,9 @@ export { createCapability, defineChatMiddleware, createChatMiddleware, + MetadataCapability, + getMetadata, + provideMetadata, } from './activities/chat/middleware/index' export type { Capability, @@ -290,6 +293,7 @@ export type { CapabilityProvider, DefinedChatMiddleware, AnyChatMiddleware, + MetadataStore, } from './activities/chat/middleware/index' // Locks are a distributed-mutex primitive — coordination, not chat state — and // live behind their own subpath: `@tanstack/ai/locks` (see ./locks.ts). diff --git a/packages/ai/tests/provider-messages.test.ts b/packages/ai/tests/provider-messages.test.ts new file mode 100644 index 0000000000..50ef604687 --- /dev/null +++ b/packages/ai/tests/provider-messages.test.ts @@ -0,0 +1,104 @@ +import { describe, expect, it } from 'vitest' +import { chat } from '../src/activities/chat/index' +import { defineChatMiddleware } from '../src/activities/chat/middleware/define' +import { collectChunks, createMockAdapter, ev, serverTool } from './test-utils' +import type { ModelMessage, StreamChunk } from '../src/types' + +describe('provider-only messages', () => { + it('changes provider input without changing the canonical transcript', async () => { + const { adapter, calls } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.textStart(), + ev.textContent('done'), + ev.textEnd(), + ev.runFinished(), + ], + ], + }) + let finalMessages: Array = [] + + const providerFilter = defineChatMiddleware({ + name: 'provider-filter', + onConfig(ctx, config) { + if (ctx.phase !== 'beforeModel') return + return { providerMessages: config.messages.slice(1) } + }, + onFinish(ctx) { + finalMessages = [...ctx.messages] + }, + }) + + await collectChunks( + chat({ + adapter, + messages: [ + { role: 'user', content: 'DROP_FROM_PROVIDER' }, + { role: 'user', content: 'KEEP_FOR_PROVIDER' }, + ], + middleware: [providerFilter], + }) as AsyncIterable, + ) + + expect(calls[0]?.messages.map((message) => message.content)).toEqual([ + 'KEEP_FOR_PROVIDER', + ]) + expect(finalMessages.map((message) => message.content)).toEqual([ + 'DROP_FROM_PROVIDER', + 'KEEP_FOR_PROVIDER', + 'done', + ]) + }) + + it('includes new tool-loop messages in later provider calls', async () => { + const { adapter, calls } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('call-1', 'lookup'), + ev.toolArgs('call-1', '{}'), + ev.runFinished('tool_calls'), + ], + [ + ev.runStarted(), + ev.textStart(), + ev.textContent('done'), + ev.textEnd(), + ev.runFinished('stop'), + ], + ], + }) + let finalMessages: Array = [] + const providerFilter = defineChatMiddleware({ + name: 'provider-filter', + onConfig(ctx, config) { + if (ctx.phase !== 'beforeModel') return + return { providerMessages: config.messages.slice(1) } + }, + onFinish(ctx) { + finalMessages = [...ctx.messages] + }, + }) + + await collectChunks( + chat({ + adapter, + messages: [ + { role: 'user', content: 'DROP_FROM_PROVIDER' }, + { role: 'user', content: 'KEEP_FOR_PROVIDER' }, + ], + tools: [serverTool('lookup', () => ({ value: 1 }))], + middleware: [providerFilter], + }) as AsyncIterable, + ) + + expect(calls[1]?.messages.map((message) => message.role)).toEqual([ + 'user', + 'assistant', + 'tool', + ]) + expect(calls[1]?.messages[0]?.content).toBe('KEEP_FOR_PROVIDER') + expect(finalMessages[0]?.content).toBe('DROP_FROM_PROVIDER') + }) +}) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 827b407241..60c9a0a7ff 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1795,6 +1795,15 @@ importers: specifier: ^8.2.1 version: 8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0) + packages/ai-compaction: + devDependencies: + '@tanstack/ai': + specifier: workspace:* + version: link:../ai + '@vitest/coverage-v8': + specifier: 4.1.10 + version: 4.1.10(vitest@4.1.10) + packages/ai-devtools: dependencies: '@tanstack/ai': @@ -2944,6 +2953,9 @@ importers: '@tanstack/ai-client': specifier: workspace:* version: link:../../packages/ai-client + '@tanstack/ai-compaction': + specifier: workspace:* + version: link:../../packages/ai-compaction '@tanstack/ai-elevenlabs': specifier: workspace:* version: link:../../packages/ai-elevenlabs @@ -3092,6 +3104,9 @@ importers: '@tanstack/ai-client': specifier: workspace:* version: link:../../packages/ai-client + '@tanstack/ai-compaction': + specifier: workspace:* + version: link:../../packages/ai-compaction '@tanstack/ai-event-client': specifier: workspace:* version: link:../../packages/ai-event-client @@ -11786,6 +11801,7 @@ packages: eslint@9.39.4: resolution: {integrity: sha512-XoMjdBOwe/esVgEvLmNsD3IRHkm7fbKIUGvrleloJXUZgDHig2IPWNniv+GwjyJXzuNqVjlr5+4yVUZjycJwfQ==} engines: {node: ^18.18.0 || ^20.9.0 || >=21.1.0} + deprecated: This version is no longer supported. Please see https://eslint.org/version-support for other options. hasBin: true peerDependencies: jiti: '*' diff --git a/testing/e2e/package.json b/testing/e2e/package.json index 0d100c5160..7d8cb50e42 100644 --- a/testing/e2e/package.json +++ b/testing/e2e/package.json @@ -24,6 +24,7 @@ "@tanstack/ai-byteplus": "workspace:*", "@tanstack/ai-claude-code": "workspace:*", "@tanstack/ai-client": "workspace:*", + "@tanstack/ai-compaction": "workspace:*", "@tanstack/ai-elevenlabs": "workspace:*", "@tanstack/ai-gemini": "workspace:*", "@tanstack/ai-grok": "workspace:*", diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 3acd134ba7..588c2a25b1 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -79,9 +79,10 @@ import { Route as ApiEmbeddingRouteImport } from './routes/api.embedding' import { Route as ApiDurableTakeoverRouteImport } from './routes/api.durable-takeover' import { Route as ApiDurableDeliveryRouteImport } from './routes/api.durable-delivery' import { Route as ApiDevtoolsMemoryRouteImport } from './routes/api.devtools-memory' +import { Route as ApiCompactionWireRouteImport } from './routes/api.compaction-wire' import { Route as ApiChatRouteImport } from './routes/api.chat' -import { Route as ApiByokChatRouteImport } from './routes/api.byok-chat' import { Route as ApiByteplusSeedance1080pWireRouteImport } from './routes/api.byteplus-seedance-1080p-wire' +import { Route as ApiByokChatRouteImport } from './routes/api.byok-chat' import { Route as ApiAudioRouteImport } from './routes/api.audio' import { Route as ApiArktypeToolWireRouteImport } from './routes/api.arktype-tool-wire' import { Route as ApiAnthropicStructuredUsageRouteImport } from './routes/api.anthropic-structured-usage' @@ -461,22 +462,27 @@ const ApiDevtoolsMemoryRoute = ApiDevtoolsMemoryRouteImport.update({ path: '/api/devtools-memory', getParentRoute: () => rootRouteImport, } as any) +const ApiCompactionWireRoute = ApiCompactionWireRouteImport.update({ + id: '/api/compaction-wire', + path: '/api/compaction-wire', + getParentRoute: () => rootRouteImport, +} as any) const ApiChatRoute = ApiChatRouteImport.update({ id: '/api/chat', path: '/api/chat', getParentRoute: () => rootRouteImport, } as any) -const ApiByokChatRoute = ApiByokChatRouteImport.update({ - id: '/api/byok-chat', - path: '/api/byok-chat', - getParentRoute: () => rootRouteImport, -} as any) const ApiByteplusSeedance1080pWireRoute = ApiByteplusSeedance1080pWireRouteImport.update({ id: '/api/byteplus-seedance-1080p-wire', path: '/api/byteplus-seedance-1080p-wire', getParentRoute: () => rootRouteImport, } as any) +const ApiByokChatRoute = ApiByokChatRouteImport.update({ + id: '/api/byok-chat', + path: '/api/byok-chat', + getParentRoute: () => rootRouteImport, +} as any) const ApiAudioRoute = ApiAudioRouteImport.update({ id: '/api/audio', path: '/api/audio', @@ -564,6 +570,7 @@ export interface FileRoutesByFullPath { '/api/byok-chat': typeof ApiByokChatRoute '/api/byteplus-seedance-1080p-wire': typeof ApiByteplusSeedance1080pWireRoute '/api/chat': typeof ApiChatRoute + '/api/compaction-wire': typeof ApiCompactionWireRoute '/api/devtools-memory': typeof ApiDevtoolsMemoryRoute '/api/durable-delivery': typeof ApiDurableDeliveryRoute '/api/durable-takeover': typeof ApiDurableTakeoverRoute @@ -650,6 +657,7 @@ export interface FileRoutesByTo { '/api/byok-chat': typeof ApiByokChatRoute '/api/byteplus-seedance-1080p-wire': typeof ApiByteplusSeedance1080pWireRoute '/api/chat': typeof ApiChatRoute + '/api/compaction-wire': typeof ApiCompactionWireRoute '/api/devtools-memory': typeof ApiDevtoolsMemoryRoute '/api/durable-delivery': typeof ApiDurableDeliveryRoute '/api/durable-takeover': typeof ApiDurableTakeoverRoute @@ -737,6 +745,7 @@ export interface FileRoutesById { '/api/byok-chat': typeof ApiByokChatRoute '/api/byteplus-seedance-1080p-wire': typeof ApiByteplusSeedance1080pWireRoute '/api/chat': typeof ApiChatRoute + '/api/compaction-wire': typeof ApiCompactionWireRoute '/api/devtools-memory': typeof ApiDevtoolsMemoryRoute '/api/durable-delivery': typeof ApiDurableDeliveryRoute '/api/durable-takeover': typeof ApiDurableTakeoverRoute @@ -825,6 +834,7 @@ export interface FileRouteTypes { | '/api/byok-chat' | '/api/byteplus-seedance-1080p-wire' | '/api/chat' + | '/api/compaction-wire' | '/api/devtools-memory' | '/api/durable-delivery' | '/api/durable-takeover' @@ -911,6 +921,7 @@ export interface FileRouteTypes { | '/api/byok-chat' | '/api/byteplus-seedance-1080p-wire' | '/api/chat' + | '/api/compaction-wire' | '/api/devtools-memory' | '/api/durable-delivery' | '/api/durable-takeover' @@ -997,6 +1008,7 @@ export interface FileRouteTypes { | '/api/byok-chat' | '/api/byteplus-seedance-1080p-wire' | '/api/chat' + | '/api/compaction-wire' | '/api/devtools-memory' | '/api/durable-delivery' | '/api/durable-takeover' @@ -1084,6 +1096,7 @@ export interface RootRouteChildren { ApiByokChatRoute: typeof ApiByokChatRoute ApiByteplusSeedance1080pWireRoute: typeof ApiByteplusSeedance1080pWireRoute ApiChatRoute: typeof ApiChatRoute + ApiCompactionWireRoute: typeof ApiCompactionWireRoute ApiDevtoolsMemoryRoute: typeof ApiDevtoolsMemoryRoute ApiDurableDeliveryRoute: typeof ApiDurableDeliveryRoute ApiDurableTakeoverRoute: typeof ApiDurableTakeoverRoute @@ -1628,6 +1641,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiDevtoolsMemoryRouteImport parentRoute: typeof rootRouteImport } + '/api/compaction-wire': { + id: '/api/compaction-wire' + path: '/api/compaction-wire' + fullPath: '/api/compaction-wire' + preLoaderRoute: typeof ApiCompactionWireRouteImport + parentRoute: typeof rootRouteImport + } '/api/chat': { id: '/api/chat' path: '/api/chat' @@ -1635,13 +1655,6 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiChatRouteImport parentRoute: typeof rootRouteImport } - '/api/byok-chat': { - id: '/api/byok-chat' - path: '/api/byok-chat' - fullPath: '/api/byok-chat' - preLoaderRoute: typeof ApiByokChatRouteImport - parentRoute: typeof rootRouteImport - } '/api/byteplus-seedance-1080p-wire': { id: '/api/byteplus-seedance-1080p-wire' path: '/api/byteplus-seedance-1080p-wire' @@ -1649,6 +1662,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiByteplusSeedance1080pWireRouteImport parentRoute: typeof rootRouteImport } + '/api/byok-chat': { + id: '/api/byok-chat' + path: '/api/byok-chat' + fullPath: '/api/byok-chat' + preLoaderRoute: typeof ApiByokChatRouteImport + parentRoute: typeof rootRouteImport + } '/api/audio': { id: '/api/audio' path: '/api/audio' @@ -1817,6 +1837,7 @@ const rootRouteChildren: RootRouteChildren = { ApiByokChatRoute: ApiByokChatRoute, ApiByteplusSeedance1080pWireRoute: ApiByteplusSeedance1080pWireRoute, ApiChatRoute: ApiChatRoute, + ApiCompactionWireRoute: ApiCompactionWireRoute, ApiDevtoolsMemoryRoute: ApiDevtoolsMemoryRoute, ApiDurableDeliveryRoute: ApiDurableDeliveryRoute, ApiDurableTakeoverRoute: ApiDurableTakeoverRoute, diff --git a/testing/e2e/src/routes/api.compaction-wire.ts b/testing/e2e/src/routes/api.compaction-wire.ts new file mode 100644 index 0000000000..8ed4c496da --- /dev/null +++ b/testing/e2e/src/routes/api.compaction-wire.ts @@ -0,0 +1,192 @@ +import { createFileRoute } from '@tanstack/react-router' +import { chat, createChatOptions, maxIterations } from '@tanstack/ai' +import { createOpenaiChat } from '@tanstack/ai-openai' +import { + clearToolResults, + evictOldest, + withCompaction, +} from '@tanstack/ai-compaction' +import type { CompactionStrategy } from '@tanstack/ai-compaction' +import type { ModelMessage } from '@tanstack/ai' +import { memoryPersistence, withPersistence } from '@tanstack/ai-persistence' + +const DUMMY_KEY = 'sk-e2e-test-dummy-key' + +function makeTextStream(callNumber: number): ReadableStream { + const encoder = new TextEncoder() + const responseId = `resp_compaction_${callNumber}` + const itemId = `msg_compaction_${callNumber}` + const events = [ + { + type: 'response.created', + response: { + id: responseId, + object: 'response', + status: 'in_progress', + output: [], + }, + }, + { + type: 'response.output_text.delta', + response_id: responseId, + item_id: itemId, + output_index: 0, + content_index: 0, + delta: 'ok', + }, + { + type: 'response.completed', + response: { + id: responseId, + object: 'response', + status: 'completed', + output: [ + { + id: itemId, + type: 'message', + role: 'assistant', + status: 'completed', + content: [{ type: 'output_text', text: 'ok' }], + }, + ], + usage: { input_tokens: 5, output_tokens: 2, total_tokens: 7 }, + }, + }, + ] + return new ReadableStream({ + start(controller) { + for (const event of events) { + controller.enqueue(encoder.encode(`data: ${JSON.stringify(event)}\n\n`)) + } + controller.enqueue(encoder.encode('data: [DONE]\n\n')) + controller.close() + }, + }) +} + +const FILLER = 'x'.repeat(160) + +// evict: oldest message carries SECRET_ALPHA_ONE, newest carries KEEP_ME_LAST. +const evictMessages: Array = [ + { role: 'user', content: `SECRET_ALPHA_ONE ${FILLER}` }, + { role: 'assistant', content: FILLER }, + { role: 'user', content: FILLER }, + { role: 'assistant', content: FILLER }, + { role: 'user', content: `KEEP_ME_LAST ${FILLER}` }, +] + +// clear: two tool results. Oldest carries SECRET_TOOL_ALPHA (should be stubbed), +// newest carries KEEP_TOOL_BETA (kept). All messages stay in place. +const clearMessages: Array = [ + { role: 'user', content: 'run the tools' }, + { + role: 'assistant', + content: '', + toolCalls: [ + { id: 'a', type: 'function', function: { name: 'f', arguments: '{}' } }, + ], + }, + { role: 'tool', content: `SECRET_TOOL_ALPHA ${FILLER}`, toolCallId: 'a' }, + { + role: 'assistant', + content: '', + toolCalls: [ + { id: 'b', type: 'function', function: { name: 'f', arguments: '{}' } }, + ], + }, + { role: 'tool', content: `KEEP_TOOL_BETA ${FILLER}`, toolCallId: 'b' }, + { role: 'user', content: 'done?' }, +] + +/** + * Wire-format verification for `withCompaction`. A capturing `fetch` records the + * outgoing request body so the spec can assert what each strategy sent. + * + * `?strategy=clear` uses `clearToolResults` on a tool-heavy history; anything + * else uses `evictOldest` on a plain chat history. + */ +export const Route = createFileRoute('/api/compaction-wire')({ + server: { + handlers: { + POST: async ({ request }) => { + const clear = + new URL(request.url).searchParams.get('strategy') === 'clear' + + const requestBodies: Array = [] + + const mockFetch: typeof fetch = async (input, init) => { + const req = + input instanceof Request ? input : new Request(input, init) + requestBodies.push(JSON.parse(await req.text())) + return new Response(makeTextStream(requestBodies.length), { + headers: { 'Content-Type': 'text/event-stream' }, + }) + } + + const messages = clear ? clearMessages : evictMessages + const strategy: CompactionStrategy = clear + ? clearToolResults({ keepRecentToolResults: 1 }) + : evictOldest({ keepRecentTokens: 45 }) + + const adapter = createOpenaiChat('gpt-5.2', DUMMY_KEY, { + fetch: mockFetch, + }) + const persistence = memoryPersistence() + let compactionCount = 0 + + try { + for await (const _ of chat({ + ...createChatOptions({ adapter }), + messages, + threadId: 'compaction-wire', + runId: 'compaction-wire-1', + middleware: [ + withPersistence(persistence), + withCompaction({ + maxTokens: 60, + strategy, + onCompact: () => compactionCount++, + }), + ], + agentLoopStrategy: maxIterations(1), + })) { + // Drain the stream. + } + + for await (const _ of chat({ + ...createChatOptions({ adapter }), + messages: [], + threadId: 'compaction-wire', + runId: 'compaction-wire-2', + middleware: [ + withPersistence(persistence), + withCompaction({ + maxTokens: 60, + strategy, + onCompact: () => compactionCount++, + }), + ], + agentLoopStrategy: maxIterations(1), + })) { + // Drain the restored run. + } + } catch (error) { + return Response.json({ + ok: false, + error: error instanceof Error ? error.message : String(error), + }) + } + + const canonicalMessages = + await persistence.stores.messages.loadThread('compaction-wire') + return Response.json({ + ok: true, + firstRequestBody: requestBodies[0], + secondRequestBody: requestBodies[1], + canonicalMessages, + compactionCount, + }) + }, + }, + }, +}) diff --git a/testing/e2e/tests/compaction-wire.spec.ts b/testing/e2e/tests/compaction-wire.spec.ts new file mode 100644 index 0000000000..4669f27c26 --- /dev/null +++ b/testing/e2e/tests/compaction-wire.spec.ts @@ -0,0 +1,64 @@ +import { expect, test } from './fixtures' + +/** + * Wire-format verification for `withCompaction`. Drives `/api/compaction-wire`, + * which sends a long history through `chat()` with a small `maxTokens` and + * captures the outgoing SDK request. The captured body must show the oldest + * message evicted, the compaction note injected, and the recent tail preserved. + */ +test.describe('withCompaction — wire format', () => { + test('evicts the old head, keeps the recent tail, injects a note', async ({ + request, + }) => { + const response = await request.post('/api/compaction-wire') + expect(response.ok()).toBe(true) + const result = (await response.json()) as { + ok: boolean + error?: string + firstRequestBody: unknown + secondRequestBody: unknown + canonicalMessages: unknown + compactionCount: number + } + if (!result.ok) throw new Error(`Route failed: ${result.error}`) + + const wire = JSON.stringify(result.firstRequestBody) + // Recent tail is preserved verbatim. + expect(wire).toContain('KEEP_ME_LAST') + // The dropped head was replaced by the eviction note. + expect(wire).toContain('omitted to save context') + // The oldest message is gone. + expect(wire).not.toContain('SECRET_ALPHA_ONE') + + // Persistence keeps the canonical transcript, while a later request reuses + // the compacted checkpoint without compacting the same prefix again. + expect(JSON.stringify(result.canonicalMessages)).toContain( + 'SECRET_ALPHA_ONE', + ) + expect(JSON.stringify(result.secondRequestBody)).not.toContain( + 'SECRET_ALPHA_ONE', + ) + expect(result.compactionCount).toBe(1) + }) + + test('clearToolResults stubs old tool output and keeps the recent one', async ({ + request, + }) => { + const response = await request.post('/api/compaction-wire?strategy=clear') + expect(response.ok()).toBe(true) + const result = (await response.json()) as { + ok: boolean + error?: string + firstRequestBody: unknown + } + if (!result.ok) throw new Error(`Route failed: ${result.error}`) + + const wire = JSON.stringify(result.firstRequestBody) + // The most recent tool result is preserved verbatim. + expect(wire).toContain('KEEP_TOOL_BETA') + // The old tool result content is replaced by the stub. + expect(wire).toContain('tool output cleared') + // The old tool result content is gone. + expect(wire).not.toContain('SECRET_TOOL_ALPHA') + }) +}) diff --git a/testing/panel/package.json b/testing/panel/package.json index 7b1e3d63be..ebbaf4e2fa 100644 --- a/testing/panel/package.json +++ b/testing/panel/package.json @@ -15,6 +15,7 @@ "@tanstack/ai": "workspace:*", "@tanstack/ai-anthropic": "workspace:*", "@tanstack/ai-client": "workspace:*", + "@tanstack/ai-compaction": "workspace:*", "@tanstack/ai-event-client": "workspace:*", "@tanstack/ai-gemini": "workspace:*", "@tanstack/ai-grok": "workspace:*", diff --git a/testing/panel/src/components/Header.tsx b/testing/panel/src/components/Header.tsx index b7711d91af..b0849b28b5 100644 --- a/testing/panel/src/components/Header.tsx +++ b/testing/panel/src/components/Header.tsx @@ -12,6 +12,7 @@ import { Menu, Mic, Package, + Scissors, Video, Volume2, X, @@ -139,6 +140,24 @@ export default function Header() { + setIsOpen(false)} + className="flex items-center gap-3 p-3 rounded-lg hover:bg-gray-800 transition-colors mb-2" + activeProps={{ + className: + 'flex items-center gap-3 p-3 rounded-lg bg-cyan-600 hover:bg-cyan-700 transition-colors mb-2', + }} + > + +
+ Compaction + + context + +
+ +

Activities diff --git a/testing/panel/src/lib/compaction-store.ts b/testing/panel/src/lib/compaction-store.ts new file mode 100644 index 0000000000..fa445843b9 --- /dev/null +++ b/testing/panel/src/lib/compaction-store.ts @@ -0,0 +1,27 @@ +import type { CompactionInfo } from '@tanstack/ai-compaction' + +/** + * Process-local record of compaction events for the `/compaction` demo. The + * chat route writes here from `withCompaction`'s `onCompact` callback; the + * inspect route reads it. Same singleton or the reader sees nothing. + */ +export interface CompactionEvent extends CompactionInfo { + /** Wall-clock time the compaction fired. */ + at: number +} + +const eventsByThread = new Map>() + +export function recordCompaction(threadId: string, info: CompactionInfo): void { + const list = eventsByThread.get(threadId) ?? [] + list.push({ ...info, at: Date.now() }) + eventsByThread.set(threadId, list) +} + +export function getCompactions(threadId: string): Array { + return eventsByThread.get(threadId) ?? [] +} + +export function clearCompactions(threadId: string): void { + eventsByThread.delete(threadId) +} diff --git a/testing/panel/src/routeTree.gen.ts b/testing/panel/src/routeTree.gen.ts index c9ce6fdaf3..8f23c3235d 100644 --- a/testing/panel/src/routeTree.gen.ts +++ b/testing/panel/src/routeTree.gen.ts @@ -18,6 +18,7 @@ import { Route as StreamDebuggerRouteImport } from './routes/stream-debugger' import { Route as SimulatorRouteImport } from './routes/simulator' import { Route as MemoryRouteImport } from './routes/memory' import { Route as ImageRouteImport } from './routes/image' +import { Route as CompactionRouteImport } from './routes/compaction' import { Route as AddonManagerRouteImport } from './routes/addon-manager' import { Route as IndexRouteImport } from './routes/index' import { Route as ApiVideoRouteImport } from './routes/api.video' @@ -31,6 +32,8 @@ import { Route as ApiMemoryChatRouteImport } from './routes/api.memory-chat' import { Route as ApiLoadTraceRouteImport } from './routes/api.load-trace' import { Route as ApiListTracesRouteImport } from './routes/api.list-traces' import { Route as ApiImageRouteImport } from './routes/api.image' +import { Route as ApiCompactionInspectRouteImport } from './routes/api.compaction-inspect' +import { Route as ApiCompactionChatRouteImport } from './routes/api.compaction-chat' import { Route as ApiChatRouteImport } from './routes/api.chat' import { Route as ApiAddonChatRouteImport } from './routes/api.addon-chat' @@ -79,6 +82,11 @@ const ImageRoute = ImageRouteImport.update({ path: '/image', getParentRoute: () => rootRouteImport, } as any) +const CompactionRoute = CompactionRouteImport.update({ + id: '/compaction', + path: '/compaction', + getParentRoute: () => rootRouteImport, +} as any) const AddonManagerRoute = AddonManagerRouteImport.update({ id: '/addon-manager', path: '/addon-manager', @@ -144,6 +152,16 @@ const ApiImageRoute = ApiImageRouteImport.update({ path: '/api/image', getParentRoute: () => rootRouteImport, } as any) +const ApiCompactionInspectRoute = ApiCompactionInspectRouteImport.update({ + id: '/api/compaction-inspect', + path: '/api/compaction-inspect', + getParentRoute: () => rootRouteImport, +} as any) +const ApiCompactionChatRoute = ApiCompactionChatRouteImport.update({ + id: '/api/compaction-chat', + path: '/api/compaction-chat', + getParentRoute: () => rootRouteImport, +} as any) const ApiChatRoute = ApiChatRouteImport.update({ id: '/api/chat', path: '/api/chat', @@ -158,6 +176,7 @@ const ApiAddonChatRoute = ApiAddonChatRouteImport.update({ export interface FileRoutesByFullPath { '/': typeof IndexRoute '/addon-manager': typeof AddonManagerRoute + '/compaction': typeof CompactionRoute '/image': typeof ImageRoute '/memory': typeof MemoryRoute '/simulator': typeof SimulatorRoute @@ -169,6 +188,8 @@ export interface FileRoutesByFullPath { '/video': typeof VideoRoute '/api/addon-chat': typeof ApiAddonChatRoute '/api/chat': typeof ApiChatRoute + '/api/compaction-chat': typeof ApiCompactionChatRoute + '/api/compaction-inspect': typeof ApiCompactionInspectRoute '/api/image': typeof ApiImageRoute '/api/list-traces': typeof ApiListTracesRoute '/api/load-trace': typeof ApiLoadTraceRoute @@ -184,6 +205,7 @@ export interface FileRoutesByFullPath { export interface FileRoutesByTo { '/': typeof IndexRoute '/addon-manager': typeof AddonManagerRoute + '/compaction': typeof CompactionRoute '/image': typeof ImageRoute '/memory': typeof MemoryRoute '/simulator': typeof SimulatorRoute @@ -195,6 +217,8 @@ export interface FileRoutesByTo { '/video': typeof VideoRoute '/api/addon-chat': typeof ApiAddonChatRoute '/api/chat': typeof ApiChatRoute + '/api/compaction-chat': typeof ApiCompactionChatRoute + '/api/compaction-inspect': typeof ApiCompactionInspectRoute '/api/image': typeof ApiImageRoute '/api/list-traces': typeof ApiListTracesRoute '/api/load-trace': typeof ApiLoadTraceRoute @@ -211,6 +235,7 @@ export interface FileRoutesById { __root__: typeof rootRouteImport '/': typeof IndexRoute '/addon-manager': typeof AddonManagerRoute + '/compaction': typeof CompactionRoute '/image': typeof ImageRoute '/memory': typeof MemoryRoute '/simulator': typeof SimulatorRoute @@ -222,6 +247,8 @@ export interface FileRoutesById { '/video': typeof VideoRoute '/api/addon-chat': typeof ApiAddonChatRoute '/api/chat': typeof ApiChatRoute + '/api/compaction-chat': typeof ApiCompactionChatRoute + '/api/compaction-inspect': typeof ApiCompactionInspectRoute '/api/image': typeof ApiImageRoute '/api/list-traces': typeof ApiListTracesRoute '/api/load-trace': typeof ApiLoadTraceRoute @@ -239,6 +266,7 @@ export interface FileRouteTypes { fullPaths: | '/' | '/addon-manager' + | '/compaction' | '/image' | '/memory' | '/simulator' @@ -250,6 +278,8 @@ export interface FileRouteTypes { | '/video' | '/api/addon-chat' | '/api/chat' + | '/api/compaction-chat' + | '/api/compaction-inspect' | '/api/image' | '/api/list-traces' | '/api/load-trace' @@ -265,6 +295,7 @@ export interface FileRouteTypes { to: | '/' | '/addon-manager' + | '/compaction' | '/image' | '/memory' | '/simulator' @@ -276,6 +307,8 @@ export interface FileRouteTypes { | '/video' | '/api/addon-chat' | '/api/chat' + | '/api/compaction-chat' + | '/api/compaction-inspect' | '/api/image' | '/api/list-traces' | '/api/load-trace' @@ -291,6 +324,7 @@ export interface FileRouteTypes { | '__root__' | '/' | '/addon-manager' + | '/compaction' | '/image' | '/memory' | '/simulator' @@ -302,6 +336,8 @@ export interface FileRouteTypes { | '/video' | '/api/addon-chat' | '/api/chat' + | '/api/compaction-chat' + | '/api/compaction-inspect' | '/api/image' | '/api/list-traces' | '/api/load-trace' @@ -318,6 +354,7 @@ export interface FileRouteTypes { export interface RootRouteChildren { IndexRoute: typeof IndexRoute AddonManagerRoute: typeof AddonManagerRoute + CompactionRoute: typeof CompactionRoute ImageRoute: typeof ImageRoute MemoryRoute: typeof MemoryRoute SimulatorRoute: typeof SimulatorRoute @@ -329,6 +366,8 @@ export interface RootRouteChildren { VideoRoute: typeof VideoRoute ApiAddonChatRoute: typeof ApiAddonChatRoute ApiChatRoute: typeof ApiChatRoute + ApiCompactionChatRoute: typeof ApiCompactionChatRoute + ApiCompactionInspectRoute: typeof ApiCompactionInspectRoute ApiImageRoute: typeof ApiImageRoute ApiListTracesRoute: typeof ApiListTracesRoute ApiLoadTraceRoute: typeof ApiLoadTraceRoute @@ -407,6 +446,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ImageRouteImport parentRoute: typeof rootRouteImport } + '/compaction': { + id: '/compaction' + path: '/compaction' + fullPath: '/compaction' + preLoaderRoute: typeof CompactionRouteImport + parentRoute: typeof rootRouteImport + } '/addon-manager': { id: '/addon-manager' path: '/addon-manager' @@ -498,6 +544,20 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiImageRouteImport parentRoute: typeof rootRouteImport } + '/api/compaction-inspect': { + id: '/api/compaction-inspect' + path: '/api/compaction-inspect' + fullPath: '/api/compaction-inspect' + preLoaderRoute: typeof ApiCompactionInspectRouteImport + parentRoute: typeof rootRouteImport + } + '/api/compaction-chat': { + id: '/api/compaction-chat' + path: '/api/compaction-chat' + fullPath: '/api/compaction-chat' + preLoaderRoute: typeof ApiCompactionChatRouteImport + parentRoute: typeof rootRouteImport + } '/api/chat': { id: '/api/chat' path: '/api/chat' @@ -518,6 +578,7 @@ declare module '@tanstack/react-router' { const rootRouteChildren: RootRouteChildren = { IndexRoute: IndexRoute, AddonManagerRoute: AddonManagerRoute, + CompactionRoute: CompactionRoute, ImageRoute: ImageRoute, MemoryRoute: MemoryRoute, SimulatorRoute: SimulatorRoute, @@ -529,6 +590,8 @@ const rootRouteChildren: RootRouteChildren = { VideoRoute: VideoRoute, ApiAddonChatRoute: ApiAddonChatRoute, ApiChatRoute: ApiChatRoute, + ApiCompactionChatRoute: ApiCompactionChatRoute, + ApiCompactionInspectRoute: ApiCompactionInspectRoute, ApiImageRoute: ApiImageRoute, ApiListTracesRoute: ApiListTracesRoute, ApiLoadTraceRoute: ApiLoadTraceRoute, diff --git a/testing/panel/src/routes/api.compaction-chat.ts b/testing/panel/src/routes/api.compaction-chat.ts new file mode 100644 index 0000000000..80923eb474 --- /dev/null +++ b/testing/panel/src/routes/api.compaction-chat.ts @@ -0,0 +1,156 @@ +import { createFileRoute } from '@tanstack/react-router' +import { + chat, + createChatOptions, + maxIterations, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { + evictOldest, + summarizeOldest, + withCompaction, +} from '@tanstack/ai-compaction' +import { anthropicText } from '@tanstack/ai-anthropic' +import { geminiText } from '@tanstack/ai-gemini' +import { grokText } from '@tanstack/ai-grok' +import { openaiText } from '@tanstack/ai-openai' +import { ollamaText } from '@tanstack/ai-ollama' +import { openRouterText } from '@tanstack/ai-openrouter' +import { recordCompaction } from '@/lib/compaction-store' +import type { AnyTextAdapter, ModelMessage } from '@tanstack/ai' +import type { Provider } from '@/lib/model-selection' + +// Provider-agnostic summary: one throwaway chat() turn on the same adapter. +async function summarizeWith( + adapter: AnyTextAdapter, + messages: Array, +): Promise { + let text = '' + for await (const chunk of chat({ + adapter, + messages: [ + ...messages, + { + role: 'user', + content: 'Summarize the conversation above in 3-4 sentences.', + }, + ], + agentLoopStrategy: maxIterations(1), + })) { + if (chunk.type === 'TEXT_MESSAGE_CONTENT') text += chunk.delta + } + return text +} + +const SYSTEM_PROMPT = `You are a helpful assistant. Keep answers reasonably long +(a paragraph or two) so this demo's context fills up quickly.` + +/** + * Chat endpoint for the `/compaction` demo. Wires `withCompaction` with a small + * `maxTokens` so the middleware fires after a couple of turns. Compaction here + * evicts the oldest messages (no `summarize` callback), keeping the recent tail + * verbatim; each event is recorded so the page can show before/after tokens. + * + * `threadId` scopes the recorded events; it is demo-only (never trust a + * client-supplied identity in production). + */ +export const Route = createFileRoute('/api/compaction-chat')({ + server: { + handlers: { + POST: async ({ request }) => { + const requestSignal = request.signal + if (requestSignal.aborted) { + return new Response(null, { status: 499 }) + } + + const abortController = new AbortController() + const body = await request.json() + const messages = body.messages + const data = body.data || {} + + const provider: Provider = data.provider || 'openai' + const model: string | undefined = data.model + const threadId: string = + typeof data.threadId === 'string' && data.threadId.length > 0 + ? data.threadId + : 'panel-default-thread' + const maxTokens: number = + typeof data.maxTokens === 'number' && data.maxTokens > 0 + ? data.maxTokens + : 400 + const strategyName: 'evict' | 'summarize' = + data.strategy === 'summarize' ? 'summarize' : 'evict' + + try { + const adapterConfig = { + anthropic: () => + createChatOptions({ + adapter: anthropicText((model || 'claude-sonnet-4-5') as any), + }), + gemini: () => + createChatOptions({ + adapter: geminiText((model || 'gemini-2.5-flash') as any), + }), + grok: () => + createChatOptions({ + adapter: grokText((model || 'grok-build-0.1') as any), + }), + ollama: () => + createChatOptions({ + adapter: ollamaText((model || 'mistral:7b') as any), + }), + openai: () => + createChatOptions({ + adapter: openaiText((model || 'gpt-4o') as any), + }), + openrouter: () => + createChatOptions({ + adapter: openRouterText((model || 'openai/gpt-4o') as any), + }), + } + + const options = adapterConfig[provider]() + const { adapter } = options + + const strategy = + strategyName === 'summarize' + ? summarizeOldest({ + summarize: (msgs) => summarizeWith(adapter, msgs), + }) + : evictOldest() + + const compaction = withCompaction({ + maxTokens, + strategy, + onCompact: (info) => recordCompaction(threadId, info), + }) + + const stream = chat({ + ...options, + adapter, + tools: [], + systemPrompts: [SYSTEM_PROMPT], + middleware: [compaction], + agentLoopStrategy: maxIterations(5), + messages, + abortController, + }) + + return toServerSentEventsResponse(stream, { abortController }) + } catch (error: any) { + console.error('[api.compaction-chat] Error:', error?.message) + if (error.name === 'AbortError' || abortController.signal.aborted) { + return new Response(null, { status: 499 }) + } + return new Response( + JSON.stringify({ error: error.message || 'An error occurred' }), + { + status: 500, + headers: { 'Content-Type': 'application/json' }, + }, + ) + } + }, + }, + }, +}) diff --git a/testing/panel/src/routes/api.compaction-inspect.ts b/testing/panel/src/routes/api.compaction-inspect.ts new file mode 100644 index 0000000000..259473f03f --- /dev/null +++ b/testing/panel/src/routes/api.compaction-inspect.ts @@ -0,0 +1,22 @@ +import { createFileRoute } from '@tanstack/react-router' +import { clearCompactions, getCompactions } from '@/lib/compaction-store' + +/** + * Read side of the `/compaction` demo. GET returns the recorded compaction + * events for a thread; DELETE clears them (used by "New thread"). + */ +export const Route = createFileRoute('/api/compaction-inspect')({ + server: { + handlers: { + GET: async ({ request }) => { + const threadId = new URL(request.url).searchParams.get('threadId') ?? '' + return Response.json({ events: getCompactions(threadId) }) + }, + DELETE: async ({ request }) => { + const threadId = new URL(request.url).searchParams.get('threadId') ?? '' + clearCompactions(threadId) + return Response.json({ ok: true }) + }, + }, + }, +}) diff --git a/testing/panel/src/routes/compaction.tsx b/testing/panel/src/routes/compaction.tsx new file mode 100644 index 0000000000..7495ed0730 --- /dev/null +++ b/testing/panel/src/routes/compaction.tsx @@ -0,0 +1,298 @@ +import { useCallback, useEffect, useMemo, useRef, useState } from 'react' +import { createFileRoute } from '@tanstack/react-router' +import { RefreshCw, RotateCcw, Send, Scissors } from 'lucide-react' +import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' +import type { UIMessage } from '@tanstack/ai-react' +import { MODEL_OPTIONS, getDefaultModelOption } from '@/lib/model-selection' +import type { ModelOption } from '@/lib/model-selection' + +const THREAD_STORAGE_KEY = 'panel-compaction-thread' + +// Mirror of /api/compaction-inspect. Kept local so the page has no build-time +// dependency on server internals. +interface CompactionEvent { + before: number + after: number + messagesBefore: number + messagesAfter: number + at: number +} +interface InspectResponse { + events: Array +} + +function getMessageText(parts: UIMessage['parts']): string { + return parts + .filter((part) => part.type === 'text' && 'content' in part && part.content) + .map((part) => (part as { type: 'text'; content: string }).content) + .join('') +} + +function CompactionPage() { + const [selectedModel, setSelectedModel] = useState( + getDefaultModelOption(), + ) + const [threadId, setThreadId] = useState('') + const [maxTokens, setMaxTokens] = useState(400) + const [strategy, setStrategy] = useState<'evict' | 'summarize'>('evict') + const [inspect, setInspect] = useState(null) + const [input, setInput] = useState('') + + useEffect(() => { + let existing = localStorage.getItem(THREAD_STORAGE_KEY) + if (!existing) { + existing = crypto.randomUUID() + localStorage.setItem(THREAD_STORAGE_KEY, existing) + } + setThreadId(existing) + }, []) + + const body = useMemo( + () => ({ + provider: selectedModel.provider, + model: selectedModel.model, + threadId, + maxTokens, + strategy, + }), + [ + selectedModel.provider, + selectedModel.model, + threadId, + maxTokens, + strategy, + ], + ) + + const { messages, sendMessage, isLoading } = useChat({ + connection: fetchServerSentEvents('/api/compaction-chat'), + body, + devtools: { name: 'Compaction' }, + }) + + const refreshInspect = useCallback(async () => { + if (!threadId) return + try { + const res = await fetch( + `/api/compaction-inspect?threadId=${encodeURIComponent(threadId)}`, + ) + if (res.ok) setInspect(await res.json()) + } catch { + // Non-fatal: read-only view. + } + }, [threadId]) + + const wasLoading = useRef(false) + useEffect(() => { + if (wasLoading.current && !isLoading) refreshInspect() + wasLoading.current = isLoading + }, [isLoading, refreshInspect]) + useEffect(() => { + refreshInspect() + }, [refreshInspect]) + + const startNewThread = async () => { + if (threadId) { + await fetch( + `/api/compaction-inspect?threadId=${encodeURIComponent(threadId)}`, + { method: 'DELETE' }, + ).catch(() => {}) + } + const next = crypto.randomUUID() + localStorage.setItem(THREAD_STORAGE_KEY, next) + setThreadId(next) + setInspect(null) + } + + const submit = () => { + const text = input.trim() + if (!text || isLoading) return + sendMessage(text) + setInput('') + } + + const events = inspect?.events ?? [] + + return ( +

+ {/* Left: chat */} +
+
+
+ + +
+
+ + setMaxTokens(parseInt(e.target.value))} + className="w-full accent-cyan-500" + /> +
+
+ + +
+
+ +
+ {messages.length === 0 ? ( +

+ Chat for a few turns. Once the running transcript passes{' '} + {maxTokens} estimated tokens, older messages get compacted away + and the events show up on the right. +

+ ) : ( + messages.map(({ id, role, parts }) => ( +
+
+ {getMessageText(parts)} +
+
+ )) + )} +
+ +
+
+ setInput(e.target.value)} + onKeyDown={(e) => { + if (e.key === 'Enter' && !e.shiftKey) { + e.preventDefault() + submit() + } + }} + placeholder="Type a message…" + disabled={isLoading} + className="flex-1 rounded-lg border border-cyan-500/20 bg-gray-800 px-3 py-2 text-sm text-white focus:outline-none focus:ring-2 focus:ring-cyan-500/50 disabled:opacity-50" + /> + +
+
+
+ + {/* Right: compaction events */} +
+
+
+

Compaction events

+

+ thread: {threadId ? threadId.slice(0, 8) : '…'} +

+
+
+ + +
+
+ +
+ {events.length === 0 ? ( +

+ No compaction yet. Lower maxTokens or keep chatting until the + transcript grows past the threshold. +

+ ) : ( + events + .slice() + .reverse() + .map((ev, i) => ( +
+
+ + Compacted {ev.messagesBefore} → {ev.messagesAfter} messages +
+
+ {ev.before} → {ev.after} tokens (− + {ev.before - ev.after}) +
+
+ {new Date(ev.at).toLocaleTimeString()} +
+
+ )) + )} +
+
+
+ ) +} + +export const Route = createFileRoute('/compaction')({ + component: CompactionPage, +})