From 6fc2a0af60f26a573d0a707cffb9bf8d1d0685d1 Mon Sep 17 00:00:00 2001
From: Gemini CLI <gemini-cli@dsmedia.com>
Date: Sun, 15 Mar 2026 22:06:02 +1100
Subject: [PATCH] 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.
---
 src/agents/pi-embedded-runner/run/attempt.ts  | 10 ++++++++
 .../stream-payload-utils.ts                   | 24 +++++++++++++++++++
 src/context-engine/types.ts                   |  7 ++++++
 3 files changed, 41 insertions(+)

diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts
index ef5a63cdcd1..f34bd8c3628 100644
--- a/src/agents/pi-embedded-runner/run/attempt.ts
+++ b/src/agents/pi-embedded-runner/run/attempt.ts
@@ -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)}`,
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 7ddd695b5b6..efb8c473b94 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 = {