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-devtools.md b/.changeset/compaction-devtools.md new file mode 100644 index 0000000000..5b4ff14942 --- /dev/null +++ b/.changeset/compaction-devtools.md @@ -0,0 +1,13 @@ +--- +'@tanstack/ai-compaction': minor +'@tanstack/ai-client': patch +'@tanstack/ai-event-client': patch +'@tanstack/ai-devtools-core': patch +--- + +Show compaction in TanStack AI DevTools. `withCompaction` injects +`compaction:started`, `compaction:state`, and `compaction:ended` CUSTOM +stream events. State includes before/after counts, the token budget, and +dropped vs sent message previews. The chat client re-emits the same three +events. The AI panel has a Compaction tab and started/state/ended steps on +the iteration. 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/.changeset/middleware-emit-custom-event.md b/.changeset/middleware-emit-custom-event.md new file mode 100644 index 0000000000..aa96e98177 --- /dev/null +++ b/.changeset/middleware-emit-custom-event.md @@ -0,0 +1,9 @@ +--- +'@tanstack/ai': minor +'@tanstack/ai-compaction': patch +--- + +Add `ctx.emitCustomEvent` on chat middleware context. The engine yields +`CUSTOM` chunks while hooks such as `onConfig` are still running, so a long +middleware step can send progress before it finishes. Compaction uses this +to emit `compaction:started` before the strategy returns. diff --git a/docs/advanced/compaction.md b/docs/advanced/compaction.md new file mode 100644 index 0000000000..874dda6117 --- /dev/null +++ b/docs/advanced/compaction.md @@ -0,0 +1,247 @@ +--- +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, or a custom eviction marker. Change it when your `summarize` function can change. | +| `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` (default `assistant`) | +| `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. +- **The canonical transcript stays complete.** Compaction writes provider-only context. Persistence and other middleware still read `ctx.messages`. + +## DevTools + +After a compaction, the chat stream includes three CUSTOM events in order: +`compaction:started`, `compaction:state`, then `compaction:ended`. +`compaction:started` is sent before the strategy runs, so a slow +`summarizeOldest` call still shows up as started on the client. The state and +ended events follow when the strategy returns. +TanStack AI DevTools has a Compaction tab on the hook. Each compact shows: + +- started, state, and ended rows +- when it ran, with token and message counts +- the `maxTokens` budget +- dropped messages +- the transcript sent to the model + +The conversation timeline also keeps `onCompactStart`, `onCompact`, and +`onCompactEnd` steps. + +Open the AI plugin in the DevTools panel (the `ts-react-chat` example mounts it). Select the Compaction hook, then open the Compaction tab. + +The `/compaction` route in `examples/ts-react-chat` uses a small `maxTokens` so this fires after a few turns. That page also shows a compact banner in the chat. The canonical transcript stays complete. The banner is example UI, not part of `useChat`. + +## 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. + +When a checkpoint is reused, a later `summarizeOldest` pass sees the previous +summary plus new messages. Then it folds the old summary into the new one. +Folding needs a metadata store and a strategy key. + +The default strategy, standard `evictOldest`, `summarizeOldest`, +`clearToolResults`, and safe compositions get a strategy key automatically. +Set `strategyKey` for custom strategies, custom estimators, or custom marker +functions. Change `strategyKey` when your `summarize` function can change. +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..7c68f5cff4 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). | @@ -686,6 +692,7 @@ Every hook receives a `ChatMiddlewareContext` as its first argument. It provides | `chunkIndex` | `number` | Running count of chunks yielded | | `signal` | `AbortSignal \| undefined` | External abort signal | | `abort(reason?)` | `function` | Abort the run from within middleware | +| `emitCustomEvent(name, value)` | `function` | Push a `CUSTOM` chunk onto the chat stream now. The engine yields it while the current hook is still running, including during `onConfig`. | | `context` | `TContext` | User-provided runtime context value | | `defer(promise)` | `function` | Register a non-blocking side-effect | @@ -968,6 +975,32 @@ See [Built-in Middleware](./built-in-middleware) for full options and examples f ## Recipes +### Live custom events + +A long `onConfig` hook can send progress to the client while it waits. Call `ctx.emitCustomEvent` when the work starts, then again when it finishes: + +```typescript +import { type ChatMiddleware } from "@tanstack/ai"; + +async function prepare() { + await new Promise((resolve) => { + setTimeout(resolve, 1); + }); +} + +const progress: ChatMiddleware = { + name: "progress", + async onConfig(ctx) { + if (ctx.phase !== "beforeModel") return; + ctx.emitCustomEvent("job:started", { step: "prepare" }); + await prepare(); + ctx.emitCustomEvent("job:ended", { step: "prepare" }); + }, +}; +``` + +The engine yields each `CUSTOM` chunk as soon as you call `emitCustomEvent`. If `RUN_STARTED` is not on the wire yet, the engine sends it first. Read these events on the client the same way as tool `emitCustomEvent` calls. See [Custom Events](../protocol/custom-events). + ### Rate Limiting Limit the number of tool calls per request: @@ -1140,6 +1173,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..638852a2c7 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", @@ -370,7 +370,7 @@ "label": "Custom Events Reference", "to": "protocol/custom-events", "addedAt": "2026-07-03", - "updatedAt": "2026-08-21" + "updatedAt": "2026-08-27" } ] }, @@ -537,7 +537,7 @@ "label": "Middleware", "to": "advanced/middleware", "addedAt": "2026-04-15", - "updatedAt": "2026-08-19" + "updatedAt": "2026-08-27" }, { "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-27" + }, { "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/docs/protocol/custom-events.md b/docs/protocol/custom-events.md index a95b649cfb..edd3e0126c 100644 --- a/docs/protocol/custom-events.md +++ b/docs/protocol/custom-events.md @@ -114,8 +114,10 @@ Read them with the same `chunk.type === "CUSTOM" && chunk.name === "..."` branch ## Your own custom events aren't in this union -Tools can emit arbitrary, application-defined events through the -`emitCustomEvent` context API: +Tools and chat middleware can emit application-defined events through +`emitCustomEvent`. + +A server tool receives `emitCustomEvent` on its execution context: ```ts import { toolDefinition } from "@tanstack/ai"; @@ -136,6 +138,30 @@ const importRows = toolDefinition({ }); ``` +Chat middleware calls the same helper on `ChatMiddlewareContext`. The engine +yields the chunk while the hook is still running, so a long `onConfig` can +send `started` before the work finishes: + +```ts +import { type ChatMiddleware } from "@tanstack/ai"; + +async function prepare() { + await new Promise((resolve) => { + setTimeout(resolve, 1); + }); +} + +const progress: ChatMiddleware = { + name: "progress", + async onConfig(ctx) { + if (ctx.phase !== "beforeModel") return; + ctx.emitCustomEvent("my-app:progress", { step: "prepare" }); + await prepare(); + ctx.emitCustomEvent("my-app:progress", { step: "ready" }); + }, +}; +``` + These flow over the wire exactly like the built-in events: same `CUSTOM` chunk shape, same runtime behavior. But `'my-app:progress'` isn't one of the literal names in `KnownCustomEvent`, so it's intentionally absent from diff --git a/examples/ts-react-chat/package.json b/examples/ts-react-chat/package.json index 15cd77c5b5..846f3745e9 100644 --- a/examples/ts-react-chat/package.json +++ b/examples/ts-react-chat/package.json @@ -25,6 +25,7 @@ "@tanstack/ai-byteplus": "workspace:*", "@tanstack/ai-claude-code": "workspace:*", "@tanstack/ai-client": "workspace:*", + "@tanstack/ai-compaction": "workspace:*", "@tanstack/ai-code-mode": "workspace:*", "@tanstack/ai-codex": "workspace:*", "@tanstack/ai-elevenlabs": "workspace:*", diff --git a/examples/ts-react-chat/src/components/Header.tsx b/examples/ts-react-chat/src/components/Header.tsx index ae2a3cbf3a..fc6d50cbb5 100644 --- a/examples/ts-react-chat/src/components/Header.tsx +++ b/examples/ts-react-chat/src/components/Header.tsx @@ -21,6 +21,7 @@ import { PauseCircle, Plug, RefreshCw, + Scissors, Server, Sparkles, Video, @@ -250,6 +251,19 @@ export default function Header() { Examples

+ setIsOpen(false)} + className="flex items-center gap-3 p-3 rounded-lg hover:bg-gray-800 transition-colors mb-1" + activeProps={{ + className: + 'flex items-center gap-3 p-3 rounded-lg bg-cyan-600 hover:bg-cyan-700 transition-colors mb-1', + }} + > + + Compaction + + setIsOpen(false)} diff --git a/examples/ts-react-chat/src/routeTree.gen.ts b/examples/ts-react-chat/src/routeTree.gen.ts index c4dd562270..7c8b6ab002 100644 --- a/examples/ts-react-chat/src/routeTree.gen.ts +++ b/examples/ts-react-chat/src/routeTree.gen.ts @@ -28,6 +28,7 @@ import { Route as ImageToolReproRouteImport } from './routes/image-tool-repro' import { Route as ImageGenRouteImport } from './routes/image-gen' import { Route as GenericInterruptsRouteImport } from './routes/generic-interrupts' import { Route as GenerationHooksRouteImport } from './routes/generation-hooks' +import { Route as CompactionRouteImport } from './routes/compaction' import { Route as CapabilityDemoRouteImport } from './routes/capability-demo' import { Route as AppStudioRouteImport } from './routes/app-studio' import { Route as IndexRouteImport } from './routes/index' @@ -65,6 +66,7 @@ import { Route as ApiInterruptsRouteImport } from './routes/api.interrupts' import { Route as ApiImageToolReproRouteImport } from './routes/api.image-tool-repro' import { Route as ApiImageGenRouteImport } from './routes/api.image-gen' import { Route as ApiGenericInterruptsRouteImport } from './routes/api.generic-interrupts' +import { Route as ApiCompactionRouteImport } from './routes/api.compaction' import { Route as ApiCapabilityDemoRouteImport } from './routes/api.capability-demo' import { Route as ApiArtifactsRouteImport } from './routes/api.artifacts' import { Route as ApiAppStudioForkRouteImport } from './routes/api.app-studio-fork' @@ -172,6 +174,11 @@ const GenerationHooksRoute = GenerationHooksRouteImport.update({ path: '/generation-hooks', getParentRoute: () => rootRouteImport, } as any) +const CompactionRoute = CompactionRouteImport.update({ + id: '/compaction', + path: '/compaction', + getParentRoute: () => rootRouteImport, +} as any) const CapabilityDemoRoute = CapabilityDemoRouteImport.update({ id: '/capability-demo', path: '/capability-demo', @@ -362,6 +369,11 @@ const ApiGenericInterruptsRoute = ApiGenericInterruptsRouteImport.update({ path: '/api/generic-interrupts', getParentRoute: () => rootRouteImport, } as any) +const ApiCompactionRoute = ApiCompactionRouteImport.update({ + id: '/api/compaction', + path: '/api/compaction', + getParentRoute: () => rootRouteImport, +} as any) const ApiCapabilityDemoRoute = ApiCapabilityDemoRouteImport.update({ id: '/api/capability-demo', path: '/api/capability-demo', @@ -423,6 +435,7 @@ export interface FileRoutesByFullPath { '/': typeof IndexRoute '/app-studio': typeof AppStudioRoute '/capability-demo': typeof CapabilityDemoRoute + '/compaction': typeof CompactionRoute '/generation-hooks': typeof GenerationHooksRoute '/generic-interrupts': typeof GenericInterruptsRoute '/image-gen': typeof ImageGenRoute @@ -446,6 +459,7 @@ export interface FileRoutesByFullPath { '/api/app-studio-fork': typeof ApiAppStudioForkRoute '/api/artifacts': typeof ApiArtifactsRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute + '/api/compaction': typeof ApiCompactionRoute '/api/generic-interrupts': typeof ApiGenericInterruptsRoute '/api/image-gen': typeof ApiImageGenRoute '/api/image-tool-repro': typeof ApiImageToolReproRoute @@ -492,6 +506,7 @@ export interface FileRoutesByTo { '/': typeof IndexRoute '/app-studio': typeof AppStudioRoute '/capability-demo': typeof CapabilityDemoRoute + '/compaction': typeof CompactionRoute '/generation-hooks': typeof GenerationHooksRoute '/generic-interrupts': typeof GenericInterruptsRoute '/image-gen': typeof ImageGenRoute @@ -515,6 +530,7 @@ export interface FileRoutesByTo { '/api/app-studio-fork': typeof ApiAppStudioForkRoute '/api/artifacts': typeof ApiArtifactsRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute + '/api/compaction': typeof ApiCompactionRoute '/api/generic-interrupts': typeof ApiGenericInterruptsRoute '/api/image-gen': typeof ApiImageGenRoute '/api/image-tool-repro': typeof ApiImageToolReproRoute @@ -562,6 +578,7 @@ export interface FileRoutesById { '/': typeof IndexRoute '/app-studio': typeof AppStudioRoute '/capability-demo': typeof CapabilityDemoRoute + '/compaction': typeof CompactionRoute '/generation-hooks': typeof GenerationHooksRoute '/generic-interrupts': typeof GenericInterruptsRoute '/image-gen': typeof ImageGenRoute @@ -585,6 +602,7 @@ export interface FileRoutesById { '/api/app-studio-fork': typeof ApiAppStudioForkRoute '/api/artifacts': typeof ApiArtifactsRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute + '/api/compaction': typeof ApiCompactionRoute '/api/generic-interrupts': typeof ApiGenericInterruptsRoute '/api/image-gen': typeof ApiImageGenRoute '/api/image-tool-repro': typeof ApiImageToolReproRoute @@ -633,6 +651,7 @@ export interface FileRouteTypes { | '/' | '/app-studio' | '/capability-demo' + | '/compaction' | '/generation-hooks' | '/generic-interrupts' | '/image-gen' @@ -656,6 +675,7 @@ export interface FileRouteTypes { | '/api/app-studio-fork' | '/api/artifacts' | '/api/capability-demo' + | '/api/compaction' | '/api/generic-interrupts' | '/api/image-gen' | '/api/image-tool-repro' @@ -702,6 +722,7 @@ export interface FileRouteTypes { | '/' | '/app-studio' | '/capability-demo' + | '/compaction' | '/generation-hooks' | '/generic-interrupts' | '/image-gen' @@ -725,6 +746,7 @@ export interface FileRouteTypes { | '/api/app-studio-fork' | '/api/artifacts' | '/api/capability-demo' + | '/api/compaction' | '/api/generic-interrupts' | '/api/image-gen' | '/api/image-tool-repro' @@ -771,6 +793,7 @@ export interface FileRouteTypes { | '/' | '/app-studio' | '/capability-demo' + | '/compaction' | '/generation-hooks' | '/generic-interrupts' | '/image-gen' @@ -794,6 +817,7 @@ export interface FileRouteTypes { | '/api/app-studio-fork' | '/api/artifacts' | '/api/capability-demo' + | '/api/compaction' | '/api/generic-interrupts' | '/api/image-gen' | '/api/image-tool-repro' @@ -841,6 +865,7 @@ export interface RootRouteChildren { IndexRoute: typeof IndexRoute AppStudioRoute: typeof AppStudioRoute CapabilityDemoRoute: typeof CapabilityDemoRoute + CompactionRoute: typeof CompactionRoute GenerationHooksRoute: typeof GenerationHooksRoute GenericInterruptsRoute: typeof GenericInterruptsRoute ImageGenRoute: typeof ImageGenRoute @@ -864,6 +889,7 @@ export interface RootRouteChildren { ApiAppStudioForkRoute: typeof ApiAppStudioForkRoute ApiArtifactsRoute: typeof ApiArtifactsRoute ApiCapabilityDemoRoute: typeof ApiCapabilityDemoRoute + ApiCompactionRoute: typeof ApiCompactionRoute ApiGenericInterruptsRoute: typeof ApiGenericInterruptsRoute ApiImageGenRoute: typeof ApiImageGenRoute ApiImageToolReproRoute: typeof ApiImageToolReproRoute @@ -1041,6 +1067,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof GenerationHooksRouteImport parentRoute: typeof rootRouteImport } + '/compaction': { + id: '/compaction' + path: '/compaction' + fullPath: '/compaction' + preLoaderRoute: typeof CompactionRouteImport + parentRoute: typeof rootRouteImport + } '/capability-demo': { id: '/capability-demo' path: '/capability-demo' @@ -1300,6 +1333,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiGenericInterruptsRouteImport parentRoute: typeof rootRouteImport } + '/api/compaction': { + id: '/api/compaction' + path: '/api/compaction' + fullPath: '/api/compaction' + preLoaderRoute: typeof ApiCompactionRouteImport + parentRoute: typeof rootRouteImport + } '/api/capability-demo': { id: '/api/capability-demo' path: '/api/capability-demo' @@ -1395,6 +1435,7 @@ const rootRouteChildren: RootRouteChildren = { IndexRoute: IndexRoute, AppStudioRoute: AppStudioRoute, CapabilityDemoRoute: CapabilityDemoRoute, + CompactionRoute: CompactionRoute, GenerationHooksRoute: GenerationHooksRoute, GenericInterruptsRoute: GenericInterruptsRoute, ImageGenRoute: ImageGenRoute, @@ -1418,6 +1459,7 @@ const rootRouteChildren: RootRouteChildren = { ApiAppStudioForkRoute: ApiAppStudioForkRoute, ApiArtifactsRoute: ApiArtifactsRoute, ApiCapabilityDemoRoute: ApiCapabilityDemoRoute, + ApiCompactionRoute: ApiCompactionRoute, ApiGenericInterruptsRoute: ApiGenericInterruptsRoute, ApiImageGenRoute: ApiImageGenRoute, ApiImageToolReproRoute: ApiImageToolReproRoute, diff --git a/examples/ts-react-chat/src/routes/api.compaction.ts b/examples/ts-react-chat/src/routes/api.compaction.ts new file mode 100644 index 0000000000..f33af94a8e --- /dev/null +++ b/examples/ts-react-chat/src/routes/api.compaction.ts @@ -0,0 +1,166 @@ +import { createFileRoute } from '@tanstack/react-router' +import { + chat, + chatParamsFromRequestBody, + 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 { groqText } from '@tanstack/ai-groq' +import { openaiText } from '@tanstack/ai-openai' +import { ollamaText } from '@tanstack/ai-ollama' +import { openRouterText } from '@tanstack/ai-openrouter' +import type { AnyTextAdapter, ModelMessage } from '@tanstack/ai' +import type { Provider } from '@/lib/model-selection' + +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.` + +function adapterFor(provider: Provider, model: string): AnyTextAdapter { + switch (provider) { + case 'anthropic': + return anthropicText( + (model || 'claude-sonnet-4-6') as 'claude-sonnet-4-6', + ) + case 'gemini': + return geminiText( + (model || 'gemini-3.1-pro-preview') as 'gemini-3.1-pro-preview', + ) + case 'grok': + return grokText((model || 'grok-build-0.1') as 'grok-build-0.1') + case 'groq': + return groqText((model || 'openai/gpt-oss-120b') as 'openai/gpt-oss-120b') + case 'ollama': + return ollamaText((model || 'mistral:7b') as 'mistral:7b') + case 'openrouter': + return openRouterText((model || 'openai/gpt-5.1') as 'openai/gpt-5.1') + case 'openai': + default: + return openaiText((model || 'gpt-5.5') as 'gpt-5.5') + } +} + +/** + * Chat endpoint for `/compaction`. Uses a small `maxTokens` so compaction + * fires after a few turns. Keys come from `examples/ts-react-chat/.env`. + */ +export const Route = createFileRoute('/api/compaction')({ + server: { + handlers: { + POST: async ({ request }) => { + const requestSignal = request.signal + if (requestSignal.aborted) { + return new Response(null, { status: 499 }) + } + + const abortController = new AbortController() + + let params + try { + params = await chatParamsFromRequestBody(await request.json()) + } catch (error) { + return new Response( + error instanceof Error ? error.message : 'Bad request', + { status: 400 }, + ) + } + + const requestedProvider = + typeof params.forwardedProps.provider === 'string' + ? params.forwardedProps.provider + : 'openai' + const model: string = + typeof params.forwardedProps.model === 'string' + ? params.forwardedProps.model + : 'gpt-5.5' + const maxTokens: number = + typeof params.forwardedProps.maxTokens === 'number' && + params.forwardedProps.maxTokens > 0 + ? params.forwardedProps.maxTokens + : 400 + const strategyName: 'evict' | 'summarize' = + params.forwardedProps.strategy === 'summarize' ? 'summarize' : 'evict' + + try { + const provider: Provider = [ + 'anthropic', + 'gemini', + 'grok', + 'groq', + 'ollama', + 'openai', + 'openrouter', + ].includes(requestedProvider) + ? (requestedProvider as Provider) + : 'openai' + + const adapter = adapterFor(provider, model) + const options = createChatOptions({ adapter }) + + const strategy = + strategyName === 'summarize' + ? summarizeOldest({ + summarize: (msgs) => summarizeWith(adapter, msgs), + }) + : evictOldest() + + const stream = chat({ + ...options, + adapter, + tools: [], + systemPrompts: [SYSTEM_PROMPT], + middleware: [withCompaction({ maxTokens, strategy })], + agentLoopStrategy: maxIterations(5), + messages: params.messages, + abortController, + }) + + return toServerSentEventsResponse(stream, { abortController }) + } catch (error) { + const message = + error instanceof Error ? error.message : 'An error occurred' + console.error('[api.compaction] Error:', message) + if ( + (error instanceof Error && error.name === 'AbortError') || + abortController.signal.aborted + ) { + return new Response(null, { status: 499 }) + } + return new Response(JSON.stringify({ error: message }), { + status: 500, + headers: { 'Content-Type': 'application/json' }, + }) + } + }, + }, + }, +}) diff --git a/examples/ts-react-chat/src/routes/compaction.tsx b/examples/ts-react-chat/src/routes/compaction.tsx new file mode 100644 index 0000000000..cd25f3b36d --- /dev/null +++ b/examples/ts-react-chat/src/routes/compaction.tsx @@ -0,0 +1,246 @@ +import { useMemo, useState } from 'react' +import { createFileRoute } from '@tanstack/react-router' +import { Send, Scissors } from 'lucide-react' +import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' +import { + COMPACTION_ENDED_EVENT, + COMPACTION_STARTED_EVENT, + COMPACTION_STATE_EVENT, +} from '@tanstack/ai-compaction' +import type { CompactionStateEventValue } from '@tanstack/ai-compaction' +import type { UIMessage } from '@tanstack/ai-react' +import { DEFAULT_MODEL_OPTION, MODEL_OPTIONS } from '@/lib/model-selection' +import type { ModelOption } from '@/lib/model-selection' + +function isCompactionState(data: unknown): data is CompactionStateEventValue { + if (!data || typeof data !== 'object') return false + if ( + !('before' in data) || + !('after' in data) || + !('messagesBefore' in data) || + !('messagesAfter' in data) + ) { + return false + } + return ( + typeof data.before === 'number' && + typeof data.after === 'number' && + typeof data.messagesBefore === 'number' && + typeof data.messagesAfter === 'number' + ) +} + +function getMessageText(parts: UIMessage['parts']): string { + return parts + .filter( + (part): part is Extract<(typeof parts)[number], { type: 'text' }> => + part.type === 'text', + ) + .map((part) => part.content) + .join('') +} + +function CompactionPage() { + const [selectedModel, setSelectedModel] = + useState(DEFAULT_MODEL_OPTION) + const [maxTokens, setMaxTokens] = useState(400) + const [strategy, setStrategy] = useState<'evict' | 'summarize'>('evict') + const [input, setInput] = useState('') + const [lastCompact, setLastCompact] = + useState(null) + const [compactPhase, setCompactPhase] = useState< + 'idle' | 'started' | 'ended' + >('idle') + const [compactEvents, setCompactEvents] = useState>([]) + + const forwardedProps = useMemo( + () => ({ + provider: selectedModel.provider, + model: selectedModel.model, + maxTokens, + strategy, + }), + [selectedModel.provider, selectedModel.model, maxTokens, strategy], + ) + + const { messages, sendMessage, isLoading, error } = useChat({ + connection: fetchServerSentEvents('/api/compaction'), + threadId: 'compaction-demo', + forwardedProps, + devtools: { name: 'Compaction' }, + onCustomEvent: (eventType, data) => { + if (eventType === COMPACTION_STARTED_EVENT) { + setCompactPhase('started') + setLastCompact(null) + setCompactEvents((events) => [...events.slice(-8), 'started']) + return + } + if (eventType === COMPACTION_STATE_EVENT) { + if (!isCompactionState(data)) return + setLastCompact(data) + setCompactEvents((events) => [...events.slice(-8), 'state']) + return + } + if (eventType === COMPACTION_ENDED_EVENT) { + setCompactPhase('ended') + setCompactEvents((events) => [...events.slice(-8), 'ended']) + } + }, + }) + + const submit = () => { + const text = input.trim() + if (!text || isLoading) return + sendMessage(text) + setInput('') + } + + return ( +
+
+
+ +

