diff --git a/.changeset/select-tools-discard-on-compaction.md b/.changeset/select-tools-discard-on-compaction.md new file mode 100644 index 000000000..70bda32d5 --- /dev/null +++ b/.changeset/select-tools-discard-on-compaction.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Progressive tool disclosure (`select_tools`, experimental): compaction now discards the loaded tool schemas instead of re-injecting them. After a compaction the boundary announcement re-lists every loadable tool name and the model re-selects what it still needs; a from-memory call to a no-longer-loaded tool is rejected with guidance to select it first. This keeps the post-compaction context at its minimal users+summary floor and removes the schema-rebuild budget heuristics. No effect unless the `tool-select` experimental flag and a `select_tools`-capable model are active. diff --git a/packages/agent-core/src/agent/compaction/full.ts b/packages/agent-core/src/agent/compaction/full.ts index 0bbc12577..ca1f19c8b 100644 --- a/packages/agent-core/src/agent/compaction/full.ts +++ b/packages/agent-core/src/agent/compaction/full.ts @@ -11,7 +11,6 @@ import { type GenerateResult, type Message, type TokenUsage, - type Tool, APIContextOverflowError, APIStatusError, createUserMessage, @@ -20,11 +19,7 @@ import { import type { Agent } from '..'; import type { GenerateOptionsWithRequestLogFields } from '../llm-request-logger'; import type { ContextMessage } from '../context/types'; -import { - collectLoadedDynamicToolNames, - DYNAMIC_TOOL_SCHEMA_VARIANT, - stripDynamicToolContext, -} from '../context/dynamic-tools'; +import { stripDynamicToolContext } from '../context/dynamic-tools'; import { isAbortError } from '../../loop/errors'; import { retryBackoffDelays, @@ -374,65 +369,6 @@ export class FullCompaction { }).trimEnd(); } - /** - * Keep-all rebuild (Phase 1): after compaction folded the history — and the - * dynamic tool schema messages with it — append ONE merged schema message so - * the model keeps calling its loaded tools without re-selecting. Schemas are - * read from the live registry, never copied from the old history, so a - * schema that changed since load self-heals here. Names whose server is - * currently disconnected have no registry schema and are not rebuilt (the - * model re-selects after reconnect); names that survived into the - * post-compaction history (none under today's users+summary rebuild, but - * guarded) are not duplicated. The message goes through the normal - * injection-origin append, so estimation and records pick it up as usual. - * - * Budget guard: the rebuilt floor (users + summary + schemas) is the one - * part of the post-compaction context that compaction itself can never - * shrink — if it lands inside the auto-compaction trigger band, every - * following step re-compacts and rebuilds in a loop. Admit schemas (in name - * order) only while the projected context stays within HALF the compaction - * trigger, so normal turn content still fits before the next compaction; - * anything dropped is simply re-selectable on demand, the same degradation - * as a disconnected server. - */ - private rebuildDynamicToolSchemas(activeBefore: ReadonlySet): void { - if (!this.agent.toolSelectEnabled) return; - if (activeBefore.size === 0) return; - const surviving = collectLoadedDynamicToolNames(this.agent.context.history); - const names = [...activeBefore] - .filter((name) => !surviving.has(name)) - .toSorted((a, b) => a.localeCompare(b)); - const candidates = names - .map((name) => this.agent.tools.getMcpToolSchema(name)) - .filter((tool): tool is NonNullable => tool !== undefined); - if (candidates.length === 0) return; - const tools: Tool[] = []; - let projected = this.tokenCountWithPending; - for (const tool of candidates) { - const toolTokens = estimateTokensForTools([tool]); - // shouldCompact is monotonic in usedSize, so doubling the projected - // size checks "within half the trigger" for both trigger branches - // (ratio and reserved-context). - if (this.strategy.shouldCompact((projected + toolTokens) * 2)) break; - tools.push(tool); - projected += toolTokens; - } - if (tools.length < candidates.length) { - this.agent.log.info('trimmed dynamic tool schema rebuild to stay clear of the compaction trigger', { - kept: tools.length, - dropped: candidates.slice(tools.length).map((tool) => tool.name), - }); - } - if (tools.length === 0) return; - this.agent.context.appendMessage({ - role: 'system', - content: [], - toolCalls: [], - tools, - origin: { kind: 'injection', variant: DYNAMIC_TOOL_SCHEMA_VARIANT }, - }); - } - private postProcessSummary(summary: string): string { const storeData = this.agent.tools.storeData(); const todos = (storeData[TODO_STORE_KEY] as readonly TodoItem[] | undefined) ?? []; @@ -450,11 +386,6 @@ export class FullCompaction { const startedAt = Date.now(); const originalHistory = [...this.agent.context.history]; const tokensBefore = estimateTokensForMessages(originalHistory); - // Loaded-tools snapshot BEFORE the rebuild below folds the history away; - // read here so the keep-all schema rebuild after applyCompaction knows - // what was active. (The ledger scans history, which applyCompaction - // replaces.) - const activeDynamicToolsBefore = new Set(this.agent.tools.loadedDynamicToolNames()); let retryCount = 0; try { await this.triggerPreCompactHook(data, tokensBefore, signal); @@ -491,11 +422,12 @@ export class FullCompaction { // Dynamic-tool protocol context (schema messages, loadable-tools // announcements) is excluded from the summarizer input entirely: it is // protocol state, not conversation — summarizing it wastes tokens and - // risks schema text leaking into the summary. Zero information loss: - // the post-compaction boundary re-announces the manifest and the - // keep-all rebuild re-carries the schemas. Must happen before project() - // (which strips the origin anchor). `originalHistory` itself stays - // untouched for the prefix-race check and `compactedCount`. + // risks schema text leaking into the summary. The post-compaction + // boundary re-announces the manifest; the schemas themselves are + // deliberately dropped (discard-on-compaction) and re-selectable on + // demand. Must happen before project() (which strips the origin + // anchor). `originalHistory` itself stays untouched for the + // prefix-race check and `compactedCount`. let historyForModel: readonly ContextMessage[] = stripDynamicToolContext(originalHistory); let droppedCount = 0; let overflowShrinkCount = 0; @@ -618,7 +550,14 @@ export class FullCompaction { tokensBefore, droppedCount: droppedCount === 0 ? undefined : droppedCount, }); - this.rebuildDynamicToolSchemas(activeDynamicToolsBefore); + // Loaded dynamic tool schemas are deliberately NOT rebuilt: compaction + // discards the loaded set entirely (the boundary announcement re-lists + // every loadable name, and the model re-selects what it still needs). + // Everything downstream already treats the empty loaded set as its + // consistent base state — the ledger scan finds no schema messages, the + // pending set was cleared by applyCompaction, deferred extras drop out + // of the executable table, and a from-memory call is rejected by + // preflight with select guidance. // Telemetry keys are snake_case, but the `context.apply_compaction` // record written below keeps its persisted camelCase field names @@ -638,17 +577,11 @@ export class FullCompaction { ? {} : { input_tokens: inputTotal(usage), output_tokens: usage.output }), }); - // Baseline the "nothing new since compaction" guard on the counter - // that includes the schema rebuild appended above. `result.tokensAfter` - // predates the rebuild (and deliberately keeps its persisted - // users+summary semantics — the rebuild message is accounted through - // the pending-estimate tail, so folding it into tokensAfter would - // double-count on both the live and the restore path). A baseline - // below the actual post-compaction floor would let checkAutoCompaction - // re-trigger even though the compacted shape cannot shrink further. - // compactionWorker raises it once more after injectAfterCompaction so - // the reinjected reminders join the floor too; this earlier capture - // stays as the fallback when reinjection throws. + // Baseline the "nothing new since compaction" guard on the live counter + // (== result.tokensAfter here, since nothing has been appended since + // applyCompaction). compactionWorker raises it once more after + // injectAfterCompaction so the reinjected reminders join the floor; + // this earlier capture stays as the fallback when reinjection throws. this.lastCompactedTokenCount = this.tokenCountWithPending; return result; } catch (error) { diff --git a/packages/agent-core/src/agent/tool/index.ts b/packages/agent-core/src/agent/tool/index.ts index f46dcd546..8803227e0 100644 --- a/packages/agent-core/src/agent/tool/index.ts +++ b/packages/agent-core/src/agent/tool/index.ts @@ -555,8 +555,8 @@ export class ToolManager { * defer-window pending set. History is the single source of truth, so the * ledger survives resume (records replay rebuilds the history), keeps its * state across undo (schema messages have `injection` origin and are not - * undone), and self-heals after compaction (the rebuild message re-carries - * the schemas). + * undone), and empties at compaction (schema messages are discarded with + * the folded history — the model re-selects what it still needs). */ loadedDynamicToolNames(): ReadonlySet { const names = collectLoadedDynamicToolNames(this.agent.context.history); @@ -580,12 +580,11 @@ export class ToolManager { } /** - * Compaction rebuilt the history: from here on the keep-all rebuild message - * (which may have trimmed or skipped schemas — budget guard, disconnected - * servers) is the sole truth about what is still loaded. A pending entry - * surviving past this boundary would report a schema the context no longer - * carries as loaded, and re-selecting it would wrongly answer - * "Already available" instead of injecting. + * Compaction rebuilt the history and discarded every loaded schema with it + * — the loaded set is empty from here on. A pending entry surviving past + * this boundary would report a schema the context no longer carries as + * loaded, and re-selecting it would wrongly answer "Already available" + * instead of injecting. */ onContextCompacted(): void { this.pendingLoadedDynamicTools.clear(); diff --git a/packages/agent-core/src/loop/run-turn.ts b/packages/agent-core/src/loop/run-turn.ts index cb0f328e4..4d55c8579 100644 --- a/packages/agent-core/src/loop/run-turn.ts +++ b/packages/agent-core/src/loop/run-turn.ts @@ -117,7 +117,7 @@ export async function runTurn(input: RunTurnInput): Promise { // Passed through unresolved: the step evaluates it AFTER beforeStep, // next to buildMessages, so the tool table and the request messages // come from the same state (beforeStep can run compaction, which - // trims loaded schemas and rewrites the ledger). + // discards loaded schemas and empties the ledger). buildTools, describeMissingTool, hooks, diff --git a/packages/agent-core/src/loop/turn-step.ts b/packages/agent-core/src/loop/turn-step.ts index 74a3899b8..df33de9a0 100644 --- a/packages/agent-core/src/loop/turn-step.ts +++ b/packages/agent-core/src/loop/turn-step.ts @@ -42,7 +42,7 @@ export interface ExecuteLoopStepDeps { * Per-step tool table builder; wins over the static `tools` snapshot. * Evaluated after `beforeStep`, next to `buildMessages`, so the executable * table and the request messages reflect the same state — `beforeStep` can - * run compaction, which trims loaded dynamic tool schemas. + * run compaction, which discards loaded dynamic tool schemas. */ readonly buildTools?: (() => readonly ExecutableTool[]) | undefined; /** See RunTurnInput.describeMissingTool. */ @@ -90,8 +90,8 @@ export async function executeLoopStep(deps: ExecuteLoopStepDeps): Promise<{ signal.throwIfAborted(); // Resolve the tool table AFTER beforeStep so it reflects the same state as - // the messages built below (beforeStep can run compaction, which trims - // loaded dynamic tool schemas out of the context and the ledger — a table + // the messages built below (beforeStep can run compaction, which discards + // loaded dynamic tool schemas from the context and the ledger — a table // captured earlier would still dispatch a tool the model no longer has). const stepTools = buildTools !== undefined ? buildTools() : tools; const messages = await buildMessages(); diff --git a/packages/agent-core/src/tools/builtin/select-tools.ts b/packages/agent-core/src/tools/builtin/select-tools.ts index f4a610b40..b54c4a083 100644 --- a/packages/agent-core/src/tools/builtin/select-tools.ts +++ b/packages/agent-core/src/tools/builtin/select-tools.ts @@ -94,8 +94,8 @@ export class SelectToolsTool implements BuiltinTool { // sorted by name for byte-stable output. History is never used as a // schema source; an already-loaded name whose registry schema has // since changed is NOT re-injected (no runtime last-wins reliance) — - // the next compaction rebuild or an explicit re-select picks up the - // new schema. + // the stale copy lasts at most until the next compaction discards + // the loaded set, after which a re-select injects the new schema. toLoad.sort((a, b) => a.localeCompare(b)); const tools = toLoad .map((name) => manager.getMcpToolSchema(name)) diff --git a/packages/agent-core/test/agent/tool-select.e2e.test.ts b/packages/agent-core/test/agent/tool-select.e2e.test.ts index 9ef5ec207..e26a0082f 100644 --- a/packages/agent-core/test/agent/tool-select.e2e.test.ts +++ b/packages/agent-core/test/agent/tool-select.e2e.test.ts @@ -490,12 +490,13 @@ describe('disclosure mode — executable table freshness', () => { }); describe('disclosure mode — compaction', () => { - it('filters protocol context from the summarizer input and rebuilds schemas after compaction', async () => { + it('filters protocol context from the summarizer input and discards loaded schemas after compaction', async () => { const ctx = await disclosureAgent(); ctx.mockNextResponse({ type: 'text', text: 'load' }, selectCall('call-1', [GRAFANA_TOOL])); ctx.mockNextResponse({ type: 'text', text: 'ok' }); await runTurn(ctx, 'load it'); + expect(schemaMessages(ctx)).toHaveLength(1); const compacted = new Promise<{ tokensAfter: number }>((resolve) => { ctx.emitter.once('context.apply_compaction', (entry: { args: { tokensAfter: number } }) => { @@ -513,36 +514,42 @@ describe('disclosure mode — compaction', () => { expect(summarizerCall.history.some((m) => m.tools !== undefined)).toBe(false); expect(JSON.stringify(summarizerCall.history)).not.toContain(''); - // Post-compaction context: one rebuild message with the registry schema, - // plus a fresh full announcement — no re-select needed. - const rebuilt = schemaMessages(ctx); - expect(rebuilt).toHaveLength(1); - expect(rebuilt[0]!.tools!.map((t) => t.name)).toEqual([GRAFANA_TOOL]); - expect(rebuilt[0]!.origin).toEqual({ kind: 'injection', variant: 'dynamic_tool_schema' }); + // Post-compaction context: the loaded set is DISCARDED — no schema + // message is rebuilt, the ledger is empty, deferred extras drop out of + // the executable table. The boundary announcement re-lists every + // loadable name so the model re-selects what it still needs. + expect(schemaMessages(ctx)).toHaveLength(0); + expect(ctx.agent.tools.loadedDynamicToolNames().has(GRAFANA_TOOL)).toBe(false); + expect(ctx.agent.tools.loopTools.map((t) => t.name)).not.toContain(GRAFANA_TOOL); expect(ctx.agent.context.history.filter(isLoadableToolsAnnouncement)).toHaveLength(1); - expect(ctx.agent.tools.loadedDynamicToolNames().has(GRAFANA_TOOL)).toBe(true); - expect(ctx.agent.tools.loopTools.map((t) => t.name)).toContain(GRAFANA_TOOL); + expect(historyText(ctx)).toContain(`\n${GRAFANA_TOOL}\n`); - // The "nothing new since compaction" guard must be baselined on the - // true post-compaction floor: summary + rebuild message + the reinjected - // announcement. result.tokensAfter predates all of it, and a baseline - // that misses any re-appended piece would let auto-compaction re-trigger - // against a floor that cannot shrink (each round strips and re-appends - // the same reminders). + // The "nothing new since compaction" guard is baselined on the true + // post-compaction floor: summary + the reinjected announcement (no + // rebuild message anymore). A baseline missing any re-appended piece + // would let auto-compaction re-trigger against a floor that cannot + // shrink (each round strips and re-appends the same reminders). const internals = ctx.agent.fullCompaction as unknown as { lastCompactedTokenCount: number | null; }; const reAnnouncement = ctx.agent.context.history.filter(isLoadableToolsAnnouncement).at(-1)!; expect(internals.lastCompactedTokenCount).toBe( - tokensAfter + estimateTokensForMessage(rebuilt[0]!) + estimateTokensForMessage(reAnnouncement), + tokensAfter + estimateTokensForMessage(reAnnouncement), ); + // Re-selecting after compaction takes the injection branch again (not + // "Already available") — the discard is total, not just cosmetic. + ctx.mockNextResponse({ type: 'text', text: 'reload' }, selectCall('call-2', [GRAFANA_TOOL])); + ctx.mockNextResponse({ type: 'text', text: 'ok' }); + await runTurn(ctx, 'load it again'); + expect(toolResultTexts(ctx)).toContainEqual(`Loaded: ${GRAFANA_TOOL}`); + expect(schemaMessages(ctx)).toHaveLength(1); + expect(ctx.agent.tools.loadedDynamicToolNames().has(GRAFANA_TOOL)).toBe(true); + // The baseline lives strictly within one turn: runOneTurn re-arms it at // every turn boundary, which is what makes cross-turn staleness (undo, // model switches, /clear while idle) structurally impossible. If this // reset ever moves, the guard's staleness analysis must be redone. - ctx.mockNextResponse({ type: 'text', text: 'next turn' }); - await runTurn(ctx, 'anything new'); expect(internals.lastCompactedTokenCount).toBeNull(); }); @@ -590,11 +597,13 @@ describe('disclosure mode — compaction', () => { expect(backCall.tools.map((t) => t.name)).not.toContain('select_tools'); }); - it('trims the schema rebuild instead of re-entering the compaction trigger band', async () => { - // A trigger far below one fat schema: without the rebuild budget guard the - // post-compaction floor (users + summary + schema) would sit permanently - // above the trigger, and every later step would re-compact and rebuild in - // a loop (with the default Infinity per-turn cap, forever). + it('discards heavy loaded schemas at compaction instead of re-entering the trigger band', async () => { + // A trigger far below one fat schema: if compaction carried loaded + // schemas back into the context, the post-compaction floor (users + + // summary + schema) would sit permanently above the trigger and every + // later step would re-compact in a loop (with the default Infinity + // per-turn cap, forever). Discard-on-compaction keeps the floor at + // users + summary, structurally outside the band. const trigger = 2_000; const ctx = testAgent({ experimentalFlags: toolSelectFlagOn(), @@ -638,10 +647,11 @@ describe('disclosure mode — compaction', () => { await ctx.rpc.setPermission({ mode: 'yolo' }); // Step 1 loads the fat schema; step 2's boundary trips the trigger and - // blocks on auto-compaction (consuming the summary mock), which trims the - // rebuild. Step 2 then calls the MCP tool directly — the executable table - // is resolved AFTER the compaction (same state as the messages), so the - // now-unloaded tool must be rejected by preflight, not dispatched. + // blocks on auto-compaction (consuming the summary mock), which discards + // the loaded schema. Step 2 then calls the MCP tool directly — the + // executable table is resolved AFTER the compaction (same state as the + // messages), so the now-unloaded tool must be rejected by preflight, not + // dispatched. const fatCallLog: unknown[] = []; (fatClient as { callTool: unknown }).callTool = async (...args: unknown[]) => { fatCallLog.push(args); @@ -653,12 +663,12 @@ describe('disclosure mode — compaction', () => { ctx.mockNextResponse({ type: 'text', text: 'done' }); await runTurn(ctx, 'load the fat tool'); - // The rebuild was trimmed away: no schema message survives, the ledger is + // The loaded set was discarded: no schema message survives, the ledger is // empty again, and the tool is simply re-selectable on demand. expect(schemaMessages(ctx)).toHaveLength(0); expect(ctx.agent.tools.loadedDynamicToolNames().has(GRAFANA_TOOL)).toBe(false); - // The direct call after the trim was rejected with select guidance and + // The direct call after the discard was rejected with select guidance and // never reached the MCP client. expect(fatCallLog).toHaveLength(0); expect(toolResultTexts(ctx).join('\n')).toContain(