feat: opt-in client-side tool search (load_tools) to cut tool-schema tokens

holmes/common/env_vars.py

@@ -62,6 +62,10 @@ def load_bool(env_var, default: Optional[bool]) -> Optional[bool]:
 
 EXTRA_HEADERS = os.environ.get("EXTRA_HEADERS", "")
 THINKING = os.environ.get("THINKING", "")
+# Client-side tool search. When enabled, heavy tool schemas (MCP toolsets) are held
+# back from the default tool list and a load_tools meta-tool is exposed so the model
+# loads them on demand instead of putting the whole catalog in context every turn.
+TOOL_SEARCH_ENABLED = load_bool("HOLMES_TOOL_SEARCH_ENABLED", False)
 REASONING_EFFORT = os.environ.get("REASONING_EFFORT", "").strip().lower()
 
 

holmes/core/tool_calling_llm.py

@@ -23,9 +23,11 @@
     LOG_LLM_USAGE_RESPONSE,
     RESET_REPEATED_TOOL_CALL_CHECK_AFTER_COMPACTION,
     TEMPERATURE,
+    TOOL_SEARCH_ENABLED,
     load_bool,
 )
 from holmes.core.llm import LLM
+from holmes.core.tool_search import LOAD_TOOLS_NAME
 from holmes.core.llm_usage import RequestStats
 from holmes.core.models import (
     FrontendToolResult,
@@ -205,6 +207,9 @@ def __init__(
         self.tool_results_dir = tool_results_dir
 
         self._skill_in_use: bool = False
+        # Names of deferred (held-back) tools the model has loaded via load_tools in
+        # this conversation. Heavy MCP tool schemas stay out of context until loaded.
+        self._loaded_tool_names: set[str] = set()
 
     def with_executor(self, tool_executor: ToolExecutor) -> "ToolCallingLLM":
         """Return a shallow copy with a different ToolExecutor.
@@ -220,15 +225,18 @@ def with_executor(self, tool_executor: ToolExecutor) -> "ToolCallingLLM":
             tracer=self.tracer,
         )
         # Preserve transient state so resumed turns keep access to
-        # skill-unlocked (restricted) tools.
+        # skill-unlocked (restricted) tools and already-loaded (tool-search) tools.
         clone._skill_in_use = self._skill_in_use
+        clone._loaded_tool_names = set(self._loaded_tool_names)
         return clone
 
     def reset_interaction_state(self) -> None:
         """
-        For interactive loop, reset skills in use
+        For interactive loop, reset transient per-conversation state (skills in
+        use, tool-search loaded tools) so a new conversation starts fresh.
         """
         self._skill_in_use = False
+        self._loaded_tool_names.clear()
 
     def _supports_vision(self) -> bool:
         """Check if vision/multimodal input is enabled.
@@ -561,6 +569,16 @@ def _get_tools(self) -> list:
         replace _connect placeholders for authenticated users.
         """
         user_id = (self._request_context or {}).get("user_id") if hasattr(self, "_request_context") else None
+        if TOOL_SEARCH_ENABLED:
+            # Progressive disclosure: hold heavy (MCP) tool schemas out of context and
+            # expose a load_tools meta-tool. Tools the model has loaded this conversation
+            # are included. The agent loop re-fetches tools each step (see call_stream),
+            # so newly loaded tools appear on the next turn.
+            return self.tool_executor.get_visible_tools_openai_format(
+                self._loaded_tool_names,
+                include_restricted=self._should_include_restricted_tools(),
+                user_id=user_id,
+            )
         return self.tool_executor.get_all_tools_openai_format(
             include_restricted=self._should_include_restricted_tools(),
             user_id=user_id,
@@ -867,6 +885,40 @@ def _log_tool_call_result(
             metadata=metadata,
         )
 
+    def _handle_load_tools(
+        self, tool_call_id: str, tool_params: Dict[str, Any]
+    ) -> ToolCallResult:
+        """Execute the load_tools meta-tool (tool search): find held-back tools matching
+        the query and mark them loaded, so the next agent step can call them."""
+        query = (tool_params or {}).get("query", "") or ""
+        matched = self.tool_executor.search_deferred_tools(query)
+        self._loaded_tool_names.update(matched)
+        if matched:
+            lines = []
+            for name in matched:
+                tool = self.tool_executor.get_tool_by_name(name)
+                desc = ""
+                if tool and tool.description:
+                    desc = tool.description.strip().splitlines()[0]
+                lines.append(f"- {name}: {desc}" if desc else f"- {name}")
+            data = (
+                f"Loaded {len(matched)} tool(s) matching '{query}'. They are now "
+                "available — call them directly on your next step:\n" + "\n".join(lines)
+            )
+        else:
+            data = (
+                f"No tools matched '{query}'. Try broader keywords or an integration "
+                "name (e.g. aws, github, pagerduty, kafka, tempo, opensearch, grafana)."
+            )
+        return ToolCallResult(
+            tool_call_id=tool_call_id,
+            tool_name=LOAD_TOOLS_NAME,
+            description=f"load_tools(query={query!r})",
+            result=StructuredToolResult(
+                status=StructuredToolResultStatus.SUCCESS, data=data
+            ),
+        )
+
     def _invoke_llm_tool_call(
         self,
         tool_to_call: ChatCompletionMessageToolCall,
@@ -915,6 +967,16 @@ def _invoke_llm_tool_call(
                     f"Failed to parse arguments for tool: {tool_name}. args: {tool_arguments}"
                 )
 
+            # load_tools is a Holmes-owned meta-tool (tool search). Handle it here so it
+            # never reaches the ToolExecutor (it isn't a toolset tool) and so the loaded
+            # set is updated before the next _get_tools() re-fetch exposes the new tools.
+            if tool_name == LOAD_TOOLS_NAME:
+                tool_call_result = self._handle_load_tools(tool_id, tool_params)
+                ToolCallingLLM._log_tool_call_result(
+                    tool_span, tool_call_result, enable_tool_approval
+                )
+                return tool_call_result
+
             tool_response = None
             if not user_approved:
                 tool_response = prevent_overly_repeated_tool_call(

holmes/core/tool_search.py

@@ -0,0 +1,54 @@
+"""Client-side tool search (progressive tool disclosure).
+
+When ``HOLMES_TOOL_SEARCH_ENABLED`` is set, heavy tool schemas (MCP toolsets) are
+NOT loaded into the model's context up front. Instead the model is given a single
+``load_tools`` function tool; it calls that with a query to discover and load the
+tools it needs, and they become available on the next step.
+
+This is a deliberate alternative to Anthropic's server-side tool-search /
+``defer_loading`` beta: that path is mis-handled by LiteLLM (server_tool_use is
+converted to a client tool_use and result blocks are dropped — BerriAI/litellm
+issues #17737 and #28083), which breaks multi-step tool use through Bedrock-backed
+gateways. ``load_tools`` is a plain function tool that HolmesGPT executes itself, so
+it works identically across the Anthropic API, Bedrock, OpenAI and any
+LiteLLM-fronted gateway — no beta header, no server tools.
+
+On large MCP deployments this keeps 10-60k tokens of tool schemas out of every turn
+until they're actually needed.
+"""
+
+# Name of the meta-tool the model calls to discover/load held-back tools.
+LOAD_TOOLS_NAME = "load_tools"
+
+# OpenAI function-tool definition for load_tools. It is a normal function tool —
+# HolmesGPT intercepts and executes it (see ToolCallingLLM._handle_load_tools).
+LOAD_TOOLS_TOOL: dict = {
+    "type": "function",
+    "function": {
+        "name": LOAD_TOOLS_NAME,
+        "description": (
+            "Search for and load additional tools that are not loaded by default. "
+            "To keep the context small, heavy integrations are hidden until needed — "
+            "for example: AWS, GitHub, PagerDuty, Kafka/MSK, Tempo tracing, "
+            "OpenSearch/Elasticsearch metrics, and Grafana. Call this with keywords "
+            "describing the capability you need (e.g. 'aws s3', 'github pull request', "
+            "'kafka consumer lag', 'pagerduty incidents', 'tempo trace', "
+            "'opensearch metrics'). The matching tools become available to call on your "
+            "next step. If nothing matches, retry with broader keywords or an "
+            "integration name. Tools already loaded in this conversation stay available."
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": (
+                        "Keywords (or a regular expression) describing the capability or "
+                        "integration you need, e.g. 'aws ec2', 'github', 'kafka lag'."
+                    ),
+                }
+            },
+            "required": ["query"],
+        },
+    },
+}

holmes/core/tools_utils/tool_executor.py

@@ -1,18 +1,27 @@
 import logging
-from typing import Any, Dict, List, Optional
+import re
+from typing import List, Optional
 
 import sentry_sdk
 
 from holmes.core.init_event import EventCallback, StatusEvent, StatusEventKind
+from holmes.core.tool_search import LOAD_TOOLS_TOOL
 from holmes.core.tools import (
     Tool,
     Toolset,
     ToolsetStatusEnum,
+    ToolsetType,
 )
 from holmes.core.tools_utils.oauth_tool_connector import OAuthToolConnector
 
 display_logger = logging.getLogger("holmes.display.tool_executor")
 
+# Toolset types whose tools are "deferred" (held back from the default tool list and
+# loaded on demand via load_tools) when tool search is enabled. MCP servers dominate
+# tool-schema bloat; built-in toolsets (kubernetes, bash, etc.) stay loaded so common
+# operations don't pay a search hop.
+DEFERRABLE_TOOLSET_TYPES = frozenset({ToolsetType.MCP})
+
 
 class ToolExecutor:
     def __init__(self, toolsets: List[Toolset], on_event: EventCallback = None):
@@ -153,3 +162,75 @@ def _get_base_tools(self, include_restricted: bool = True) -> list:
                 continue
             tools.append(tool.get_openai_format())
         return tools
+
+    # ── Tool search (progressive disclosure) ───────────────────────────
+
+    def _is_deferrable(self, tool_name: str) -> bool:
+        """Whether a tool is held back from the default list until loaded on demand."""
+        toolset = self._tool_to_toolset.get(tool_name)
+        return toolset is not None and toolset.type in DEFERRABLE_TOOLSET_TYPES
+
+    def get_visible_tools_openai_format(
+        self,
+        loaded_tool_names: Optional[set] = None,
+        include_restricted: bool = True,
+        user_id: Optional[str] = None,
+    ) -> list:
+        """Tool-search mode: return core (non-deferrable) tools + the ``load_tools``
+        meta-tool + any deferrable tools that have already been loaded. Heavy MCP tool
+        schemas stay out of context until the model loads them via ``load_tools``.
+        """
+        loaded = loaded_tool_names or set()
+        tools = []
+        for tool in self.tools_by_name.values():
+            if not include_restricted and tool._is_restricted():
+                continue
+            if self._is_deferrable(tool.name) and tool.name not in loaded:
+                continue
+            tools.append(tool.get_openai_format())
+        tools.append(LOAD_TOOLS_TOOL)
+        return self.oauth_connector.apply_user_tools(tools, user_id, self._tool_to_toolset)
+
+    def search_deferred_tools(self, query: str, limit: int = 8) -> List[str]:
+        """Return names of held-back (deferrable) tools matching ``query``.
+
+        ``query`` is treated as a regular expression (Python ``re.search``); if it isn't
+        valid regex it falls back to a case-insensitive substring match. Each tool's
+        name, description and owning toolset's name/description are searched, so a query
+        like ``"kafka"`` matches every tool in the kafka toolset.
+        """
+        query = (query or "").strip()
+        if not query:
+            return []
+        try:
+            pattern = re.compile(query, re.IGNORECASE)
+
+            def matches(text: str) -> bool:
+                return bool(pattern.search(text))
+        except re.error:
+            needle = query.lower()
+
+            def matches(text: str) -> bool:
+                return needle in text.lower()
+
+        results: List[str] = []
+        for name, tool in self.tools_by_name.items():
+            if not self._is_deferrable(name):
+                continue
+            ts = self._tool_to_toolset.get(name)
+            haystack = " ".join(
+                filter(
+                    None,
+                    [
+                        name,
+                        tool.description or "",
+                        ts.name if ts else "",
+                        (ts.description or "") if ts else "",
+                    ],
+                )
+            )
+            if matches(haystack):
+                results.append(name)
+                if len(results) >= limit:
+                    break
+        return results

tests/core/test_tool_executor.py

@@ -106,3 +106,80 @@ def test_ensure_toolset_initialized_failure_blocks_subsequent_calls():
 
     # The callable should only have been invoked once (lazy init is not retried)
     mock_callable.assert_called_once()
+
+
+# ── Tool search (progressive disclosure) ─────────────────────────────────────
+
+from holmes.core.tool_search import LOAD_TOOLS_NAME  # noqa: E402
+from holmes.core.tools import (  # noqa: E402
+    StructuredToolResult,
+    StructuredToolResultStatus,
+    Tool,
+    ToolInvokeContext,
+    Toolset,
+    ToolsetType,
+)
+
+
+class _SearchTool(Tool):
+    def _invoke(self, params, context: ToolInvokeContext) -> StructuredToolResult:
+        return StructuredToolResult(status=StructuredToolResultStatus.SUCCESS)
+
+    def get_parameterized_one_liner(self, params) -> str:
+        return ""
+
+
+class _SearchToolset(Toolset):
+    def __init__(self, name, toolset_type, tool_specs, **kw):
+        super().__init__(name=name, description=f"{name} toolset", type=toolset_type, **kw)
+        self.tools = [_SearchTool(name=n, description=d) for n, d in tool_specs]
+
+
+def _executor_with_search() -> ToolExecutor:
+    core = _SearchToolset(
+        "kubernetes/core", ToolsetType.BUILTIN, [("kubectl_get", "get k8s resources")]
+    )
+    aws = _SearchToolset(
+        "aws_api", ToolsetType.MCP, [("call_aws", "run an aws cli command")]
+    )
+    kafka = _SearchToolset(
+        "kafka", ToolsetType.MCP, [("list_clusters", "list MSK clusters")]
+    )
+    for ts in (core, aws, kafka):
+        ts.status = ToolsetStatusEnum.ENABLED
+    return ToolExecutor(toolsets=[core, aws, kafka])
+
+
+def test_visible_tools_hide_mcp_and_include_load_tools():
+    ex = _executor_with_search()
+    names = {t["function"]["name"] for t in ex.get_visible_tools_openai_format(set())}
+    assert "kubectl_get" in names  # built-in core stays loaded
+    assert LOAD_TOOLS_NAME in names  # meta-tool exposed
+    assert "call_aws" not in names  # MCP tools held back
+    assert "list_clusters" not in names
+
+
+def test_visible_tools_include_loaded_mcp_tool():
+    ex = _executor_with_search()
+    names = {
+        t["function"]["name"]
+        for t in ex.get_visible_tools_openai_format({"call_aws"})
+    }
+    assert "call_aws" in names  # loaded on demand → now visible
+    assert "list_clusters" not in names  # other MCP tool still hidden
+
+
+def test_search_deferred_tools_matches_mcp_only():
+    ex = _executor_with_search()
+    assert ex.search_deferred_tools("aws") == ["call_aws"]
+    assert ex.search_deferred_tools("kafka") == ["list_clusters"]  # toolset-name match
+    assert ex.search_deferred_tools("kubectl") == []  # core tool isn't deferrable
+    assert ex.search_deferred_tools("no-such-thing") == []
+
+
+def test_search_deferred_tools_falls_back_for_invalid_regex():
+    # A malformed regex (e.g. the model emits "[" or "(") must not raise — it falls
+    # back to a case-insensitive substring match (no match here → empty list).
+    ex = _executor_with_search()
+    assert ex.search_deferred_tools("[") == []
+    assert ex.search_deferred_tools("(") == []

tests/core/test_tool_search.py

@@ -0,0 +1,10 @@
+from holmes.core.tool_search import LOAD_TOOLS_NAME, LOAD_TOOLS_TOOL
+
+
+def test_load_tools_tool_is_a_function_tool():
+    # Must be a normal OpenAI function tool so HolmesGPT can execute it and so it
+    # works across every provider (no Anthropic server-tool / beta dependency).
+    assert LOAD_TOOLS_TOOL["type"] == "function"
+    assert LOAD_TOOLS_TOOL["function"]["name"] == LOAD_TOOLS_NAME
+    assert "query" in LOAD_TOOLS_TOOL["function"]["parameters"]["properties"]
+    assert LOAD_TOOLS_TOOL["function"]["parameters"]["required"] == ["query"]