feat(agent-core): discard loaded tool schemas on compaction (#1471)

Align progressive tool disclosure with the discard-on-compaction model:
compaction no longer rebuilds loaded dynamic tool schemas. The boundary
announcement re-lists every loadable name, the model re-selects what it
still needs, and a from-memory call to a no-longer-loaded tool is
rejected by preflight with select guidance.

This removes the keep-all rebuild and its half-trigger budget heuristics
entirely: the post-compaction floor is back to users + summary, which is
structurally outside the auto-compaction trigger band, and the guard
baseline degenerates to summary + reinjected reminders. Every downstream
mechanism already treated the empty loaded set as its consistent base
state (ledger scan, pending clear at the compaction boundary, deferred
extras, preflight wording), so this is a strict simplification.

Co-authored-by: fengchenchen <fengchenchen@moonshot.ai>
This commit is contained in:
STAR-QUAKE 2026-07-07 23:06:29 +08:00 committed by GitHub
parent 131700097a
commit 9b76e5bff6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 77 additions and 130 deletions

View file

@ -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.

View file

@ -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<string>): 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<typeof tool> => 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) {

View file

@ -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<string> {
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();

View file

@ -117,7 +117,7 @@ export async function runTurn(input: RunTurnInput): Promise<TurnResult> {
// 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,

View file

@ -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();

View file

@ -94,8 +94,8 @@ export class SelectToolsTool implements BuiltinTool<SelectToolsInput> {
// 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))

View file

@ -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('<tools_added>');
// 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(`<tools_added>\n${GRAFANA_TOOL}\n</tools_added>`);
// 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(