From 8dbe7bd634dd920a807c3109444e7b077f401c2e Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 1 Oct 2026 12:40:44 +0000 Subject: [PATCH 1/2] [EPIC-W][LOU-W2] Context compaction: prune old tool results behind a pluggable strategy Add src/context/compaction.ts (exported from the package root): - CompactionStrategy / CompactionResult: a pluggable compaction interface. - pruneToolResultsStrategy(): replaces tool results older than the protected tail (newest protectedTokens, default 40,000, plus the latest assistant turn's results) with a "[pruned: result, N chars]" marker. System, user and assistant messages are never changed, so the transcript stays valid; markers and short results are left alone, so it is idempotent. - compactMessages(): pure, manual compaction. - createCompactionHook(): an AgentHook that compacts before a model call estimated above thresholdPercent (default 0.9) of the context window (registry lookup by request.model, else 128,000), reported through onCompaction. GenerateHookContext.request.messages is the run's state.messages array, so the hook rewrites it in place and pruning persists in later steps, checkpoints and result.messages. Docs: docs/compaction.md, linked from docs/api-overview.md and the README; llms files regenerated. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_016bY29ibQgv3MDdjAbUT4YR --- CHANGELOG.md | 1 + README.md | 1 + docs/api-overview.md | 11 ++ docs/compaction.md | 110 +++++++++++++++ llms-full.txt | 125 +++++++++++++++++ llms.txt | 1 + src/context/compaction.test.ts | 250 +++++++++++++++++++++++++++++++++ src/context/compaction.ts | 194 +++++++++++++++++++++++++ src/context/index.ts | 5 + src/index.ts | 3 + 10 files changed, 701 insertions(+) create mode 100644 docs/compaction.md create mode 100644 src/context/compaction.test.ts create mode 100644 src/context/compaction.ts create mode 100644 src/context/index.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index 8dd32e13f..b1d22c254 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] - 2026-09-28 ### Added +- Context compaction (LOU-W2): `createCompactionHook()` returns an `AgentHook` that, before a model call estimated above `thresholdPercent` (default 0.9) of the model's context window, replaces tool results older than the newest `protectedTokens` (default 40,000) with a `[pruned: result, N chars]` marker, in the run's transcript itself. `compactMessages()` compacts a conversation by hand; `CompactionStrategy` / `pruneToolResultsStrategy()` make the strategy pluggable. See docs/compaction.md. - `AgentSpec.mcpServers` (LOU-D20): a validated map of MCP servers (stdio `command`/`args`/`env` or HTTP `url`/`headers`) in agent spec files. `loadSpec()` reports a bad entry with its name, `loushy doctor` reads the validated field instead of the raw file, and `specToAgent()` exposes the parsed servers as `agent.mcpServers` (connecting them is TODO(LOU-D20.2)). Existing specs are unaffected. - `createAgent()` agents can pause for approval (LOU-D21): a `needsApproval` tool no longer fails the run with "requires approval but no approvalStore". The run pauses (`finishReason: 'awaiting-approval'`) in a per-agent `InMemoryApprovalStore` (or the new `approvalStore` option), and `agent.approvals.list()` / `agent.approvals.resolve({ id, approved, note? })` continue it, in its session if it paused in one. The `approve` option decides calls in code without pausing. diff --git a/README.md b/README.md index 0fea82fe6..9f45857df 100644 --- a/README.md +++ b/README.md @@ -259,6 +259,7 @@ console.log(session.id, session.messages.length); - [Sessions](docs/sessions.md) - multi-turn conversations: `agent.session()`, `MemorySessionStore`, `FileSessionStore` - [Skills](docs/skills.md) - on-demand instructions: `defineSkill()`, `loadSkills()`, how they save context - [Sub-agents](docs/sub-agents.md) - the `subagents` option and its `task` tool, what sub-agents inherit, approvals inside sub-agents +- [Context compaction](docs/compaction.md) - keep long runs under the context window: prune old tool results with `createCompactionHook()` - [Agent directories](docs/agent-directories.md) - define an agent as a folder: layout, mapping to `createAgent()`, moving to code, security - [Agent Forge](docs/agent-forge.md) - the visual dashboard (`loushy studio`): quickstart, first-agent walkthrough, hook authoring - Full guides site: [linuxdevil.github.io/agent-sdk-docs](https://linuxdevil.github.io/agent-sdk-docs/) diff --git a/docs/api-overview.md b/docs/api-overview.md index 04ab1e11b..214a95c55 100644 --- a/docs/api-overview.md +++ b/docs/api-overview.md @@ -413,6 +413,17 @@ registerModel({ }); ``` +### Context compaction + +`createCompactionHook({ thresholdPercent?, contextWindow?, protectedTokens?, strategy?, onCompaction? })` +returns an `AgentHook` that, before each model call above 90% (by default) of +the model's context window, replaces tool results older than the newest +40,000 tokens with a `[pruned: result, N chars]` marker. It edits the +run's transcript in place, so pruning persists in checkpoints and +`result.messages`. `compactMessages(messages, options)` does the same once, +by hand, and `pruneToolResultsStrategy()` / `CompactionStrategy` are the +built-in and pluggable strategies. See [Context compaction](./compaction.md). + ## Flows, evals, observability and security - `FlowBuilder` / `FlowExecutor` - multi-step workflow graphs. diff --git a/docs/compaction.md b/docs/compaction.md new file mode 100644 index 000000000..8f239878a --- /dev/null +++ b/docs/compaction.md @@ -0,0 +1,110 @@ +# Context compaction + +A long agent run fills its context window mostly with tool results: file +contents, search hits, API responses the model read once and no longer needs +in full. Compaction keeps such a run under the model's limit by replacing old +tool results with a short marker before the request is sent. + +## Pruning old tool results in a run + +`createCompactionHook()` returns an `AgentHook` named `compaction`. Register it +and pass the registry to `AgentExecutor.execute()`: + +```ts +import { AgentExecutor, HookRegistry, createCompactionHook } from '@loushy/build-ai-agent'; + +const hooks = new HookRegistry(); +hooks.register( + createCompactionHook({ + thresholdPercent: 0.9, // compact above 90% of the context window (the default) + protectedTokens: 40_000, // never prune the newest 40K tokens (the default) + onCompaction: ({ tokensBefore, tokensAfter, prunedToolCallIds }) => + console.log(`compacted ${tokensBefore} -> ${tokensAfter} tokens (${prunedToolCallIds.length} results)`), + }) +); + +const result = await AgentExecutor.execute({ agent, input, provider, toolRegistry, hooks }); +``` + +Before every model call the hook estimates the request's size with +`estimateTokens` (see [Models, tokens and cost](./api-overview.md#models-tokens-and-cost)). +When it is above `thresholdPercent` of the context window, the hook runs its +strategy and calls `onCompaction` with the token counts, the pruned +`toolCallId`s and the strategy name. + +| Option | Default | Meaning | +| --- | --- | --- | +| `thresholdPercent` | `0.9` | Compact when the estimated request is above this share of the context window. Must be in (0, 1]. | +| `contextWindow` | registry, else `128_000` | Context window in tokens. By default it is looked up for `request.model` with `getModelInfo()`; register your own models with `registerModel()`. | +| `protectedTokens` | `40_000` | The newest messages that fit in this many tokens are never changed. | +| `strategy` | `pruneToolResultsStrategy()` | How to compact (see below). | +| `onCompaction` | none | Called after each compaction that changed the conversation. | + +## What pruning changes + +`pruneToolResultsStrategy()` replaces the `content` of each tool result older +than the protected tail with a marker such as +`[pruned: search result, 18234 chars]`. It never changes: + +- system messages, user messages (the first one included) or assistant turns, + so every tool call keeps its result message and the transcript stays valid + for every provider; +- the results of the latest assistant turn, which the model has not read yet, + even when they are bigger than `protectedTokens`; +- results that are already a marker, or shorter than one. + +The hook rewrites the run's transcript itself, not a copy: the request's +`messages` array is the run's message list, and the hook changes it in place. +A pruned result therefore stays pruned in later steps, in checkpoints (a +resumed session does not bring the full result back) and in +`result.messages`, and it is not pruned again on the next step. Pruning is +lossy: if the model needs a pruned result again, it has to call the tool +again. + +## Compacting by hand + +`compactMessages(messages, options)` runs a strategy once, whatever the size +of the conversation, and returns a new array (its input is not modified). Use +it to shrink a stored transcript before you continue it: + +```ts +import { compactMessages, type Message } from '@loushy/build-ai-agent'; + +declare const history: Message[]; + +const { messages, tokensBefore, tokensAfter, prunedToolCallIds } = compactMessages(history, { + protectedTokens: 8_000, + model: 'gpt-4o-mini', // for the context-window lookup and the token estimator +}); +console.log(`${tokensBefore} -> ${tokensAfter} tokens, pruned ${prunedToolCallIds.length} results`); +``` + +## Writing a strategy + +A strategy is a name and a `compact()` function. It receives the messages, a +token counter for the request's model, the context window and +`protectedTokens`, and returns the new messages with before/after token counts. +Return the input array unchanged when there is nothing to do: the hook then +leaves the run alone and does not call `onCompaction`. + +```ts +import { createCompactionHook, type CompactionStrategy } from '@loushy/build-ai-agent'; + +// Keep the system prompt, the first user message and the last 20 messages. +const keepRecent: CompactionStrategy = { + name: 'keep-recent', + compact({ messages, estimateTokens }) { + const tokensBefore = estimateTokens(messages); + if (messages.length <= 22) return { messages, tokensBefore, tokensAfter: tokensBefore, prunedToolCallIds: [] }; + // A real strategy must not split a tool call from its result. + const kept = [...messages.slice(0, 2), ...messages.slice(-20)]; + return { messages: kept, tokensBefore, tokensAfter: estimateTokens(kept), prunedToolCallIds: [] }; + }, +}; + +const hook = createCompactionHook({ strategy: keepRecent }); +``` + +A summarizing strategy (replace old turns with a model-written summary) and +typed `compaction.*` stream events are planned; `onCompaction` is where they +will plug in. diff --git a/llms-full.txt b/llms-full.txt index 2202587c9..e74e94d32 100644 --- a/llms-full.txt +++ b/llms-full.txt @@ -257,6 +257,7 @@ console.log(session.id, session.messages.length); - [Sessions](https://github.com/LinuxDevil/agent-sdk/blob/main/docs/sessions.md) - multi-turn conversations: `agent.session()`, `MemorySessionStore`, `FileSessionStore` - [Skills](https://github.com/LinuxDevil/agent-sdk/blob/main/docs/skills.md) - on-demand instructions: `defineSkill()`, `loadSkills()`, how they save context - [Sub-agents](https://github.com/LinuxDevil/agent-sdk/blob/main/docs/sub-agents.md) - the `subagents` option and its `task` tool, what sub-agents inherit, approvals inside sub-agents +- [Context compaction](https://github.com/LinuxDevil/agent-sdk/blob/main/docs/compaction.md) - keep long runs under the context window: prune old tool results with `createCompactionHook()` - [Agent directories](https://github.com/LinuxDevil/agent-sdk/blob/main/docs/agent-directories.md) - define an agent as a folder: layout, mapping to `createAgent()`, moving to code, security - [Agent Forge](https://github.com/LinuxDevil/agent-sdk/blob/main/docs/agent-forge.md) - the visual dashboard (`loushy studio`): quickstart, first-agent walkthrough, hook authoring - Full guides site: [linuxdevil.github.io/agent-sdk-docs](https://linuxdevil.github.io/agent-sdk-docs/) @@ -1617,6 +1618,17 @@ registerModel({ }); ``` +### Context compaction + +`createCompactionHook({ thresholdPercent?, contextWindow?, protectedTokens?, strategy?, onCompaction? })` +returns an `AgentHook` that, before each model call above 90% (by default) of +the model's context window, replaces tool results older than the newest +40,000 tokens with a `[pruned: result, N chars]` marker. It edits the +run's transcript in place, so pruning persists in checkpoints and +`result.messages`. `compactMessages(messages, options)` does the same once, +by hand, and `pruneToolResultsStrategy()` / `CompactionStrategy` are the +built-in and pluggable strategies. See [Context compaction](https://github.com/LinuxDevil/agent-sdk/blob/main/docs/compaction.md). + ## Flows, evals, observability and security - `FlowBuilder` / `FlowExecutor` - multi-step workflow graphs. @@ -2222,6 +2234,119 @@ npm run test:e2e:studio # from the repo root, or apps/agent-forge/ This builds Agent Forge first (`build:studio`), then runs Playwright against it. +# Context compaction + +Source: docs/compaction.md + +A long agent run fills its context window mostly with tool results: file +contents, search hits, API responses the model read once and no longer needs +in full. Compaction keeps such a run under the model's limit by replacing old +tool results with a short marker before the request is sent. + +## Pruning old tool results in a run + +`createCompactionHook()` returns an `AgentHook` named `compaction`. Register it +and pass the registry to `AgentExecutor.execute()`: + +```ts +import { AgentExecutor, HookRegistry, createCompactionHook } from '@loushy/build-ai-agent'; + +const hooks = new HookRegistry(); +hooks.register( + createCompactionHook({ + thresholdPercent: 0.9, // compact above 90% of the context window (the default) + protectedTokens: 40_000, // never prune the newest 40K tokens (the default) + onCompaction: ({ tokensBefore, tokensAfter, prunedToolCallIds }) => + console.log(`compacted ${tokensBefore} -> ${tokensAfter} tokens (${prunedToolCallIds.length} results)`), + }) +); + +const result = await AgentExecutor.execute({ agent, input, provider, toolRegistry, hooks }); +``` + +Before every model call the hook estimates the request's size with +`estimateTokens` (see [Models, tokens and cost](https://github.com/LinuxDevil/agent-sdk/blob/main/docs/api-overview.md#models-tokens-and-cost)). +When it is above `thresholdPercent` of the context window, the hook runs its +strategy and calls `onCompaction` with the token counts, the pruned +`toolCallId`s and the strategy name. + +| Option | Default | Meaning | +| --- | --- | --- | +| `thresholdPercent` | `0.9` | Compact when the estimated request is above this share of the context window. Must be in (0, 1]. | +| `contextWindow` | registry, else `128_000` | Context window in tokens. By default it is looked up for `request.model` with `getModelInfo()`; register your own models with `registerModel()`. | +| `protectedTokens` | `40_000` | The newest messages that fit in this many tokens are never changed. | +| `strategy` | `pruneToolResultsStrategy()` | How to compact (see below). | +| `onCompaction` | none | Called after each compaction that changed the conversation. | + +## What pruning changes + +`pruneToolResultsStrategy()` replaces the `content` of each tool result older +than the protected tail with a marker such as +`[pruned: search result, 18234 chars]`. It never changes: + +- system messages, user messages (the first one included) or assistant turns, + so every tool call keeps its result message and the transcript stays valid + for every provider; +- the results of the latest assistant turn, which the model has not read yet, + even when they are bigger than `protectedTokens`; +- results that are already a marker, or shorter than one. + +The hook rewrites the run's transcript itself, not a copy: the request's +`messages` array is the run's message list, and the hook changes it in place. +A pruned result therefore stays pruned in later steps, in checkpoints (a +resumed session does not bring the full result back) and in +`result.messages`, and it is not pruned again on the next step. Pruning is +lossy: if the model needs a pruned result again, it has to call the tool +again. + +## Compacting by hand + +`compactMessages(messages, options)` runs a strategy once, whatever the size +of the conversation, and returns a new array (its input is not modified). Use +it to shrink a stored transcript before you continue it: + +```ts +import { compactMessages, type Message } from '@loushy/build-ai-agent'; + +declare const history: Message[]; + +const { messages, tokensBefore, tokensAfter, prunedToolCallIds } = compactMessages(history, { + protectedTokens: 8_000, + model: 'gpt-4o-mini', // for the context-window lookup and the token estimator +}); +console.log(`${tokensBefore} -> ${tokensAfter} tokens, pruned ${prunedToolCallIds.length} results`); +``` + +## Writing a strategy + +A strategy is a name and a `compact()` function. It receives the messages, a +token counter for the request's model, the context window and +`protectedTokens`, and returns the new messages with before/after token counts. +Return the input array unchanged when there is nothing to do: the hook then +leaves the run alone and does not call `onCompaction`. + +```ts +import { createCompactionHook, type CompactionStrategy } from '@loushy/build-ai-agent'; + +// Keep the system prompt, the first user message and the last 20 messages. +const keepRecent: CompactionStrategy = { + name: 'keep-recent', + compact({ messages, estimateTokens }) { + const tokensBefore = estimateTokens(messages); + if (messages.length <= 22) return { messages, tokensBefore, tokensAfter: tokensBefore, prunedToolCallIds: [] }; + // A real strategy must not split a tool call from its result. + const kept = [...messages.slice(0, 2), ...messages.slice(-20)]; + return { messages: kept, tokensBefore, tokensAfter: estimateTokens(kept), prunedToolCallIds: [] }; + }, +}; + +const hook = createCompactionHook({ strategy: keepRecent }); +``` + +A summarizing strategy (replace old turns with a model-written summary) and +typed `compaction.*` stream events are planned; `onCompaction` is where they +will plug in. + # Configuration Source: docs/configuration.md diff --git a/llms.txt b/llms.txt index b633e3c40..797f2151d 100644 --- a/llms.txt +++ b/llms.txt @@ -9,6 +9,7 @@ - [Quick Start](https://github.com/LinuxDevil/agent-sdk/blob/main/docs/quick-start.md): Every TypeScript snippet on this page is a complete, standalone ES module (`.mts`). - [API Overview](https://github.com/LinuxDevil/agent-sdk/blob/main/docs/api-overview.md): The package root (`@loushy/build-ai-agent`) exports everything below. - [Agent directories](https://github.com/LinuxDevil/agent-sdk/blob/main/docs/agent-directories.md): An agent can be defined as a directory. +- [Context compaction](https://github.com/LinuxDevil/agent-sdk/blob/main/docs/compaction.md): A long agent run fills its context window mostly with tool results: file contents, search hits, API responses the model read once and no longer needs in full. - [Configuration](https://github.com/LinuxDevil/agent-sdk/blob/main/docs/configuration.md): A declarative agent is a YAML (`.yaml`/`.yml`) or JSON (`.json`) file, validated with zod by `loadSpec()` (`src/spec/schema.ts`). - [Durable execution](https://github.com/LinuxDevil/agent-sdk/blob/main/docs/durable-execution.md): Pass a `sessionId` and a `checkpointStore` to `AgentExecutor.execute()` and the run survives a crash, a restart, an abort or a human approval that takes days. - [Evals](https://github.com/LinuxDevil/agent-sdk/blob/main/docs/evals.md): Evals are regression tests for agent *behaviour*: which tools the agent called, in what order, with which arguments, how many steps it took, and what it said. diff --git a/src/context/compaction.test.ts b/src/context/compaction.test.ts new file mode 100644 index 000000000..7392f6195 --- /dev/null +++ b/src/context/compaction.test.ts @@ -0,0 +1,250 @@ +import { describe, it, expect, vi } from 'vitest'; +import { z } from 'zod'; +import { + compactMessages, + createCompactionHook, + pruneToolResultsStrategy, + type CompactionInfo, + type CompactionStrategy, +} from './compaction'; +import type { GenerateOptions, Message } from '../providers'; +import { estimateTokens, registerModel } from '../models'; +import { HookRegistry } from '../execution/hooks'; +import { AgentExecutor } from '../execution/AgentExecutor'; +import { defineTool } from '../tools/defineTool'; +import { ToolRegistry } from '../tools'; +import { mockModel } from '../testing'; +import { AgentType } from '../types'; + +const BIG = 'x'.repeat(4_000); // about 1,000 tokens + +/** system, user, then `rounds` of assistant tool call + its (large) result. */ +function transcript(rounds: number, firstUser = 'start'): Message[] { + const messages: Message[] = [ + { role: 'system', content: 'You are helpful.' }, + { role: 'user', content: firstUser }, + ]; + for (let i = 1; i <= rounds; i++) { + const id = `call_${i}`; + messages.push({ + role: 'assistant', + content: '', + toolCalls: [{ id, type: 'function', function: { name: 'search', arguments: '{}' } }], + }); + messages.push({ role: 'tool', toolCallId: id, toolName: 'search', content: `${i}:${BIG}` }); + } + return messages; +} + +const isMarker = (m: Message) => /^\[pruned: search result, \d+ chars\]$/.test(m.content); + +/** Every assistant tool call is followed by a tool message answering it. */ +function expectValidTranscript(messages: Message[]): void { + messages.forEach((message, index) => { + for (const call of message.toolCalls ?? []) { + const answer = messages.slice(index + 1).find((m) => m.role === 'tool' && m.toolCallId === call.id); + expect(answer, `result for ${call.id}`).toBeDefined(); + } + }); +} + +describe('pruneToolResultsStrategy / compactMessages', () => { + it('replaces old tool results with a marker and keeps the protected tail intact and in order', () => { + const original = transcript(10); + const result = compactMessages(original, { protectedTokens: 3_000, strategy: pruneToolResultsStrategy() }); + + expect(result.messages).toHaveLength(original.length); + expect(result.messages.map((m) => m.role)).toEqual(original.map((m) => m.role)); + expect(result.messages.map((m) => m.toolCallId)).toEqual(original.map((m) => m.toolCallId)); + + const firstKept = result.messages.findIndex((m, i) => i > 1 && m === original[i] && m.role === 'tool'); + // Everything from the first kept result on is the original tail, untouched and in order... + expect(result.messages.slice(firstKept)).toEqual(original.slice(firstKept)); + expect(estimateTokens(original.slice(firstKept + 1))).toBeLessThanOrEqual(3_000); + // ...and every older result is a marker. + const older = result.messages.slice(0, firstKept).filter((m) => m.role === 'tool'); + expect(older.length).toBeGreaterThan(0); + expect(older.every(isMarker)).toBe(true); + expect(result.prunedToolCallIds).toEqual(older.map((m) => m.toolCallId)); + expect(result.messages[3].content).toBe(`[pruned: search result, ${original[3].content.length} chars]`); + + expect(result.tokensBefore).toBe(estimateTokens(original)); + expect(result.tokensAfter).toBe(estimateTokens(result.messages)); + expect(result.tokensAfter).toBeLessThan(result.tokensBefore / 2); + expectValidTranscript(result.messages); + }); + + it('never touches system messages, user messages or assistant turns', () => { + const original = transcript(4, BIG); + original.splice(4, 0, { role: 'user', content: BIG }); + const result = compactMessages(original, { protectedTokens: 0 }); + + result.messages.forEach((m, i) => { + if (m.role !== 'tool') expect(m).toBe(original[i]); + }); + expect(result.prunedToolCallIds).toEqual(['call_1', 'call_2', 'call_3']); + }); + + it('never prunes results of the latest assistant turn, even with no protected tokens', () => { + const original = transcript(3); + const result = compactMessages(original, { protectedTokens: 0 }); + expect(result.messages.at(-1)).toBe(original.at(-1)); + expect(result.prunedToolCallIds).toEqual(['call_1', 'call_2']); + }); + + it('is pure and idempotent: markers and short results are left alone', () => { + const original = transcript(5); + original.push({ role: 'assistant', content: '', toolCalls: [{ id: 'tiny', type: 'function', function: { name: 'n', arguments: '{}' } }] }); + original.push({ role: 'tool', toolCallId: 'tiny', toolName: 'n', content: '"ok"' }); + original.push({ role: 'assistant', content: 'done' }); + const snapshot = structuredClone(original); + + const first = compactMessages(original, { protectedTokens: 0 }); + expect(original).toEqual(snapshot); + expect(first.prunedToolCallIds).toEqual(['call_1', 'call_2', 'call_3', 'call_4', 'call_5']); + + const second = compactMessages(first.messages, { protectedTokens: 0 }); + expect(second.messages).toBe(first.messages); + expect(second.prunedToolCallIds).toEqual([]); + expect(second.tokensAfter).toBe(second.tokensBefore); + }); + + it('uses the given strategy', () => { + const strategy: CompactionStrategy = { + name: 'drop-all', + compact: ({ messages, estimateTokens: count, contextWindow }) => ({ + messages: messages.slice(0, 2), + tokensBefore: count(messages), + tokensAfter: contextWindow, + prunedToolCallIds: [], + }), + }; + const result = compactMessages(transcript(2), { strategy, model: 'gpt-4o-mini' }); + expect(result.messages).toHaveLength(2); + expect(result.tokensAfter).toBe(128_000); // gpt-4o-mini's window from the registry + }); +}); + +function request(messages: Message[], model = 'compaction-test-model'): GenerateOptions { + return { model, messages }; +} + +describe('createCompactionHook', () => { + registerModel({ id: 'compaction-test-model', provider: 'test', contextWindow: 10_000 }); + + async function runHook(hook = createCompactionHook(), req = request(transcript(10))) { + const hooks = new HookRegistry(); + hooks.register(hook); + const before = req.messages; + await hooks.runPreGenerate({ messages: req.messages, request: req }); + expect(req.messages).toBe(before); // rewritten in place, never replaced + return req; + } + + it('is named compaction and rejects a threshold outside (0, 1]', () => { + expect(createCompactionHook().name).toBe('compaction'); + expect(() => createCompactionHook({ thresholdPercent: 0 })).toThrow(RangeError); + expect(() => createCompactionHook({ thresholdPercent: 1.5 })).toThrow(RangeError); + }); + + it('compacts only above thresholdPercent of the context window', async () => { + const messages = transcript(4); + const tokens = estimateTokens(messages); + const onCompaction = vi.fn(); + + // Exactly at the threshold: nothing happens. + const atThreshold = await runHook( + createCompactionHook({ contextWindow: tokens * 2, thresholdPercent: 0.5, protectedTokens: 0, onCompaction }), + request(structuredClone(messages)) + ); + expect(atThreshold.messages).toEqual(messages); + expect(onCompaction).not.toHaveBeenCalled(); + + // One token over: compacts. + const over = await runHook( + createCompactionHook({ contextWindow: (tokens - 1) * 2, thresholdPercent: 0.5, protectedTokens: 0, onCompaction }), + request(structuredClone(messages)) + ); + expect(over.messages.filter(isMarker)).toHaveLength(3); + expect(onCompaction).toHaveBeenCalledTimes(1); + }); + + it('uses the registry context window for request.model and the default 0.9 threshold', async () => { + // compaction-test-model has a 10,000-token window: 8 rounds (~8,100 tokens) stay, 10 (~10,100) compact. + const onCompaction = vi.fn<(info: CompactionInfo) => void>(); + const small = await runHook(createCompactionHook({ protectedTokens: 2_000, onCompaction }), request(transcript(8))); + expect(small.messages.some(isMarker)).toBe(false); + + const large = await runHook(createCompactionHook({ protectedTokens: 2_000, onCompaction }), request(transcript(10))); + expect(large.messages.filter(isMarker).length).toBeGreaterThan(5); + expectValidTranscript(large.messages); + + const info = onCompaction.mock.calls[0][0]; + expect(info.strategy).toBe('prune-tool-results'); + expect(info.tokensBefore).toBeGreaterThan(9_000); + expect(info.tokensAfter).toBe(estimateTokens(large.messages)); + expect(info.prunedToolCallIds).toEqual(large.messages.filter(isMarker).map((m) => m.toolCallId)); + + // An unknown model falls back to a 128,000-token window. + const unknown = await runHook(createCompactionHook({ protectedTokens: 0 }), request(transcript(10), 'no-such-model')); + expect(unknown.messages.some(isMarker)).toBe(false); + }); + + it('does not report a compaction that changed nothing', async () => { + const onCompaction = vi.fn(); + const hook = createCompactionHook({ contextWindow: 100, protectedTokens: 0, onCompaction }); + const req = await runHook(hook); + await runHook(hook, req); + expect(onCompaction).toHaveBeenCalledTimes(1); + }); +}); + +describe('compaction in a run', () => { + it('prunes old tool results before each model call and keeps them pruned in the transcript', async () => { + const fetchPage = defineTool({ + name: 'fetch_page', + description: 'Fetch a page', + input: z.object({ n: z.number() }), + execute: async ({ n }) => `${n}:${BIG}`, + }); + const toolRegistry = new ToolRegistry(); + toolRegistry.register(fetchPage); + const call = (n: number) => ({ toolCalls: [{ name: 'fetch_page', args: { n }, id: `call_${n}` }] }); + const model = mockModel([call(1), call(2), call(3), call(4), 'done']); + + const compactions: CompactionInfo[] = []; + const hooks = new HookRegistry(); + hooks.register( + createCompactionHook({ contextWindow: 3_000, protectedTokens: 1_500, onCompaction: (info) => compactions.push(info) }) + ); + + const result = await AgentExecutor.execute({ + agent: { + id: 'a', + name: 'Agent', + agentType: AgentType.SmartAssistant, + prompt: 'Read pages.', + tools: { fetch_page: { tool: 'fetch_page' } }, + }, + input: 'read four pages', + provider: model, + toolRegistry, + hooks, + maxSteps: 10, + }); + + expect(result.text).toBe('done'); + // Each request stays under the threshold once compacted; the newest result is always intact. + for (const req of model.calls.slice(1)) { + expect(estimateTokens(req.messages as Message[])).toBeLessThan(3_000 * 0.9); + expect(req.messages.at(-1)?.content).toContain(BIG); + } + // Pruning persists in the run's transcript, and no result is pruned twice. + const pruned = compactions.flatMap((c) => c.prunedToolCallIds); + expect(new Set(pruned).size).toBe(pruned.length); + expect(pruned).toEqual(['call_1', 'call_2']); // call_3 and call_4 fit under the threshold + const markers = result.messages.filter((m) => m.role === 'tool' && m.content.startsWith('[pruned: fetch_page result')); + expect(markers.map((m) => m.toolCallId)).toEqual(pruned); + expectValidTranscript(result.messages); + }); +}); diff --git a/src/context/compaction.ts b/src/context/compaction.ts new file mode 100644 index 000000000..f26f38c1e --- /dev/null +++ b/src/context/compaction.ts @@ -0,0 +1,194 @@ +/** + * Context compaction (LOU-W2): keep a long run under its model's context + * window by pruning old tool results. + * + * Ships as an `AgentHook` (`createCompactionHook()`), so it needs no changes + * to the execution loop: `preGenerate` runs before every model call, and + * `ctx.request.messages` IS the run's transcript array (AgentExecutor passes + * `state.messages` itself into the request, see `prepareGenerateRequest()` in + * src/execution/generateStep.ts). The hook rewrites that array in place, so a + * pruned result stays pruned in later steps, in checkpoints and in + * `ExecutionResult.messages`, and is not pruned again on the next step. + */ + +import type { Message } from '../providers/llm'; +import type { AgentHook, GenerateHookContext } from '../execution/hooks'; +import { estimateTokens } from '../models/estimateTokens'; +import { getModelInfo } from '../models/registry'; + +/** Counts the tokens of a message or a conversation. */ +export type CompactionTokenCounter = (input: Message | Message[]) => number; + +/** What a {@link CompactionStrategy} gets to work with. */ +export interface CompactionInput { + /** The conversation to compact. Do not mutate it; return a new array. */ + messages: Message[]; + /** Token counter for the model the request is for. */ + estimateTokens: CompactionTokenCounter; + /** The model's context window, in tokens. */ + contextWindow: number; + /** How many of the most recent tokens to leave untouched. */ + protectedTokens: number; +} + +/** What a {@link CompactionStrategy} returns. */ +export interface CompactionResult { + /** The compacted conversation (the input array itself when nothing changed). */ + messages: Message[]; + tokensBefore: number; + tokensAfter: number; + /** `toolCallId`s whose results were replaced by a marker. */ + prunedToolCallIds: string[]; +} + +/** A way of shrinking a conversation. */ +export interface CompactionStrategy { + name: string; + compact(input: CompactionInput): CompactionResult; +} + +const DEFAULT_PROTECTED_TOKENS = 40_000; +/** Context window assumed for a model the registry does not know. */ +const FALLBACK_CONTEXT_WINDOW = 128_000; +const DEFAULT_THRESHOLD_PERCENT = 0.9; + +const PRUNED_MARKER = /^\[pruned: .* result, \d+ chars\]$/; + +function prunedMarker(message: Message): string { + return `[pruned: ${message.toolName ?? 'tool'} result, ${message.content.length} chars]`; +} + +/** + * Index of the first message of the protected tail: the newest messages + * whose tokens fit in `protectedTokens`, plus everything after the last + * assistant turn (results the model has not seen yet are never pruned). + */ +function protectedTailStart(messages: Message[], count: CompactionTokenCounter, protectedTokens: number): number { + let start = messages.length; + let tokens = 0; + for (let i = messages.length - 1; i >= 0; i--) { + tokens += count(messages[i]); + if (tokens > protectedTokens) break; + start = i; + } + const lastAssistant = messages.map((m) => m.role).lastIndexOf('assistant'); + return Math.min(start, lastAssistant + 1); +} + +/** + * Replaces the content of tool results older than the protected tail with a + * short marker such as `[pruned: search result, 18234 chars]`. Only `tool` + * messages change (system, user and assistant messages, tool calls included, + * are kept as they are), so every tool call still has its result and the + * transcript stays valid. A result that is already a marker, or shorter than + * one, is left alone, so running it again changes nothing. + */ +export function pruneToolResultsStrategy(): CompactionStrategy { + return { + name: 'prune-tool-results', + compact({ messages, estimateTokens: count, protectedTokens }) { + const tokensBefore = count(messages); + const tailStart = protectedTailStart(messages, count, protectedTokens); + const prunedToolCallIds: string[] = []; + let changed = false; + const compacted = messages.map((message, index) => { + if (index >= tailStart || message.role !== 'tool' || PRUNED_MARKER.test(message.content)) return message; + const marker = prunedMarker(message); + if (marker.length >= message.content.length) return message; + changed = true; + if (message.toolCallId) prunedToolCallIds.push(message.toolCallId); + return { ...message, content: marker }; + }); + if (!changed) return { messages, tokensBefore, tokensAfter: tokensBefore, prunedToolCallIds }; + return { messages: compacted, tokensBefore, tokensAfter: count(compacted), prunedToolCallIds }; + }, + }; +} + +/** Options for {@link compactMessages}. */ +export interface CompactMessagesOptions { + /** Defaults to {@link pruneToolResultsStrategy}. */ + strategy?: CompactionStrategy; + /** Defaults to the model's window in the model registry, else 128,000. */ + contextWindow?: number; + /** How many recent tokens to keep intact. Defaults to 40,000. */ + protectedTokens?: number; + /** Model id, used for the context-window lookup and passed to the token estimator. */ + model?: string; +} + +function prepare(messages: Message[], options: CompactMessagesOptions) { + const { model, protectedTokens = DEFAULT_PROTECTED_TOKENS } = options; + const strategy = options.strategy ?? pruneToolResultsStrategy(); + const contextWindow = + options.contextWindow ?? (model ? getModelInfo(model)?.contextWindow : undefined) ?? FALLBACK_CONTEXT_WINDOW; + const count: CompactionTokenCounter = (input) => estimateTokens(input, { model }); + return { strategy, contextWindow, input: { messages, estimateTokens: count, contextWindow, protectedTokens } }; +} + +/** + * Compacts a conversation now, whatever its size. Pure: `messages` is not + * modified. Use it to shrink a stored transcript by hand; inside a run use + * {@link createCompactionHook}. + * + * @example + * const { messages: smaller, tokensBefore, tokensAfter } = compactMessages(history, { protectedTokens: 8_000 }); + */ +export function compactMessages(messages: Message[], options: CompactMessagesOptions = {}): CompactionResult { + const { strategy, input } = prepare(messages, options); + return strategy.compact(input); +} + +/** Reported to `onCompaction` each time the hook compacts a request. */ +export interface CompactionInfo { + tokensBefore: number; + tokensAfter: number; + prunedToolCallIds: string[]; + /** The strategy's `name`. */ + strategy: string; +} + +/** Options for {@link createCompactionHook}. */ +export interface CompactionHookOptions extends Omit { + /** Compact once the request is estimated above this share of the context window. Defaults to 0.9. */ + thresholdPercent?: number; + /** Called after each compaction that changed the conversation. */ + onCompaction?: (info: CompactionInfo) => void; +} + +/** + * An `AgentHook` named `compaction` that compacts the conversation before a + * model call whose estimated size is above `thresholdPercent` of the context + * window (looked up by `request.model` unless `contextWindow` is given). The + * compacted messages replace the run's transcript in place, so they are what + * later steps, checkpoints and the result see. + * + * @example + * const hooks = new HookRegistry(); + * hooks.register(createCompactionHook({ onCompaction: (info) => console.log(info) })); + * await AgentExecutor.execute({ agent, input, provider, hooks }); + */ +export function createCompactionHook(options: CompactionHookOptions = {}): AgentHook { + const { thresholdPercent = DEFAULT_THRESHOLD_PERCENT, onCompaction } = options; + if (!(thresholdPercent > 0 && thresholdPercent <= 1)) { + throw new RangeError(`createCompactionHook: thresholdPercent must be in (0, 1], got ${thresholdPercent}`); + } + return { + name: 'compaction', + preGenerate(ctx: GenerateHookContext) { + const messages = ctx.request.messages; + const { strategy, contextWindow, input } = prepare(messages, { ...options, model: ctx.request.model }); + if (input.estimateTokens(messages) <= thresholdPercent * contextWindow) return; + const result = strategy.compact(input); + if (result.messages === messages) return; + // In place: request.messages is the run's transcript (see the file comment). + messages.splice(0, messages.length, ...result.messages); + onCompaction?.({ + tokensBefore: result.tokensBefore, + tokensAfter: result.tokensAfter, + prunedToolCallIds: result.prunedToolCallIds, + strategy: strategy.name, + }); + }, + }; +} diff --git a/src/context/index.ts b/src/context/index.ts new file mode 100644 index 000000000..27cfed2bb --- /dev/null +++ b/src/context/index.ts @@ -0,0 +1,5 @@ +/** + * Context management (LOU-W2): compaction of long conversations. + */ + +export * from './compaction'; diff --git a/src/index.ts b/src/index.ts index d58086678..df1385875 100644 --- a/src/index.ts +++ b/src/index.ts @@ -47,6 +47,9 @@ export * from './templates'; // Token estimation and model registry (LOU-W1) export * from './models'; +// Context compaction (LOU-W2) +export * from './context'; + // Utils export * from './utils'; From 7de7b9036f0916ae676cdcf26513bc0946419dd8 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 1 Oct 2026 12:44:52 +0000 Subject: [PATCH 2/2] [EPIC-W][LOU-W2] Simplify the prune step to pass the fallow health gate Move the per-message pruning decision into prunedToolResult() so the strategy's map callback stays under fallow's CRAP threshold. No behavior change. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_016bY29ibQgv3MDdjAbUT4YR --- src/context/compaction.ts | 22 +++++++++------------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/src/context/compaction.ts b/src/context/compaction.ts index f26f38c1e..96e427700 100644 --- a/src/context/compaction.ts +++ b/src/context/compaction.ts @@ -54,8 +54,11 @@ const DEFAULT_THRESHOLD_PERCENT = 0.9; const PRUNED_MARKER = /^\[pruned: .* result, \d+ chars\]$/; -function prunedMarker(message: Message): string { - return `[pruned: ${message.toolName ?? 'tool'} result, ${message.content.length} chars]`; +/** The message with its result replaced by a marker, or `undefined` when it is not a tool result worth pruning. */ +function prunedToolResult(message: Message): Message | undefined { + if (message.role !== 'tool' || PRUNED_MARKER.test(message.content)) return undefined; + const marker = `[pruned: ${message.toolName ?? 'tool'} result, ${message.content.length} chars]`; + return marker.length < message.content.length ? { ...message, content: marker } : undefined; } /** @@ -89,17 +92,10 @@ export function pruneToolResultsStrategy(): CompactionStrategy { compact({ messages, estimateTokens: count, protectedTokens }) { const tokensBefore = count(messages); const tailStart = protectedTailStart(messages, count, protectedTokens); - const prunedToolCallIds: string[] = []; - let changed = false; - const compacted = messages.map((message, index) => { - if (index >= tailStart || message.role !== 'tool' || PRUNED_MARKER.test(message.content)) return message; - const marker = prunedMarker(message); - if (marker.length >= message.content.length) return message; - changed = true; - if (message.toolCallId) prunedToolCallIds.push(message.toolCallId); - return { ...message, content: marker }; - }); - if (!changed) return { messages, tokensBefore, tokensAfter: tokensBefore, prunedToolCallIds }; + const compacted = messages.map((message, index) => (index < tailStart && prunedToolResult(message)) || message); + const pruned = compacted.filter((message, index) => message !== messages[index]); + const prunedToolCallIds = pruned.flatMap((message) => (message.toolCallId ? [message.toolCallId] : [])); + if (pruned.length === 0) return { messages, tokensBefore, tokensAfter: tokensBefore, prunedToolCallIds }; return { messages: compacted, tokensBefore, tokensAfter: count(compacted), prunedToolCallIds }; }, };