feat(context-engine): add toolsOverride field to AssembleResult

Allow context-engine plugins to filter tool schemas sent to the model.
When assemble() returns toolsOverride, only those tool schemas are
included in model requests, reducing token usage while keeping all
tool executors available.
This commit is contained in:
Gemini CLI 2026-03-15 22:06:02 +11:00
parent 843e3c1efb
commit 6fc2a0af60
3 changed files with 41 additions and 0 deletions

View File

@ -127,6 +127,7 @@ import { dropThinkingBlocks } from "../thinking.js";
import { collectAllowedToolNames } from "../tool-name-allowlist.js";
import { installToolResultContextGuard } from "../tool-result-context-guard.js";
import { splitSdkTools } from "../tool-split.js";
import { wrapStreamFnWithToolsOverride } from "../stream-payload-utils.js";
import { describeUnknownError, mapThinkingLevel } from "../utils.js";
import { flushPendingToolResultsAfterIdle } from "../wait-for-idle-before-flush.js";
import { waitForCompactionRetryWithAggregateTimeout } from "./compaction-retry-aggregate-timeout.js";
@ -2123,6 +2124,15 @@ export async function runEmbeddedAttempt(
`context engine: prepended system prompt addition (${assembled.systemPromptAddition.length} chars)`,
);
}
if (assembled.toolsOverride && assembled.toolsOverride.length > 0) {
activeSession.agent.streamFn = wrapStreamFnWithToolsOverride(
activeSession.agent.streamFn,
assembled.toolsOverride,
);
log.debug(
`context engine: applied tools override (${assembled.toolsOverride.length} tools)`,
);
}
} catch (assembleErr) {
log.warn(
`context engine assemble failed, using pipeline messages: ${String(assembleErr)}`,

View File

@ -1,4 +1,28 @@
import type { StreamFn } from "@mariozechner/pi-agent-core";
import type { AnyAgentTool } from "../pi-tools.types.js";
/**
* Creates a streamFn wrapper that filters the tools context to only include
* tools whose names are in the override set. This reduces token usage by
* sending fewer tool schemas to the model while keeping all executors available.
*/
export function wrapStreamFnWithToolsOverride(
underlying: StreamFn,
toolsOverride: AnyAgentTool[],
): StreamFn {
const allowedNames = new Set(toolsOverride.map((tool) => tool.name));
return (model, context, options) => {
const ctx = context as { tools?: unknown[] };
if (!Array.isArray(ctx.tools)) {
return underlying(model, context, options);
}
const filteredTools = ctx.tools.filter((tool) => {
const name = (tool as { name?: string })?.name;
return typeof name === "string" && allowedNames.has(name);
});
return underlying(model, { ...context, tools: filteredTools }, options);
};
}
export function streamWithPayloadPatch(
underlying: StreamFn,

View File

@ -1,4 +1,5 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { AnyAgentTool } from "../agents/pi-tools.types.js";
// Result types
@ -9,6 +10,12 @@ export type AssembleResult = {
estimatedTokens: number;
/** Optional context-engine-provided instructions prepended to the runtime system prompt */
systemPromptAddition?: string;
/**
* Optional filtered list of tools to send to the model instead of the default full set.
* When present, only the schemas for these tools are included in the model request,
* reducing token usage. Tool executors for all registered tools remain available.
*/
toolsOverride?: AnyAgentTool[];
};
export type CompactResult = {