diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index d785218f819..fb48f62391b 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -137,6 +137,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"; @@ -2441,6 +2442,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)}`, diff --git a/src/agents/pi-embedded-runner/stream-payload-utils.ts b/src/agents/pi-embedded-runner/stream-payload-utils.ts index 580bf5b1391..7121996a791 100644 --- a/src/agents/pi-embedded-runner/stream-payload-utils.ts +++ b/src/agents/pi-embedded-runner/stream-payload-utils.ts @@ -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, diff --git a/src/context-engine/types.ts b/src/context-engine/types.ts index 03401fdf3f2..5f4f8731128 100644 --- a/src/context-engine/types.ts +++ b/src/context-engine/types.ts @@ -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 = {