Compaction

+
+

+ Chat until the transcript passes maxTokens. Then open TanStack + DevTools (bottom-right), pick the AI plugin, and open the Compaction + tab. API keys come from .env. + You do not paste a key here. +

+
+ + +
+
+ + setMaxTokens(parseInt(e.target.value))} + className="w-full accent-cyan-500" + /> +
+
+ + +
+ {compactPhase === 'started' && !lastCompact && ( +
+ Compacting… +
+ )} + {lastCompact && ( +
+ {compactPhase === 'started' ? 'Compacting… ' : 'Compacted '} + {lastCompact.messagesBefore} → {lastCompact.messagesAfter} messages + ({lastCompact.before} → {lastCompact.after} tokens) + {lastCompact.strategyKey ? ` · ${lastCompact.strategyKey}` : ''} + {compactEvents.length > 0 ? ` · ${compactEvents.join(' → ')}` : ''} +
+ )} +
+ +
+ {messages.length === 0 ? ( +

+ Send a few long messages. Once the running transcript passes{' '} + {maxTokens} estimated tokens, older messages are compacted for the + model only. The chat still shows the full transcript. +

+ ) : ( + messages.map(({ id, role, parts }) => ( +
+
+ {getMessageText(parts)} +
+
+ )) + )} +
+ + {error && ( +
+ {error.message} +
+ )} + +
+
+ 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" + /> + +
+
+
+ ) +} + +export const Route = createFileRoute('/compaction')({ + component: CompactionPage, +}) diff --git a/examples/ts-react-chat/vite.config.ts b/examples/ts-react-chat/vite.config.ts index e951c7758e..d91178a03b 100644 --- a/examples/ts-react-chat/vite.config.ts +++ b/examples/ts-react-chat/vite.config.ts @@ -1,3 +1,5 @@ +import { existsSync } from 'node:fs' +import { resolve } from 'node:path' import { defineConfig } from 'vite' import { tanstackStart } from '@tanstack/react-start/plugin/vite' import viteReact from '@vitejs/plugin-react' @@ -6,6 +8,13 @@ import { nitro } from 'nitro/vite' import { devtools } from '@tanstack/devtools-vite' import { webSocketChatPlugin } from './src/lib/websocket-chat-plugin.ts' +// Server routes read process.env. Vite does not copy unprefixed .env keys +// into the SSR process. Load them here. Already-set vars win. +for (const name of ['.env.local', '.env']) { + const path = resolve(import.meta.dirname, name) + if (existsSync(path)) process.loadEnvFile(path) +} + // `dockerode` is a server-only dependency that pulls in optional native addons // (`ssh2` → `cpu-features`, a `.node` binary that this install does not compile). // At runtime `ssh2` catches the missing addon, but the bundlers don't: diff --git a/packages/ai-client/src/chat-client.ts b/packages/ai-client/src/chat-client.ts index ac7afe8194..0f4fede535 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -820,6 +820,13 @@ export class ChatClient< if (eventType === 'memory:state') { this.devtoolsBridge.recordMemoryState(data) } + if ( + eventType === 'compaction:started' || + eventType === 'compaction:state' || + eventType === 'compaction:ended' + ) { + this.devtoolsBridge.recordCompactionEvent(eventType, data) + } this.callbacksRef.current.onCustomEvent(eventType, data, context) }, }, diff --git a/packages/ai-client/src/devtools-noop.ts b/packages/ai-client/src/devtools-noop.ts index 6a1f1ce00b..cd1e0ca011 100644 --- a/packages/ai-client/src/devtools-noop.ts +++ b/packages/ai-client/src/devtools-noop.ts @@ -87,6 +87,8 @@ export class NoOpChatDevtoolsBridge { } observeChunk(_chunk: StreamChunk): void {} recordMemoryState(_value: unknown): void {} + recordCompactionEvent(_eventType: string, _value: unknown): void {} + recordCompactionState(_value: unknown): void {} beginRun(_runId: string, _threadId: string): void {} getCurrentRunEventContext(): ChatClientRunEventContext | undefined { return undefined diff --git a/packages/ai-client/src/devtools.ts b/packages/ai-client/src/devtools.ts index a5bee8a6a8..d06fc2dfdd 100644 --- a/packages/ai-client/src/devtools.ts +++ b/packages/ai-client/src/devtools.ts @@ -8,6 +8,7 @@ import { DefaultChatClientEventEmitter } from './events' import type { AnyClientTool, StreamChunk } from '@tanstack/ai/client' import type { AIDevtoolsEventVisibility, + CompactionMessagePreview, MemoryScopeLite, } from '@tanstack/ai-event-client' import type { @@ -55,6 +56,127 @@ interface MemoryStateEventValue { } } +function readCompactionBoundaryValue(rawValue: unknown): { + before?: number + after?: number + messagesBefore?: number + messagesAfter?: number + reusedCheckpoint?: boolean + maxTokens?: number + strategyKey?: string + durationMs?: number +} { + if (!rawValue || typeof rawValue !== 'object') return {} + return { + ...('before' in rawValue && typeof rawValue.before === 'number' + ? { before: rawValue.before } + : {}), + ...('after' in rawValue && typeof rawValue.after === 'number' + ? { after: rawValue.after } + : {}), + ...('messagesBefore' in rawValue && + typeof rawValue.messagesBefore === 'number' + ? { messagesBefore: rawValue.messagesBefore } + : {}), + ...('messagesAfter' in rawValue && + typeof rawValue.messagesAfter === 'number' + ? { messagesAfter: rawValue.messagesAfter } + : {}), + ...('reusedCheckpoint' in rawValue && + typeof rawValue.reusedCheckpoint === 'boolean' + ? { reusedCheckpoint: rawValue.reusedCheckpoint } + : {}), + ...('maxTokens' in rawValue && typeof rawValue.maxTokens === 'number' + ? { maxTokens: rawValue.maxTokens } + : {}), + ...('strategyKey' in rawValue && typeof rawValue.strategyKey === 'string' + ? { strategyKey: rawValue.strategyKey } + : {}), + ...('durationMs' in rawValue && typeof rawValue.durationMs === 'number' + ? { durationMs: rawValue.durationMs } + : {}), + } +} + +function readPreviewList( + value: unknown, +): Array | undefined { + if (!Array.isArray(value)) return undefined + const previews: Array = [] + for (const item of value) { + if (!item || typeof item !== 'object') continue + if (!('role' in item) || !('tokens' in item) || !('text' in item)) continue + if ( + typeof item.role !== 'string' || + typeof item.tokens !== 'number' || + typeof item.text !== 'string' + ) { + continue + } + previews.push({ + role: item.role, + tokens: item.tokens, + text: item.text, + }) + } + return previews +} + +function readCompactionStateValue(rawValue: unknown): { + before: number + after: number + messagesBefore: number + messagesAfter: number + reusedCheckpoint: boolean + maxTokens?: number + strategyKey?: string + dropped?: Array + result?: Array +} | null { + if (!rawValue || typeof rawValue !== 'object') return null + if ( + !('before' in rawValue) || + !('after' in rawValue) || + !('messagesBefore' in rawValue) || + !('messagesAfter' in rawValue) + ) { + return null + } + if ( + typeof rawValue.before !== 'number' || + typeof rawValue.after !== 'number' || + typeof rawValue.messagesBefore !== 'number' || + typeof rawValue.messagesAfter !== 'number' + ) { + return null + } + const reusedCheckpoint = + 'reusedCheckpoint' in rawValue && rawValue.reusedCheckpoint === true + const maxTokens = + 'maxTokens' in rawValue && typeof rawValue.maxTokens === 'number' + ? rawValue.maxTokens + : undefined + const strategyKey = + 'strategyKey' in rawValue && typeof rawValue.strategyKey === 'string' + ? rawValue.strategyKey + : undefined + const dropped = + 'dropped' in rawValue ? readPreviewList(rawValue.dropped) : undefined + const result = + 'result' in rawValue ? readPreviewList(rawValue.result) : undefined + return { + before: rawValue.before, + after: rawValue.after, + messagesBefore: rawValue.messagesBefore, + messagesAfter: rawValue.messagesAfter, + reusedCheckpoint, + ...(maxTokens !== undefined ? { maxTokens } : {}), + ...(strategyKey ? { strategyKey } : {}), + ...(dropped ? { dropped } : {}), + ...(result ? { result } : {}), + } +} + export interface AIDevtoolsClientMetadata extends AIDevtoolsDisplayOptions { framework?: string hookName: string @@ -712,6 +834,9 @@ export class ClientDevtoolsBridge { | 'memory:retrieve:started' | 'memory:retrieve:completed' | 'memory:snapshot' + | 'compaction:started' + | 'compaction:state' + | 'compaction:ended' | AIDevtoolsRunEventType, visibility: AIDevtoolsEventVisibility = 'client-state', context: { runId?: string } = {}, @@ -770,6 +895,9 @@ export class ChatDevtoolsBridge extends ClientDevtoolsBridge = + [] constructor(options: ChatDevtoolsBridgeOptions) { super({ @@ -926,10 +1054,62 @@ export class ChatDevtoolsBridge extends ClientDevtoolsBridge 60) { + this.lastCompactionEvents.splice(0, this.lastCompactionEvents.length - 60) + } + this.emitCompactionEvent(eventType, rawValue) + } + + recordCompactionState(rawValue: unknown): void { + this.recordCompactionEvent('compaction:state', rawValue) + } + + private emitCompactionEvent(eventType: string, rawValue: unknown): void { + const runContext = this.currentRunId ? { runId: this.currentRunId } : {} + if (eventType === 'compaction:started') { + const value = readCompactionBoundaryValue(rawValue) + emitAIDevtoolsEvent('compaction:started', { + ...this.createEnvelope( + 'compaction:started', + 'client-state', + runContext, + ), + ...value, + }) + return + } + if (eventType === 'compaction:ended') { + const value = readCompactionBoundaryValue(rawValue) + emitAIDevtoolsEvent('compaction:ended', { + ...this.createEnvelope('compaction:ended', 'client-state', runContext), + ...value, + }) + return + } + if (eventType === 'compaction:state') { + const value = readCompactionStateValue(rawValue) + if (!value) return + emitAIDevtoolsEvent('compaction:state', { + ...this.createEnvelope('compaction:state', 'client-state', runContext), + ...value, + }) + } + } + protected override onReplayState(): void { if (this.lastMemoryStateValue != null) { this.emitMemoryState(this.lastMemoryStateValue) } + for (const event of this.lastCompactionEvents) { + this.emitCompactionEvent(event.eventType, event.value) + } } getCurrentRunEventContext(): ChatClientRunEventContext | undefined { diff --git a/packages/ai-client/tests/devtools.test.ts b/packages/ai-client/tests/devtools.test.ts index 64198cf6d6..06a1bf6793 100644 --- a/packages/ai-client/tests/devtools.test.ts +++ b/packages/ai-client/tests/devtools.test.ts @@ -1423,6 +1423,127 @@ describe('ChatClient devtools bridge', () => { client.dispose() }) + it('re-emits compaction started, state, and ended from transported CUSTOM chunks', async () => { + const runContexts: Array = [] + const chunks: Array = [ + runStartedChunk({ threadId: 'thread-1', runId: 'run-cmp' }), + { + type: EventType.CUSTOM, + metadata: { tanstack: { model: 'test' } }, + timestamp: Date.now(), + name: 'compaction:started', + value: { + before: 400, + messagesBefore: 8, + reusedCheckpoint: false, + maxTokens: 400, + strategyKey: 'evict-oldest:half:maxTokens=400', + }, + }, + { + type: EventType.CUSTOM, + metadata: { tanstack: { model: 'test' } }, + timestamp: Date.now(), + name: 'compaction:state', + value: { + before: 400, + after: 180, + messagesBefore: 8, + messagesAfter: 3, + reusedCheckpoint: false, + maxTokens: 400, + strategyKey: 'evict-oldest:half:maxTokens=400', + dropped: [{ role: 'user', tokens: 40, text: 'old turn' }], + result: [{ role: 'user', tokens: 10, text: 'omitted' }], + }, + }, + { + type: EventType.CUSTOM, + metadata: { tanstack: { model: 'test' } }, + timestamp: Date.now(), + name: 'compaction:ended', + value: { + after: 180, + messagesAfter: 3, + reusedCheckpoint: false, + maxTokens: 400, + durationMs: 12, + strategyKey: 'evict-oldest:half:maxTokens=400', + }, + }, + textContentChunk({ + messageId: 'msg-cmp', + delta: 'ok', + content: 'ok', + }), + runFinishedChunk({ threadId: 'thread-1', runId: 'run-cmp' }), + ] + const client = createClient({ + connection: createRunTrackingAdapter([chunks], runContexts), + }) + vi.clearAllMocks() + + await client.sendMessage('keep going') + await waitForCondition( + () => eventClientMock.emitted('compaction:ended').length > 0, + ) + + expect(eventClientMock.emitted('compaction:started')).toEqual([ + [ + 'compaction:started', + expect.objectContaining({ + before: 400, + messagesBefore: 8, + reusedCheckpoint: false, + maxTokens: 400, + }), + ], + ]) + expect(eventClientMock.emitted('compaction:state')).toEqual([ + [ + 'compaction:state', + expect.objectContaining({ + before: 400, + after: 180, + messagesBefore: 8, + messagesAfter: 3, + reusedCheckpoint: false, + maxTokens: 400, + strategyKey: 'evict-oldest:half:maxTokens=400', + dropped: [{ role: 'user', tokens: 40, text: 'old turn' }], + result: [{ role: 'user', tokens: 10, text: 'omitted' }], + }), + ], + ]) + expect(eventClientMock.emitted('compaction:ended')).toEqual([ + [ + 'compaction:ended', + expect.objectContaining({ + after: 180, + messagesAfter: 3, + durationMs: 12, + }), + ], + ]) + + vi.clearAllMocks() + eventClientMock.dispatch('devtools:request-state', {}) + await waitForCondition( + () => eventClientMock.emitted('compaction:ended').length > 0, + ) + expect(eventClientMock.emitted('compaction:started')).toEqual([ + ['compaction:started', expect.objectContaining({ before: 400 })], + ]) + expect(eventClientMock.emitted('compaction:state')).toEqual([ + ['compaction:state', expect.objectContaining({ after: 180 })], + ]) + expect(eventClientMock.emitted('compaction:ended')).toEqual([ + ['compaction:ended', expect.objectContaining({ durationMs: 12 })], + ]) + + client.dispose() + }) + it('batches structured output update events while preserving final state', async () => { const runContexts: Array = [] const finalObject = { title: 'Pasta', servings: 2 } diff --git a/packages/ai-compaction/README.md b/packages/ai-compaction/README.md new file mode 100644 index 0000000000..6e7b6fb2d1 --- /dev/null +++ b/packages/ai-compaction/README.md @@ -0,0 +1,135 @@ +# @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` (default `assistant`) | +| `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. + +TanStack AI DevTools has a Compaction tab with started, state, and ended +events, before/after counts, and dropped vs sent message previews. Those +stats ride the chat stream as `compaction:started`, `compaction:state`, and +`compaction:ended` CUSTOM events. 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..f9d15eba80 --- /dev/null +++ b/packages/ai-compaction/src/index.test.ts @@ -0,0 +1,526 @@ +import { describe, expect, it, vi } from 'vitest' +import type { + ChatMiddlewareConfig, + ChatMiddlewareContext, + MetadataStore, + ModelMessage, + ToolCall, +} from '@tanstack/ai' +import { provideMetadata } from '@tanstack/ai' +import { + COMPACTION_ENDED_EVENT, + COMPACTION_STARTED_EVENT, + COMPACTION_STATE_EVENT, + clearToolResults, + composeStrategies, + estimateMessageTokens, + evictOldest, + summarizeOldest, + withCompaction, +} from './index' + +interface RecordedCustom { + name: string + value: Record +} + +function recordingContext( + phase: ChatMiddlewareContext['phase'] = 'beforeModel', + extras: Partial = {}, +): { ctx: ChatMiddlewareContext; events: Array } { + const events: Array = [] + // oxlint-disable-next-line eslint-js/no-restricted-syntax -- focused hook stub + const ctx = { + phase, + emitCustomEvent: (name: string, value: Record) => { + events.push({ name, value }) + }, + ...extras, + } as unknown as ChatMiddlewareContext + return { ctx, events } +} + +function runOnConfig( + mw: ReturnType, + messages: Array, + ctx: ChatMiddlewareContext = recordingContext().ctx, +) { + const config: ChatMiddlewareConfig = { + messages, + systemPrompts: [], + tools: [], + } + return mw.onConfig?.(ctx, config) +} + +function checkpointContext( + store: MetadataStore, + options: { aborted?: boolean; phase?: ChatMiddlewareContext['phase'] } = {}, +): ChatMiddlewareContext { + const recorded = recordingContext(options.phase) + // oxlint-disable-next-line eslint-js/no-restricted-syntax -- focused hook stub + const ctx = { + ...recorded.ctx, + threadId: 'thread-1', + signal: options.aborted ? AbortSignal.abort() : undefined, + capabilities: { markProvided: () => undefined }, + } as unknown as ChatMiddlewareContext + provideMetadata(ctx, store) + return ctx +} + +function memoryStore(): MetadataStore { + const values = new Map() + return { + 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}`) + }, + } +} + +function phaseContext( + phase: ChatMiddlewareContext['phase'], +): ChatMiddlewareContext { + return recordingContext(phase).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('emits started, state, and ended custom events when compacting', async () => { + const mw = withCompaction({ maxTokens: 100 }) + const { ctx, events } = recordingContext('beforeModel') + const msgs = [big('user'), big('assistant'), big('user'), big('assistant')] + await runOnConfig(mw, msgs, ctx) + expect(events.map((event) => event.name)).toEqual([ + COMPACTION_STARTED_EVENT, + COMPACTION_STATE_EVENT, + COMPACTION_ENDED_EVENT, + ]) + const stateValue = events[1]?.value + expect(stateValue).toMatchObject({ + reusedCheckpoint: false, + maxTokens: 100, + }) + expect(Array.isArray(stateValue?.dropped)).toBe(true) + expect(Array.isArray(stateValue?.result)).toBe(true) + expect( + Array.isArray(stateValue?.dropped) ? stateValue.dropped.length : 0, + ).toBeGreaterThan(0) + expect(typeof events[2]?.value.durationMs).toBe('number') + }) + + it('emits started before summarizeOldest finishes', async () => { + let release!: (summary: string) => void + const gate = new Promise((resolve) => { + release = resolve + }) + const mw = withCompaction({ + maxTokens: 100, + strategy: summarizeOldest({ summarize: () => gate }), + }) + const { ctx, events } = recordingContext('beforeModel') + const pending = runOnConfig( + mw, + [big('user'), big('assistant'), big('user'), big('assistant')], + ctx, + ) + await vi.waitFor(() => { + expect(events.map((event) => event.name)).toEqual([ + COMPACTION_STARTED_EVENT, + ]) + }) + release('the gist') + await pending + expect(events.map((event) => event.name)).toEqual([ + COMPACTION_STARTED_EVENT, + COMPACTION_STATE_EVENT, + COMPACTION_ENDED_EVENT, + ]) + }) + + it('does not emit custom events when under the token budget', async () => { + const mw = withCompaction({ maxTokens: 1000 }) + const { ctx, events } = recordingContext('beforeModel') + await runOnConfig(mw, [text('user', 'hi')], ctx) + expect(events).toEqual([]) + }) + + it('does not compact during init', async () => { + const onCompact = vi.fn() + const mw = withCompaction({ maxTokens: 100, onCompact }) + const result = await runOnConfig( + mw, + [big('user'), big('assistant'), big('user'), big('assistant')], + phaseContext('init'), + ) + expect(result).toBeUndefined() + expect(onCompact).not.toHaveBeenCalled() + }) + + it('reuses a persisted checkpoint for an unchanged canonical prefix', async () => { + const store = memoryStore() + 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 store = memoryStore() + 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) + }) + + it('keeps the trailing assistant plus tool result when the transcript ends in a tool', 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] + const mw = withCompaction({ + maxTokens: 100, + strategy: evictOldest({ keepRecentTokens: 45 }), + }) + const out = (await runOnConfig(mw, msgs))?.providerMessages ?? [] + expect(out.at(-2)).toBe(assistantCall) + expect(out.at(-1)).toBe(toolResult) + expect(out.some((m) => m.role === 'tool')).toBe(true) + }) + + it('keeps a trailing parallel tool-result group with its assistant', async () => { + const assistantCall: ModelMessage = { + role: 'assistant', + content: 'x'.repeat(160), + toolCalls: [ + call, + { + id: 't2', + type: 'function', + function: { name: 'g', arguments: '{}' }, + }, + ], + } + const toolA: ModelMessage = { + role: 'tool', + content: 'x'.repeat(160), + toolCallId: 't1', + } + const toolB: ModelMessage = { + role: 'tool', + content: 'x'.repeat(160), + toolCallId: 't2', + } + const msgs = [big('user'), assistantCall, toolA, toolB] + const mw = withCompaction({ + maxTokens: 100, + strategy: evictOldest({ keepRecentTokens: 45 }), + }) + const out = (await runOnConfig(mw, msgs))?.providerMessages ?? [] + expect(out.slice(-3)).toEqual([assistantCall, toolA, toolB]) + }) +}) + +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]?.role).toBe('assistant') + expect(result?.providerMessages?.[0]?.content).toBe( + '\nthe gist\n', + ) + }) + + it('reuses a checkpoint without an explicit strategyKey', async () => { + const store = memoryStore() + 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 }), + } + + await runOnConfig( + withCompaction(options), + messages, + checkpointContext(store), + ) + const second = await runOnConfig( + withCompaction(options), + [...messages, text('user', 'new')], + checkpointContext(store), + ) + + expect(summarize).toHaveBeenCalledOnce() + expect(second?.providerMessages?.at(-1)?.content).toBe('new') + }) +}) + +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..e467ecfa09 --- /dev/null +++ b/packages/ai-compaction/src/index.ts @@ -0,0 +1,618 @@ +/** + * `@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, + ChatMiddlewareContext, + ModelMessage, +} from '@tanstack/ai' + +/** CUSTOM stream event: compaction is about to run. */ +export const COMPACTION_STARTED_EVENT = 'compaction:started' +/** CUSTOM stream event: compaction result (counts and previews). */ +export const COMPACTION_STATE_EVENT = 'compaction:state' +/** CUSTOM stream event: compaction finished. */ +export const COMPACTION_ENDED_EVENT = 'compaction:ended' + +export type CompactionStreamEventName = + | typeof COMPACTION_STARTED_EVENT + | typeof COMPACTION_STATE_EVENT + | typeof COMPACTION_ENDED_EVENT + +const PREVIEW_CHARS = 4000 +const MAX_PREVIEWS = 24 + +/** One message in a `compaction:state` preview list. */ +export interface CompactionMessagePreview { + role: string + tokens: number + text: string +} + +/** Payload of {@link COMPACTION_STARTED_EVENT}. */ +export interface CompactionStartedEventValue { + before: number + messagesBefore: number + reusedCheckpoint: boolean + maxTokens: number + strategyKey?: string +} + +/** Payload of {@link COMPACTION_STATE_EVENT}. */ +export interface CompactionStateEventValue { + before: number + after: number + messagesBefore: number + messagesAfter: number + reusedCheckpoint: boolean + maxTokens: number + strategyKey?: string + /** Messages removed or rewritten. */ + dropped?: Array + /** Messages the model will see after compaction. */ + result?: Array +} + +/** Payload of {@link COMPACTION_ENDED_EVENT}. */ +export interface CompactionEndedEventValue { + after: number + messagesAfter: number + reusedCheckpoint: boolean + maxTokens: number + durationMs: number + strategyKey?: string +} + +function emitCompactionStarted( + ctx: ChatMiddlewareContext, + value: CompactionStartedEventValue, +) { + ctx.emitCustomEvent(COMPACTION_STARTED_EVENT, value) +} + +function emitCompactionState( + ctx: ChatMiddlewareContext, + value: CompactionStateEventValue, +) { + ctx.emitCustomEvent(COMPACTION_STATE_EVENT, value) +} + +function emitCompactionEnded( + ctx: ChatMiddlewareContext, + value: CompactionEndedEventValue, +) { + ctx.emitCustomEvent(COMPACTION_ENDED_EVENT, value) +} + +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) + ) +} + +function messagePreviewText(message: ModelMessage): string { + if (typeof message.content === 'string') return message.content + return JSON.stringify(message.content ?? '') +} + +function toMessagePreview( + message: ModelMessage, + estimate: (message: ModelMessage) => number, +): CompactionMessagePreview { + const text = messagePreviewText(message) + return { + role: message.role, + tokens: estimate(message), + text: + text.length > PREVIEW_CHARS ? `${text.slice(0, PREVIEW_CHARS)}…` : text, + } +} + +function previewList( + messages: ReadonlyArray, + estimate: (message: ModelMessage) => number, +): Array { + const mapped = messages.map((message) => toMessagePreview(message, estimate)) + if (mapped.length <= MAX_PREVIEWS) return mapped + return mapped.slice(0, MAX_PREVIEWS) +} + +function droppedMessages( + before: ReadonlyArray, + after: ReadonlyArray, +): Array { + const afterKeys = new Set(after.map((message) => JSON.stringify(message))) + return before.filter((message) => !afterKeys.has(JSON.stringify(message))) +} + +function compactionStateValue(args: { + before: number + after: number + messagesBefore: number + messagesAfter: number + reusedCheckpoint: boolean + maxTokens: number + strategyKey?: string + beforeMessages?: ReadonlyArray + afterMessages?: ReadonlyArray + estimate: (message: ModelMessage) => number +}): CompactionStateEventValue { + const value: CompactionStateEventValue = { + before: args.before, + after: args.after, + messagesBefore: args.messagesBefore, + messagesAfter: args.messagesAfter, + reusedCheckpoint: args.reusedCheckpoint, + maxTokens: args.maxTokens, + ...(args.strategyKey ? { strategyKey: args.strategyKey } : {}), + } + if (args.afterMessages) { + value.result = previewList(args.afterMessages, args.estimate) + } + if (args.beforeMessages && args.afterMessages) { + value.dropped = previewList( + droppedMessages(args.beforeMessages, args.afterMessages), + args.estimate, + ) + } + return value +} + +/** Rough token estimate for one message. Default: characters / 4. */ +export function estimateMessageTokens(message: ModelMessage): number { + let text = messagePreviewText(message) + 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++ + // Trailing tool results: skipping orphans would drop the whole tail (the + // normal agent-loop state). Keep those results and the message that owns them. + if (cut >= messages.length) { + cut = messages.length + while (cut > 0 && messages[cut - 1]?.role === 'tool') cut-- + if (cut > 0) 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) + // 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 `'assistant'`. */ + summaryRole?: 'user' | 'assistant' +}): CompactionStrategy { + const strategy: CompactionStrategy = 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 ?? 'assistant', + content: `\n${summary}\n`, + }, + ...messages.slice(cut), + ] + } + return identifyStrategy( + strategy, + `summarize-oldest:${options.keepRecentTokens ?? 'half'}:${options.summaryRole ?? 'assistant'}`, + ) +} + +/** + * 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) { + // init is discarded by the engine rebuild and can run before persistence + // hydrates the thread. Compact only on model-bound phases. + if (ctx.phase === 'init') return + + const startedAt = Date.now() + 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) + const startedValue: CompactionStartedEventValue = { + before, + messagesBefore: workingMessages.length, + reusedCheckpoint, + maxTokens: options.maxTokens, + ...(checkpointStrategyKey + ? { strategyKey: checkpointStrategyKey } + : {}), + } + + if (before <= options.maxTokens) { + if (reusedCheckpoint) { + emitCompactionStarted(ctx, startedValue) + const stateValue = compactionStateValue({ + before, + after: before, + messagesBefore: workingMessages.length, + messagesAfter: workingMessages.length, + reusedCheckpoint: true, + maxTokens: options.maxTokens, + strategyKey: checkpointStrategyKey, + afterMessages: workingMessages, + estimate, + }) + emitCompactionState(ctx, stateValue) + emitCompactionEnded(ctx, { + after: before, + messagesAfter: workingMessages.length, + reusedCheckpoint: true, + maxTokens: options.maxTokens, + durationMs: Date.now() - startedAt, + ...(checkpointStrategyKey + ? { strategyKey: checkpointStrategyKey } + : {}), + }) + return { providerMessages: workingMessages } + } + return + } + + emitCompactionStarted(ctx, startedValue) + const next = await strategy(workingMessages, { + maxTokens: options.maxTokens, + estimate, + }) + if (!next || next === workingMessages) { + emitCompactionEnded(ctx, { + after: before, + messagesAfter: workingMessages.length, + reusedCheckpoint, + maxTokens: options.maxTokens, + durationMs: Date.now() - startedAt, + ...(checkpointStrategyKey + ? { strategyKey: checkpointStrategyKey } + : {}), + }) + if (reusedCheckpoint) { + return { providerMessages: workingMessages } + } + return + } + + const info = { + before, + after: sum(next, estimate), + messagesBefore: workingMessages.length, + messagesAfter: next.length, + } + options.onCompact?.(info) + emitCompactionState( + ctx, + compactionStateValue({ + ...info, + reusedCheckpoint, + maxTokens: options.maxTokens, + strategyKey: checkpointStrategyKey, + beforeMessages: workingMessages, + afterMessages: next, + estimate, + }), + ) + emitCompactionEnded(ctx, { + after: info.after, + messagesAfter: info.messagesAfter, + reusedCheckpoint, + maxTokens: options.maxTokens, + durationMs: Date.now() - startedAt, + ...(checkpointStrategyKey + ? { strategyKey: checkpointStrategyKey } + : {}), + }) + + 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-devtools/src/components/conversation/IterationCard.tsx b/packages/ai-devtools/src/components/conversation/IterationCard.tsx index 1654746eb8..2d46141979 100644 --- a/packages/ai-devtools/src/components/conversation/IterationCard.tsx +++ b/packages/ai-devtools/src/components/conversation/IterationCard.tsx @@ -169,6 +169,9 @@ const MiddlewareStep: Component<{ if (ev().wasDropped) return 'DROP' if (ev().hookName === 'onChunk' && ev().hasTransform) return 'TRANSFORM' if (ev().hookName === 'onConfig' && ev().hasTransform) return 'TRANSFORM' + if (ev().hookName === 'onCompactStart') return 'START' + if (ev().hookName === 'onCompact') return 'COMPACT' + if (ev().hookName === 'onCompactEnd') return 'END' if (ev().hookName === 'onBeforeToolCall' && ev().hasTransform) return 'DECISION' return null @@ -181,7 +184,13 @@ const MiddlewareStep: Component<{ return ( <> -
+
{ + if (hasChanges()) setExpanded(!expanded()) + }} + > Middleware @@ -196,10 +205,7 @@ const MiddlewareStep: Component<{ {suffix()} - setExpanded(!expanded())} - > + {expanded() ? 'hide changes' : 'show changes'} diff --git a/packages/ai-devtools/src/components/hooks/CompactionPanel.tsx b/packages/ai-devtools/src/components/hooks/CompactionPanel.tsx new file mode 100644 index 0000000000..e18fcb54a2 --- /dev/null +++ b/packages/ai-devtools/src/components/hooks/CompactionPanel.tsx @@ -0,0 +1,223 @@ +import { For, Show, createMemo, createSignal } from 'solid-js' +import { JsonTree } from '@tanstack/devtools-ui' +import { useAIStore } from '../../store/ai-context' +import { useStyles } from '../../styles/use-styles' +import { compactionEventsForHook } from '../../store/compaction-registry' +import type { Component } from 'solid-js' +import type { + CompactionEventRecord, + CompactionMessagePreview, +} from '../../store/compaction-registry' +import type { HookRecord } from '../../store/hook-registry' + +function formatTime(value: number): string { + const date = new Date(value) + if (Number.isNaN(date.getTime())) return '' + return date.toLocaleTimeString() +} + +function shortStrategy(event: CompactionEventRecord): string { + const key = event.strategyKey + if (!key) return event.reusedCheckpoint ? 'checkpoint' : 'compaction' + const cut = key.indexOf(':maxTokens=') + return cut === -1 ? key : key.slice(0, cut) +} + +const PreviewMessage: Component<{ + preview: CompactionMessagePreview +}> = (props) => { + const styles = useStyles() + return ( +
+
+ {props.preview.role} + 0}> + {` · ${props.preview.tokens} tok`} + +
+
+ {props.preview.text} +
+
+ ) +} + +const PreviewColumn: Component<{ + title: string + previews: Array + empty: string +}> = (props) => { + const styles = useStyles() + const s = () => styles().iterationTimeline + return ( +
+
+ {props.title} ({props.previews.length}) +
+ 0} + fallback={ +
{props.empty}
+ } + > +
+ + {(preview) => } + +
+ +
+ ) +} + +function kindLabel(kind: CompactionEventRecord['kind']): string { + if (kind === 'started') return 'Started' + if (kind === 'ended') return 'Ended' + return 'State' +} + +const CompactEvent: Component<{ + event: CompactionEventRecord + defaultOpen?: boolean +}> = (props) => { + const styles = useStyles() + const s = () => styles().iterationTimeline + const [expanded, setExpanded] = createSignal(props.defaultOpen === true) + const event = () => props.event + const isState = () => event().kind === 'state' + + const stats = () => ({ + kind: event().kind, + before: event().before, + after: event().after, + messagesBefore: event().messagesBefore, + messagesAfter: event().messagesAfter, + maxTokens: event().maxTokens, + strategyKey: event().strategyKey, + reusedCheckpoint: event().reusedCheckpoint, + durationMs: event().durationMs, + }) + + const countLabel = () => { + if (event().kind === 'started') { + return `${event().messagesBefore ?? 0} msgs · ${event().before ?? 0} tok` + } + if (event().kind === 'ended') { + return `${event().messagesAfter ?? 0} msgs · ${event().after ?? 0} tok` + } + return `${event().messagesBefore ?? 0} → ${event().messagesAfter ?? 0} msgs` + } + + return ( + <> +
setExpanded(!expanded())} + > + + {kindLabel(event().kind)} + + + {shortStrategy(event())} + + {countLabel()} + + + {event().before} → {event().after} tok + + + + {event().durationMs}ms + + + checkpoint + + {formatTime(event().timestamp)} + + {'\u25B6'} + +
+ +
+ +
+ +
+ + +
+ + + + ) +} + +export const CompactionPanel: Component<{ hook: HookRecord }> = (props) => { + const { state, clearCompaction } = useAIStore() + const styles = useStyles() + const s = () => styles().iterationTimeline + + const events = createMemo(() => + compactionEventsForHook(state.compaction, props.hook), + ) + const lastStateIndex = createMemo(() => { + const list = events() + for (let i = list.length - 1; i >= 0; i--) { + if (list[i]?.kind === 'state') return i + } + return -1 + }) + + return ( +
+ 0} + fallback={ +
+ No compaction yet. When the transcript passes maxTokens, each + compact appears here. +
+ } + > +
+ + {events().length} event{events().length === 1 ? '' : 's'} + + +
+
+ + {(event, index) => ( + + )} + +
+ +
+ ) +} diff --git a/packages/ai-devtools/src/components/hooks/HookDashboard.tsx b/packages/ai-devtools/src/components/hooks/HookDashboard.tsx index 8c5bf4f329..dcbc441232 100644 --- a/packages/ai-devtools/src/components/hooks/HookDashboard.tsx +++ b/packages/ai-devtools/src/components/hooks/HookDashboard.tsx @@ -55,8 +55,9 @@ export const HookDashboard: Component = () => { const handleSelect = (hook: HookRecord) => { selectHook(hook.id) - if (state.conversations[hook.id]) { - selectConversation(hook.id) + const conversationId = hook.id || hook.clientId || hook.threadId + if (conversationId && state.conversations[conversationId]) { + selectConversation(conversationId) } } diff --git a/packages/ai-devtools/src/components/hooks/HookDetails.tsx b/packages/ai-devtools/src/components/hooks/HookDetails.tsx index 8ae7860e04..143607c26c 100644 --- a/packages/ai-devtools/src/components/hooks/HookDetails.tsx +++ b/packages/ai-devtools/src/components/hooks/HookDetails.tsx @@ -39,6 +39,7 @@ import { } from './preview-messages' import { GenerationPanel, GenerationPreview } from './GenerationPanel' import { MemoryPanel } from './MemoryPanel' +import { CompactionPanel } from './CompactionPanel' import type { HoverOrigin, HoverTarget, PreviewJsonItem } from './preview-model' import type { HookRecord, @@ -49,7 +50,7 @@ import type { import type { Conversation, Message, ToolCall } from '../../store/ai-store' import type { Component, Setter } from 'solid-js' -type DetailTab = 'conversation' | 'tools' | 'state' | 'memory' +type DetailTab = 'conversation' | 'tools' | 'state' | 'memory' | 'compaction' type MessagePart = NonNullable[number] const scrollAnimations = new WeakMap() @@ -119,7 +120,8 @@ export const HookDetails: Component = () => { const hook = createMemo((): HookRecord | undefined => { const id = state.hooks.activeHookId - return id ? state.hooks.hooks[id] : undefined + if (id == null) return undefined + return state.hooks.hooks[id] }) const conversation = createMemo(() => { @@ -151,7 +153,9 @@ export const HookDetails: Component = () => { // while one of them is selected, fall back to the conversation view. if ( isGenerationHook() && - (activeTab() === 'tools' || activeTab() === 'memory') + (activeTab() === 'tools' || + activeTab() === 'memory' || + activeTab() === 'compaction') ) { setActiveTab('conversation') } @@ -204,8 +208,10 @@ export const HookDetails: Component = () => { { selectHook(selectedHook.id) - if (state.conversations[selectedHook.id]) { - selectConversation(selectedHook.id) + const conversationId = + selectedHook.id || selectedHook.clientId || selectedHook.threadId + if (conversationId && state.conversations[conversationId]) { + selectConversation(conversationId) } }} /> @@ -255,6 +261,12 @@ export const HookDetails: Component = () => { activeTab={activeTab()} onSelect={setActiveTab} /> + @@ -302,6 +314,9 @@ export const HookDetails: Component = () => { + + + diff --git a/packages/ai-devtools/src/components/hooks/index.ts b/packages/ai-devtools/src/components/hooks/index.ts index 73a7fafb70..6dd1a19143 100644 --- a/packages/ai-devtools/src/components/hooks/index.ts +++ b/packages/ai-devtools/src/components/hooks/index.ts @@ -2,4 +2,5 @@ export { HookDashboard } from './HookDashboard' export { HookDetails } from './HookDetails' export { GenerationPanel, GenerationPreview } from './GenerationPanel' export { MemoryPanel } from './MemoryPanel' +export { CompactionPanel } from './CompactionPanel' export { ToolFixtureForm } from './ToolFixtureForm' diff --git a/packages/ai-devtools/src/store/ai-context.tsx b/packages/ai-devtools/src/store/ai-context.tsx index 3a95463587..5c9f5cf2b8 100644 --- a/packages/ai-devtools/src/store/ai-context.tsx +++ b/packages/ai-devtools/src/store/ai-context.tsx @@ -6,9 +6,9 @@ import { applyHookEvent, clearHookRegistry, createHookRegistryState, + markHookViewed, removeSavedFixture, replaceSavedFixtures, - setActiveHook, } from './hook-registry' import { createClientToolCallMessage, @@ -20,6 +20,11 @@ import { clearMemoryRegistry, createMemoryRegistryState, } from './memory-registry' +import { + applyCompactionEvent, + clearCompactionRegistry, + createCompactionRegistryState, +} from './compaction-registry' import type { ContentPartSource, TokenUsage } from '@tanstack/ai' import type { DevtoolsToolFixtureApplyEvent, @@ -27,6 +32,7 @@ import type { } from '@tanstack/ai-event-client' import type { HookRegistryState, ToolFixtureRecord } from './hook-registry' import type { MemoryRegistryState } from './memory-registry' +import type { CompactionRegistryState } from './compaction-registry' import type { ParentComponent } from 'solid-js' interface MessagePart { @@ -235,6 +241,7 @@ interface AIStoreState { activeConversationId: string | null hooks: HookRegistryState memory: MemoryRegistryState + compaction: CompactionRegistryState } interface AIContextValue { @@ -243,6 +250,7 @@ interface AIContextValue { selectConversation: (id: string) => void clearHooks: () => void clearMemory: () => void + clearCompaction: () => void selectHook: (id: string | null) => void saveToolFixture: (fixture: ToolFixtureRecord) => void deleteToolFixture: (fixtureId: string) => void @@ -265,6 +273,7 @@ export const AIProvider: ParentComponent = (props) => { activeConversationId: null, hooks: createHookRegistryState(), memory: createMemoryRegistryState(), + compaction: createCompactionRegistryState(), }) const streamToConversation = new Map() @@ -660,15 +669,27 @@ export const AIProvider: ParentComponent = (props) => { ) } - function selectHook(id: string | null) { + function clearCompaction() { setState( - 'hooks', - produce((hooks: HookRegistryState) => { - setActiveHook(hooks, id) + 'compaction', + produce((compaction: CompactionRegistryState) => { + clearCompactionRegistry(compaction) }), ) } + function selectHook(id: string | null) { + setState('hooks', 'activeHookId', id) + if (id) { + setState( + 'hooks', + produce((hooks: HookRegistryState) => { + markHookViewed(hooks, id) + }), + ) + } + } + function saveToolFixture(fixture: ToolFixtureRecord) { const fixtures = mergeFixtures(fixture, state.hooks.fixtures) setState( @@ -2957,6 +2978,90 @@ export const AIProvider: ParentComponent = (props) => { }), ) + const recordCompactionLifecycle = ( + kind: 'started' | 'state' | 'ended', + hookName: 'onCompactStart' | 'onCompact' | 'onCompactEnd', + payload: { + timestamp: number + requestId?: string + streamId?: string + clientId?: string + hookId?: string + threadId?: string + runId?: string + eventId?: string + before?: number + after?: number + messagesBefore?: number + messagesAfter?: number + reusedCheckpoint?: boolean + maxTokens?: number + strategyKey?: string + durationMs?: number + dropped?: CompactionRegistryState['events'][number]['dropped'] + result?: CompactionRegistryState['events'][number]['result'] + }, + ) => { + setState( + 'compaction', + produce((compaction: CompactionRegistryState) => { + applyCompactionEvent(compaction, kind, payload) + }), + ) + + const { requestId, streamId, clientId } = payload + const conversationId = + (clientId && state.conversations[clientId] ? clientId : undefined) || + (streamId ? streamToConversation.get(streamId) : undefined) || + (requestId ? requestToConversation.get(requestId) : undefined) + if (!conversationId || !state.conversations[conversationId]) return + + const conv = state.conversations[conversationId] + const iterIndex = findLatestIterationIndex(conv, requestId) + if (iterIndex < 0) return + + const mwEvent: MiddlewareEvent = { + id: `mw-cmp-${kind}-${Date.now()}-${Math.random()}`, + middlewareName: 'compaction', + hookName, + timestamp: payload.timestamp, + hasTransform: kind === 'state', + configChanges: { + before: payload.before, + after: payload.after, + messagesBefore: payload.messagesBefore, + messagesAfter: payload.messagesAfter, + reusedCheckpoint: payload.reusedCheckpoint, + maxTokens: payload.maxTokens, + strategyKey: payload.strategyKey, + durationMs: payload.durationMs, + }, + } + + setState( + 'conversations', + conversationId, + 'iterations', + iterIndex, + 'middlewareEvents', + produce((arr: Array) => { + arr.push(mwEvent) + }), + ) + } + + cleanupFns.push( + aiEventClient.on('compaction:started', (e) => { + recordCompactionLifecycle('started', 'onCompactStart', e.payload) + }), + aiEventClient.on('compaction:state', (e) => { + recordCompactionLifecycle('state', 'onCompact', e.payload) + }), + aiEventClient.on('compaction:ended', (e) => { + recordCompactionLifecycle('ended', 'onCompactEnd', e.payload) + }), + ) + cleanupFns.push( aiEventClient.on('summarize:request:started', (e) => { const { requestId, model, inputLength, timestamp, clientId } = e.payload @@ -3464,6 +3569,7 @@ export const AIProvider: ParentComponent = (props) => { selectConversation, clearHooks, clearMemory, + clearCompaction, selectHook, saveToolFixture, deleteToolFixture, diff --git a/packages/ai-devtools/src/store/compaction-registry.ts b/packages/ai-devtools/src/store/compaction-registry.ts new file mode 100644 index 0000000000..97ac7e2549 --- /dev/null +++ b/packages/ai-devtools/src/store/compaction-registry.ts @@ -0,0 +1,121 @@ +const MAX_EVENTS = 50 + +export type CompactionEventKind = 'started' | 'state' | 'ended' + +export interface CompactionMessagePreview { + role: string + tokens: number + text: string +} + +export interface CompactionEventRecord { + id: string + kind: CompactionEventKind + timestamp: number + hookId?: string + clientId?: string + threadId?: string + runId?: string + before?: number + after?: number + messagesBefore?: number + messagesAfter?: number + reusedCheckpoint?: boolean + maxTokens?: number + strategyKey?: string + durationMs?: number + dropped?: Array + result?: Array +} + +export interface CompactionRegistryState { + events: Array +} + +export function createCompactionRegistryState(): CompactionRegistryState { + return { events: [] } +} + +export interface CompactionLifecycleInput { + timestamp: number + eventId?: string + hookId?: string + clientId?: string + threadId?: string + runId?: string + before?: number + after?: number + messagesBefore?: number + messagesAfter?: number + reusedCheckpoint?: boolean + maxTokens?: number + strategyKey?: string + durationMs?: number + dropped?: Array + result?: Array +} + +function eventId( + kind: CompactionEventKind, + event: CompactionLifecycleInput, +): string { + return [ + 'cmp', + kind, + String(event.timestamp), + event.hookId ?? event.clientId ?? event.threadId ?? 'unknown', + String(event.before ?? event.after ?? ''), + ].join('-') +} + +/** Append one compaction lifecycle event. Newest stays at the end. */ +export function applyCompactionEvent( + state: CompactionRegistryState, + kind: CompactionEventKind, + event: CompactionLifecycleInput, +): void { + state.events.push({ + id: event.eventId ?? eventId(kind, event), + kind, + timestamp: event.timestamp, + hookId: event.hookId, + clientId: event.clientId, + threadId: event.threadId, + runId: event.runId, + before: event.before, + after: event.after, + messagesBefore: event.messagesBefore, + messagesAfter: event.messagesAfter, + reusedCheckpoint: event.reusedCheckpoint, + maxTokens: event.maxTokens, + strategyKey: event.strategyKey, + durationMs: event.durationMs, + dropped: event.dropped, + result: event.result, + }) + if (state.events.length > MAX_EVENTS) { + state.events.splice(0, state.events.length - MAX_EVENTS) + } +} + +export function clearCompactionRegistry(state: CompactionRegistryState): void { + state.events = [] +} + +export function compactionEventsForHook( + state: CompactionRegistryState, + hook: { id: string; clientId?: string; threadId?: string }, +): Array { + const matched = state.events.filter((event) => { + if (event.hookId && event.hookId === hook.id) return true + if (event.clientId && hook.clientId && event.clientId === hook.clientId) { + return true + } + if (event.threadId && hook.threadId && event.threadId === hook.threadId) { + return true + } + return false + }) + if (matched.length > 0) return matched + return state.events +} diff --git a/packages/ai-devtools/src/styles/use-styles.ts b/packages/ai-devtools/src/styles/use-styles.ts index 43e003d9ab..82c1b77914 100644 --- a/packages/ai-devtools/src/styles/use-styles.ts +++ b/packages/ai-devtools/src/styles/use-styles.ts @@ -3652,6 +3652,7 @@ const stylesFactory = (theme: 'light' | 'dark') => { grid-template-columns: repeat(2, minmax(320px, 1fr)); gap: ${size[3]}; align-items: start; + padding-left: ${size[3]}; padding-right: ${size[3]}; @media (max-width: 760px) { grid-template-columns: 1fr; diff --git a/packages/ai-devtools/tests/compaction-registry.test.ts b/packages/ai-devtools/tests/compaction-registry.test.ts new file mode 100644 index 0000000000..9e8146641f --- /dev/null +++ b/packages/ai-devtools/tests/compaction-registry.test.ts @@ -0,0 +1,103 @@ +import { describe, expect, it } from 'vitest' +import { + applyCompactionEvent, + clearCompactionRegistry, + compactionEventsForHook, + createCompactionRegistryState, +} from '../src/store/compaction-registry' +import type { CompactionLifecycleInput } from '../src/store/compaction-registry' + +function input( + overrides: Partial & { timestamp?: number }, +): CompactionLifecycleInput { + return { + timestamp: 1, + messagesBefore: 8, + messagesAfter: 3, + reusedCheckpoint: false, + ...overrides, + } +} + +describe('compaction registry', () => { + it('accumulates started, state, and ended events', () => { + const state = createCompactionRegistryState() + applyCompactionEvent( + state, + 'started', + input({ + before: 400, + messagesBefore: 8, + timestamp: 10, + hookId: 'hook-1', + }), + ) + applyCompactionEvent( + state, + 'state', + input({ + before: 400, + after: 180, + hookId: 'hook-1', + timestamp: 11, + strategyKey: 'evict-oldest:half:maxTokens=400', + maxTokens: 400, + dropped: [{ role: 'user', tokens: 40, text: 'old' }], + result: [{ role: 'user', tokens: 10, text: 'omitted' }], + }), + ) + applyCompactionEvent( + state, + 'ended', + input({ + after: 180, + messagesAfter: 3, + durationMs: 12, + hookId: 'hook-1', + timestamp: 12, + }), + ) + expect(state.events.map((event) => event.kind)).toEqual([ + 'started', + 'state', + 'ended', + ]) + expect(state.events[1]).toMatchObject({ + kind: 'state', + before: 400, + after: 180, + strategyKey: 'evict-oldest:half:maxTokens=400', + }) + expect(state.events[1]?.dropped?.[0]?.text).toBe('old') + expect(state.events[2]?.durationMs).toBe(12) + }) + + it('filters events for a hook and falls back to all', () => { + const state = createCompactionRegistryState() + applyCompactionEvent( + state, + 'state', + input({ before: 1, after: 1, hookId: 'a', timestamp: 1 }), + ) + applyCompactionEvent( + state, + 'state', + input({ before: 2, after: 1, hookId: 'b', timestamp: 2 }), + ) + expect( + compactionEventsForHook(state, { id: 'b' }).map((event) => event.hookId), + ).toEqual(['b']) + expect( + compactionEventsForHook(state, { id: 'missing' }).map( + (event) => event.hookId, + ), + ).toEqual(['a', 'b']) + }) + + it('clears the registry', () => { + const state = createCompactionRegistryState() + applyCompactionEvent(state, 'started', input({ before: 1 })) + clearCompactionRegistry(state) + expect(state.events).toEqual([]) + }) +}) diff --git a/packages/ai-event-client/src/index.ts b/packages/ai-event-client/src/index.ts index 37546bbcfe..b184b97341 100644 --- a/packages/ai-event-client/src/index.ts +++ b/packages/ai-event-client/src/index.ts @@ -979,6 +979,61 @@ export interface VideoUsageEvent extends BaseEventContext { usage: TokenUsage } +// --------------------------------------------------------------------------- +// Compaction events +// --------------------------------------------------------------------------- + +/** One message in a compaction preview list. */ +export interface CompactionMessagePreview { + role: string + tokens: number + text: string +} + +/** Emitted when `withCompaction` starts rewriting provider context. */ +export interface CompactionStartedEvent extends BaseEventContext { + before?: number + messagesBefore?: number + reusedCheckpoint?: boolean + maxTokens?: number + strategyKey?: string +} + +/** Emitted when `withCompaction` has a compacted transcript to inspect. */ +export interface CompactionStateEvent extends BaseEventContext { + /** Estimated tokens before compaction. */ + before: number + /** Estimated tokens after compaction. */ + after: number + /** Message count before compaction. */ + messagesBefore: number + /** Message count after compaction. */ + messagesAfter: number + /** True when a persisted checkpoint supplied the compacted prefix. */ + reusedCheckpoint: boolean + /** Token budget that triggered compaction. */ + maxTokens?: number + /** Strategy identity from `withCompaction`. */ + strategyKey?: string + /** Messages removed or rewritten. */ + dropped?: Array + /** Messages the model will see after compaction. */ + result?: Array +} + +/** Emitted when `withCompaction` finishes rewriting provider context. */ +export interface CompactionEndedEvent extends BaseEventContext { + after?: number + messagesAfter?: number + reusedCheckpoint?: boolean + maxTokens?: number + strategyKey?: string + durationMs?: number +} + +/** @deprecated Use {@link CompactionStateEvent}. */ +export type CompactionAppliedEvent = CompactionStateEvent + // --------------------------------------------------------------------------- // Memory events // --------------------------------------------------------------------------- @@ -1322,6 +1377,12 @@ export interface AIDevtoolsEventMap { 'client:reloaded': ClientReloadedEvent 'client:stopped': ClientStoppedEvent + // Compaction events + 'compaction:started': CompactionStartedEvent + 'compaction:state': CompactionStateEvent + 'compaction:ended': CompactionEndedEvent + 'compaction:applied': CompactionAppliedEvent + // Memory events 'memory:retrieve:started': MemoryRetrieveStartedEvent 'memory:retrieve:completed': MemoryRetrieveCompletedEvent 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-sandbox/tests/fakes.ts b/packages/ai-sandbox/tests/fakes.ts index c38f2fa6b8..a0e33284c6 100644 --- a/packages/ai-sandbox/tests/fakes.ts +++ b/packages/ai-sandbox/tests/fakes.ts @@ -264,6 +264,7 @@ export function makeMiddlewareCtx(input: { chunkIndex: 0, signal: controller.signal, abort: (reason) => controller.abort(reason), + emitCustomEvent: () => {}, context: {}, defer: () => {}, activity: 'chat', diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 91f84d3490..6c0aaa4af6 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 @@ -846,6 +847,9 @@ class TextEngine< > private readonly middlewareCtx: ChatMiddlewareContext private readonly sandboxFileQueue: Array = [] + private readonly middlewareCustomQueue: Array = [] + private middlewareCustomWaiters: Array<() => void> = [] + private drainingMiddlewareCustom = false private readonly deferredPromises: Array> = [] private abortReason?: string private readonly middlewareAbortController?: AbortController @@ -936,6 +940,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 || []) @@ -994,6 +999,14 @@ class TextEngine< this.abortReason = reason this.middlewareAbortController?.abort(reason) }, + emitCustomEvent: (name, value) => { + this.middlewareCustomQueue.push( + this.createCustomEventChunk(name, value), + ) + const waiters = this.middlewareCustomWaiters + this.middlewareCustomWaiters = [] + for (const waiter of waiters) waiter() + }, context: config.context as TContext, defer: (promise: Promise) => { this.deferredPromises.push(promise) @@ -1129,21 +1142,24 @@ class TextEngine< try { // Provision capabilities before any consumer (onConfig onward) can read them - await this.middlewareRunner.runSetup(this.middlewareCtx) + yield* this.runWhileYielding( + this.middlewareRunner.runSetup(this.middlewareCtx), + ) // Run initial onConfig (phase = init) this.middlewareCtx.phase = 'init' const initialConfig = this.buildMiddlewareConfig() - const transformedConfig = await this.middlewareRunner.runOnConfig( - this.middlewareCtx, - initialConfig, + const transformedConfig = yield* this.runWhileYielding( + this.middlewareRunner.runOnConfig(this.middlewareCtx, initialConfig), ) this.applyMiddlewareConfig(transformedConfig) await this.applyEphemeralInterruptResume(transformedConfig) await this.applyDurableGenericInterruptResolution() // Run onStart (devtools middleware emits text:request:started and initial messages here) - await this.middlewareRunner.runOnStart(this.middlewareCtx) + yield* this.runWhileYielding( + this.middlewareRunner.runOnStart(this.middlewareCtx), + ) if (this.earlyTermination) { yield* this.emitSuccessfulEarlyTermination() @@ -1190,18 +1206,16 @@ class TextEngine< iteration: this.middlewareCtx.iteration, }) - await this.beginCycle() + yield* this.runWhileYielding(this.beginCycle()) if (this.cyclePhase === 'processText') { // Run onConfig before each model call (phase = beforeModel) this.middlewareCtx.phase = 'beforeModel' this.middlewareCtx.iteration = this.iterationCount const iterConfig = this.buildMiddlewareConfig() - const iterTransformedConfig = - await this.middlewareRunner.runOnConfig( - this.middlewareCtx, - iterConfig, - ) + const iterTransformedConfig = yield* this.runWhileYielding( + this.middlewareRunner.runOnConfig(this.middlewareCtx, iterConfig), + ) this.applyMiddlewareConfig(iterTransformedConfig) if ( @@ -1240,7 +1254,7 @@ class TextEngine< } this.endCycle() - } while (await this.shouldContinue()) + } while (yield* this.runWhileYielding(this.shouldContinue())) } this.logger.agentLoop('run finished', { @@ -1498,7 +1512,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, @@ -1633,6 +1647,7 @@ class TextEngine< continue } if (spec.type === EventType.RUN_STARTED) { + if (this.hasPublicRunStarted) continue this.hasPublicRunStarted = true } this.logger.output(`type=${spec.type}`, { chunk: spec }) @@ -1647,6 +1662,7 @@ class TextEngine< // Drain any sandbox.file events emitted while processing this chunk. yield* this.drainSandboxFileQueue() + yield* this.drainMiddlewareCustomQueue() if (this.earlyTermination) { break @@ -1655,6 +1671,7 @@ class TextEngine< // Drain any remaining sandbox.file events emitted after the stream ended. yield* this.drainSandboxFileQueue() + yield* this.drainMiddlewareCustomQueue() } private handleStreamChunk(chunk: AdapterYieldChunk): void { @@ -2048,12 +2065,14 @@ class TextEngine< const allResults = [...executionResult.results, ...deferredErrorResults] // Notify middleware of tool phase completion (devtools emits aggregate events here) - await this.middlewareRunner.runOnToolPhaseComplete(this.middlewareCtx, { - toolCalls: pendingToolCalls, - results: allResults, - needsApproval: executionResult.needsApproval, - needsClientExecution: executionResult.needsClientExecution, - }) + yield* this.runWhileYielding( + this.middlewareRunner.runOnToolPhaseComplete(this.middlewareCtx, { + toolCalls: pendingToolCalls, + results: allResults, + needsApproval: executionResult.needsApproval, + needsClientExecution: executionResult.needsClientExecution, + }), + ) if ( executionResult.needsApproval.length > 0 || @@ -2229,23 +2248,26 @@ class TextEngine< const allResults = [...executionResult.results, ...deferredErrorResults] // Notify middleware of tool phase completion (devtools emits aggregate events here) - await this.middlewareRunner.runOnToolPhaseComplete(this.middlewareCtx, { - toolCalls, - results: allResults, - needsApproval: executionResult.needsApproval, - needsClientExecution: executionResult.needsClientExecution, - }) + yield* this.runWhileYielding( + this.middlewareRunner.runOnToolPhaseComplete(this.middlewareCtx, { + toolCalls, + results: allResults, + needsApproval: executionResult.needsApproval, + needsClientExecution: executionResult.needsClientExecution, + }), + ) const afterToolBoundaryChunks = this.buildToolResultChunks( allResults, finishEvent, ) - const afterToolRequests = - await this.middlewareRunner.runOnInterruptBoundary( + const afterToolRequests = yield* this.runWhileYielding( + this.middlewareRunner.runOnInterruptBoundary( this.middlewareCtx as ChatMiddlewareContext & { phase: 'afterTools' }, - ) + ), + ) if (afterToolRequests.length > 0) { for (const chunk of afterToolBoundaryChunks) { yield* this.pipeThroughMiddleware(chunk) @@ -2925,10 +2947,12 @@ class TextEngine< this.middlewareCtx.phase = phase const boundaryRequests = requests ?? - (await this.middlewareRunner.runOnInterruptBoundary( - this.middlewareCtx as ChatMiddlewareContext & { - phase: typeof phase - }, + (yield* this.runWhileYielding( + this.middlewareRunner.runOnInterruptBoundary( + this.middlewareCtx as ChatMiddlewareContext & { + phase: typeof phase + }, + ), )) if (boundaryRequests.length === 0) return false for (const request of boundaryRequests) { @@ -3446,9 +3470,11 @@ class TextEngine< } // 1) onStructuredOutputConfig — middleware can transform messages, options, outputSchema - structuredConfig = await this.middlewareRunner.runOnStructuredOutputConfig( - this.middlewareCtx, - structuredConfig, + structuredConfig = yield* this.runWhileYielding( + this.middlewareRunner.runOnStructuredOutputConfig( + this.middlewareCtx, + structuredConfig, + ), ) // 2) onConfig — phase-aware general-purpose middleware re-runs at the @@ -3457,9 +3483,11 @@ class TextEngine< // call — same constraint applies — but the view is consistent with the // ChatMiddlewareConfig shape). const { outputSchema: pinnedSchema, ...chatConfigSlice } = structuredConfig - const postOnConfig = await this.middlewareRunner.runOnConfig( - this.middlewareCtx, - { ...chatConfigSlice, tools: baseConfig.tools }, + const postOnConfig = yield* this.runWhileYielding( + this.middlewareRunner.runOnConfig(this.middlewareCtx, { + ...chatConfigSlice, + tools: baseConfig.tools, + }), ) // Apply merged config back to engine state @@ -3471,7 +3499,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 +3978,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 +4397,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 @@ -4400,6 +4430,7 @@ class TextEngine< for (const spec of normalizeStreamChunk(output as AdapterYieldChunk)) { restorePublicUsage(spec) if (spec.type === EventType.RUN_STARTED) { + if (this.hasPublicRunStarted) continue this.hasPublicRunStarted = true } yield spec @@ -4420,6 +4451,69 @@ class TextEngine< chunk, ) yield* this.emitPublicChunks(afterMw) + if (!this.drainingMiddlewareCustom) { + yield* this.drainMiddlewareCustomQueue() + } + } + + /** + * Drain CUSTOM chunks pushed by `ctx.emitCustomEvent` through middleware + * and into the public stream. If the run has not yet sent `RUN_STARTED`, + * emit that first so CUSTOM events are not the first wire event. + */ + private async *drainMiddlewareCustomQueue(): AsyncGenerator { + if (this.drainingMiddlewareCustom) return + if (this.middlewareCustomQueue.length === 0) return + this.drainingMiddlewareCustom = true + try { + yield* this.emitSyntheticRunStarted(this.createSyntheticFinishedEvent()) + while (this.middlewareCustomQueue.length > 0) { + const chunk = this.middlewareCustomQueue.shift() + if (chunk) yield* this.pipeThroughMiddleware(chunk) + } + } finally { + this.drainingMiddlewareCustom = false + } + } + + /** + * Await `work` while yielding any `emitCustomEvent` chunks as they arrive. + */ + private async *runWhileYielding( + work: Promise, + ): AsyncGenerator { + let settled = false + let result: T | undefined + let error: unknown + const done = work.then( + (value) => { + settled = true + result = value + }, + (err: unknown) => { + settled = true + error = err + }, + ) + + while (!settled) { + yield* this.drainMiddlewareCustomQueue() + if (settled) break + await Promise.race([ + done, + new Promise((resolve) => { + if (this.middlewareCustomQueue.length > 0) { + resolve() + return + } + this.middlewareCustomWaiters.push(resolve) + }), + ]) + } + + yield* this.drainMiddlewareCustomQueue() + if (error !== undefined) throw error + return result as T } /** @@ -4456,17 +4550,18 @@ class TextEngine< }, void > { - let next = await generator.next() - while (!next.done) { + let pending = generator.next() + while (true) { + const next = yield* this.runWhileYielding(pending) + if (next.done) return next.value yield* this.pipeThroughMiddleware(next.value) - next = await generator.next() + pending = generator.next() } - return next.value } private createCustomEventChunk( eventName: string, - value: Record, + value: Record, ): CustomEvent { return { type: EventType.CUSTOM, 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..933c4ea271 100644 --- a/packages/ai/src/activities/chat/middleware/types.ts +++ b/packages/ai/src/activities/chat/middleware/types.ts @@ -213,6 +213,12 @@ export interface ChatMiddlewareContext { signal?: AbortSignal /** Abort the chat run with a reason */ abort: (reason?: string) => void + /** + * Push a `CUSTOM` chunk onto the chat stream immediately. + * The engine yields it as soon as it can (including while `onConfig` + * is still awaiting work such as a summarize call). + */ + emitCustomEvent: (name: string, value: Record) => void /** Runtime context provided by chat() options */ context: TContext /** @@ -305,7 +311,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/middleware-capabilities.test.ts b/packages/ai/tests/middleware-capabilities.test.ts index 3c02a4bb44..4344a4bc33 100644 --- a/packages/ai/tests/middleware-capabilities.test.ts +++ b/packages/ai/tests/middleware-capabilities.test.ts @@ -85,6 +85,7 @@ function makeRunnerCtx(): ChatMiddlewareContext { iteration: 0, chunkIndex: 0, abort: () => {}, + emitCustomEvent: () => {}, context: undefined, defer: () => {}, provider: 'test', diff --git a/packages/ai/tests/middleware-interrupt.test.ts b/packages/ai/tests/middleware-interrupt.test.ts index acd7946e2a..fd3628592e 100644 --- a/packages/ai/tests/middleware-interrupt.test.ts +++ b/packages/ai/tests/middleware-interrupt.test.ts @@ -38,6 +38,7 @@ const context: ChatMiddlewareContext = { iteration: 0, chunkIndex: 0, abort() {}, + emitCustomEvent() {}, context: undefined, defer() {}, activity: 'chat', diff --git a/packages/ai/tests/middleware.test.ts b/packages/ai/tests/middleware.test.ts index f0b784616d..1929a591f9 100644 --- a/packages/ai/tests/middleware.test.ts +++ b/packages/ai/tests/middleware.test.ts @@ -165,6 +165,84 @@ describe('chat() middleware', () => { expect(onFinish).not.toHaveBeenCalled() }) + it('yields emitCustomEvent chunks during onConfig before adapter chunks', async () => { + let release!: () => void + const gate = new Promise((resolve) => { + release = resolve + }) + + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.textStart(), + ev.textContent('Hello'), + ev.textEnd(), + ev.runFinished('stop'), + ], + ], + }) + + const middleware: ChatMiddleware = { + name: 'live-emit', + async onConfig(ctx) { + if (ctx.phase !== 'beforeModel') return + ctx.emitCustomEvent('test:started', { ok: true }) + await gate + ctx.emitCustomEvent('test:ended', { ok: true }) + }, + } + + const stream = chat({ + adapter, + messages: [{ role: 'user', content: 'Hi' }], + middleware: [middleware], + }) as AsyncIterable + const iter = stream[Symbol.asyncIterator]() + + const beforeText: Array = [] + while (true) { + const next = await iter.next() + expect(next.done).toBe(false) + const chunk = next.value + expect(chunk).toBeDefined() + if (!chunk) break + expect(chunk.type).not.toBe(EventType.TEXT_MESSAGE_CONTENT) + beforeText.push(chunk) + if (chunk.type === EventType.CUSTOM && chunk.name === 'test:started') { + break + } + } + + expect( + beforeText.some((chunk) => chunk.type === EventType.RUN_STARTED), + ).toBe(true) + + release() + + const rest: Array = [] + while (true) { + const next = await iter.next() + if (next.done) break + rest.push(next.value) + } + + const endedIndex = rest.findIndex( + (chunk) => + chunk.type === EventType.CUSTOM && chunk.name === 'test:ended', + ) + const textIndex = rest.findIndex( + (chunk) => chunk.type === EventType.TEXT_MESSAGE_CONTENT, + ) + expect(endedIndex).toBeGreaterThanOrEqual(0) + expect(textIndex).toBeGreaterThan(endedIndex) + + const started = [...beforeText, ...rest].filter( + (chunk) => chunk.type === EventType.RUN_STARTED, + ) + expect(started).toHaveLength(1) + }) + it('should call exactly one terminal hook per run', async () => { const onStart = vi.fn() const onFinish = vi.fn() diff --git a/packages/ai/tests/middlewares/fake-otel.ts b/packages/ai/tests/middlewares/fake-otel.ts index b3ad6bf222..08b373e30b 100644 --- a/packages/ai/tests/middlewares/fake-otel.ts +++ b/packages/ai/tests/middlewares/fake-otel.ts @@ -261,6 +261,7 @@ export function makeCtx( iteration: 0, chunkIndex: 0, abort: () => {}, + emitCustomEvent: () => {}, context: undefined, defer: () => {}, provider: 'openai', 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..e489e973be 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -697,6 +697,9 @@ importers: '@tanstack/ai-codex': specifier: workspace:* version: link:../../packages/ai-codex + '@tanstack/ai-compaction': + specifier: workspace:* + version: link:../../packages/ai-compaction '@tanstack/ai-elevenlabs': specifier: workspace:* version: link:../../packages/ai-elevenlabs @@ -1795,6 +1798,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 +2956,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 +3107,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 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, +})