diff --git a/CHANGELOG.md b/CHANGELOG.md index 4127bdaf3..fe55a059b 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. - Tools own their contract (LOU-D22): `ToolDescriptor` gains optional `inputSchema` (zod) and `execute(args, ctx)`, which are now the canonical fields. `defineTool()` sets both and no longer calls `ai`'s `tool()`. Argument validation, the schema sent to the model, and tool execution read them first and fall back to `tool.parameters` / `tool.execute`. `ToolDescriptor.tool` (the `ai` v4 `Tool`) is now legacy: it is still built by `defineTool()` and still accepted on hand-written descriptors this release, but new code should set `inputSchema` and `execute`. diff --git a/README.md b/README.md index 0c404b2c5..df8d60110 100644 --- a/README.md +++ b/README.md @@ -258,6 +258,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 b21f75102..f43a10c1f 100644 --- a/docs/api-overview.md +++ b/docs/api-overview.md @@ -419,6 +419,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 6fd834db5..db221a852 100644 --- a/llms-full.txt +++ b/llms-full.txt @@ -256,6 +256,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/) @@ -1618,6 +1619,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. @@ -2223,6 +2235,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..96e427700 --- /dev/null +++ b/src/context/compaction.ts @@ -0,0 +1,190 @@ +/** + * 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\]$/; + +/** 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; +} + +/** + * 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 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 }; + }, + }; +} + +/** 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';