Skip to content

feat: Support built-in toolset white list by --allowed-builtin-toolsets #835

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
nilo19 wants to merge 1 commit into from
Closed

feat: Support built-in toolset white list by --allowed-builtin-toolsets #835

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 34 additions & 1 deletion holmes/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from typing import TYPE_CHECKING, Any, List, Optional, Union

import yaml # type: ignore
from pydantic import BaseModel, ConfigDict, FilePath, SecretStr
from pydantic import BaseModel, ConfigDict, FilePath, SecretStr, field_validator

from holmes.common.env_vars import ROBUSTA_AI, ROBUSTA_API_ENDPOINT, ROBUSTA_CONFIG_PATH
from holmes.core.tools_utils.tool_executor import ToolExecutor
Expand Down Expand Up @@ -112,6 +112,28 @@ class Config(RobustaBaseConfig):
# custom_toolsets_from_cli is passed from CLI option `--custom-toolsets` as 'experimental' custom toolsets.
# The status of toolset here won't be cached, so the toolset from cli will always be loaded when specified in the CLI.
custom_toolsets_from_cli: Optional[List[FilePath]] = None

# Optional filter for builtin toolsets - specifies which builtin toolsets should be loaded
allowed_builtin_toolsets: Optional[List[str]] = None

@field_validator("allowed_builtin_toolsets")
@classmethod
def _validate_allowed_builtin_toolsets(
cls, v: Optional[List[str]]
) -> Optional[List[str]]:
"""Validate allowed_builtin_toolsets field."""
if v is None:
return v

# Filter out empty strings and whitespace-only strings
if isinstance(v, list):
filtered = [name.strip() for name in v if name and name.strip()]
# Important: preserve empty list as empty list (different from None)
# None = no filtering, [] = filter to nothing
return filtered

return v

should_try_robusta_ai: bool = False # if True, we will try to load the Robusta AI model, in cli we aren't trying to load it.

toolsets: Optional[dict[str, dict[str, Any]]] = None
Expand All @@ -125,6 +147,7 @@ class Config(RobustaBaseConfig):
def toolset_manager(self) -> ToolsetManager:
if not self._toolset_manager:
self._toolset_manager = ToolsetManager(
config=self, # Pass self as config parameter
toolsets=self.toolsets,
mcp_servers=self.mcp_servers,
custom_toolsets=self.custom_toolsets,
Expand Down Expand Up @@ -182,6 +205,16 @@ def load_from_file(cls, config_file: Optional[Path], **kwargs) -> "Config":

cli_options = {k: v for k, v in kwargs.items() if v is not None and v != []}

# Parse CLI option for allowed_builtin_toolsets if provided
if "allowed_builtin_toolsets" in cli_options and isinstance(
cli_options["allowed_builtin_toolsets"], str
):
cli_options["allowed_builtin_toolsets"] = [
name.strip()
for name in cli_options["allowed_builtin_toolsets"].split(",")
if name.strip()
]

if config_from_file is None:
result = cls(**cli_options)
else:
Expand Down
12 changes: 11 additions & 1 deletion holmes/core/toolset_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,12 +25,14 @@ class ToolsetManager:

def __init__(
self,
config: Optional[Any] = None, # Config instance
toolsets: Optional[dict[str, dict[str, Any]]] = None,
mcp_servers: Optional[dict[str, dict[str, Any]]] = None,
custom_toolsets: Optional[List[FilePath]] = None,
custom_toolsets_from_cli: Optional[List[FilePath]] = None,
toolset_status_location: Optional[FilePath] = None,
):
self._config = config # Store config instance
self.toolsets = toolsets
self.toolsets = toolsets or {}
if mcp_servers is not None:
Expand Down Expand Up @@ -65,6 +67,12 @@ def server_tool_tags(self) -> List[ToolsetTag]:
"""
return [ToolsetTag.CORE, ToolsetTag.CLUSTER]

def _get_allowed_builtin_toolsets(self) -> Optional[List[str]]:
"""Get allowed builtin toolsets from config."""
if self._config is None:
return None
return getattr(self._config, "allowed_builtin_toolsets", None)

def _list_all_toolsets(
self,
dal: Optional[SupabaseDal] = None,
Expand All @@ -81,7 +89,9 @@ def _list_all_toolsets(
3. custom toolset from config can override both built-in and add new custom toolsets # for backward compatibility
"""
# Load built-in toolsets
builtin_toolsets = load_builtin_toolsets(dal)
builtin_toolsets = load_builtin_toolsets(
dal, allowed_builtin_toolsets=self._get_allowed_builtin_toolsets()
)
toolsets_by_name: dict[str, Toolset] = {
toolset.name: toolset for toolset in builtin_toolsets
}
Expand Down
17 changes: 17 additions & 0 deletions holmes/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,11 @@
"-r",
help="Path to a custom runbooks (can specify -r multiple times to add multiple runbooks)",
)
opt_allowed_builtin_toolsets: Optional[str] = typer.Option(
None,
"--allowed-builtin-toolsets",
help="Comma-separated list of builtin toolsets to allow (e.g., 'kubernetes/core,prometheus/core')",
)
opt_max_steps: Optional[int] = typer.Option(
10,
"--max-steps",
Expand Down Expand Up @@ -174,6 +179,7 @@ def ask(
model: Optional[str] = opt_model,
config_file: Optional[Path] = opt_config_file,
custom_toolsets: Optional[List[Path]] = opt_custom_toolsets,
allowed_builtin_toolsets: Optional[str] = opt_allowed_builtin_toolsets,
max_steps: Optional[int] = opt_max_steps,
verbose: Optional[List[bool]] = opt_verbose,
# semi-common options
Expand Down Expand Up @@ -240,6 +246,7 @@ def ask(
model=model,
max_steps=max_steps,
custom_toolsets_from_cli=custom_toolsets,
allowed_builtin_toolsets=allowed_builtin_toolsets,
slack_token=slack_token,
slack_channel=slack_channel,
)
Expand Down Expand Up @@ -371,6 +378,7 @@ def alertmanager(
config_file: Optional[Path] = opt_config_file, # type: ignore
custom_toolsets: Optional[List[Path]] = opt_custom_toolsets,
custom_runbooks: Optional[List[Path]] = opt_custom_runbooks,
allowed_builtin_toolsets: Optional[str] = opt_allowed_builtin_toolsets,
max_steps: Optional[int] = opt_max_steps,
verbose: Optional[List[bool]] = opt_verbose,
# advanced options for this command
Expand Down Expand Up @@ -402,6 +410,7 @@ def alertmanager(
slack_channel=slack_channel,
custom_toolsets_from_cli=custom_toolsets,
custom_runbooks=custom_runbooks,
allowed_builtin_toolsets=allowed_builtin_toolsets,
)

ai = config.create_console_issue_investigator() # type: ignore
Expand Down Expand Up @@ -508,6 +517,7 @@ def jira(
config_file: Optional[Path] = opt_config_file, # type: ignore
custom_toolsets: Optional[List[Path]] = opt_custom_toolsets,
custom_runbooks: Optional[List[Path]] = opt_custom_runbooks,
allowed_builtin_toolsets: Optional[str] = opt_allowed_builtin_toolsets,
max_steps: Optional[int] = opt_max_steps,
verbose: Optional[List[bool]] = opt_verbose,
json_output_file: Optional[str] = opt_json_output_file,
Expand All @@ -532,6 +542,7 @@ def jira(
jira_query=jira_query,
custom_toolsets_from_cli=custom_toolsets,
custom_runbooks=custom_runbooks,
allowed_builtin_toolsets=allowed_builtin_toolsets,
)
ai = config.create_console_issue_investigator() # type: ignore
source = config.create_jira_source()
Expand Down Expand Up @@ -699,6 +710,7 @@ def github(
config_file: Optional[Path] = opt_config_file, # type: ignore
custom_toolsets: Optional[List[Path]] = opt_custom_toolsets,
custom_runbooks: Optional[List[Path]] = opt_custom_runbooks,
allowed_builtin_toolsets: Optional[str] = opt_allowed_builtin_toolsets,
max_steps: Optional[int] = opt_max_steps,
verbose: Optional[List[bool]] = opt_verbose,
# advanced options for this command
Expand All @@ -723,6 +735,7 @@ def github(
github_query=github_query,
custom_toolsets_from_cli=custom_toolsets,
custom_runbooks=custom_runbooks,
allowed_builtin_toolsets=allowed_builtin_toolsets,
)
ai = config.create_console_issue_investigator()
source = config.create_github_source()
Expand Down Expand Up @@ -784,6 +797,7 @@ def pagerduty(
config_file: Optional[Path] = opt_config_file, # type: ignore
custom_toolsets: Optional[List[Path]] = opt_custom_toolsets,
custom_runbooks: Optional[List[Path]] = opt_custom_runbooks,
allowed_builtin_toolsets: Optional[str] = opt_allowed_builtin_toolsets,
max_steps: Optional[int] = opt_max_steps,
verbose: Optional[List[bool]] = opt_verbose,
json_output_file: Optional[str] = opt_json_output_file,
Expand All @@ -807,6 +821,7 @@ def pagerduty(
pagerduty_incident_key=pagerduty_incident_key,
custom_toolsets_from_cli=custom_toolsets,
custom_runbooks=custom_runbooks,
allowed_builtin_toolsets=allowed_builtin_toolsets,
)
ai = config.create_console_issue_investigator()
source = config.create_pagerduty_source()
Expand Down Expand Up @@ -870,6 +885,7 @@ def opsgenie(
config_file: Optional[Path] = opt_config_file, # type: ignore
custom_toolsets: Optional[List[Path]] = opt_custom_toolsets,
custom_runbooks: Optional[List[Path]] = opt_custom_runbooks,
allowed_builtin_toolsets: Optional[str] = opt_allowed_builtin_toolsets,
max_steps: Optional[int] = opt_max_steps,
verbose: Optional[List[bool]] = opt_verbose,
# advanced options for this command
Expand All @@ -893,6 +909,7 @@ def opsgenie(
opsgenie_query=opsgenie_query,
custom_toolsets_from_cli=custom_toolsets,
custom_runbooks=custom_runbooks,
allowed_builtin_toolsets=allowed_builtin_toolsets,
)
ai = config.create_console_issue_investigator()
source = config.create_opsgenie_source()
Expand Down
42 changes: 37 additions & 5 deletions holmes/plugins/toolsets/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,18 @@ def load_python_toolsets(dal: Optional[SupabaseDal]) -> List[Toolset]:
return toolsets


def load_builtin_toolsets(dal: Optional[SupabaseDal] = None) -> List[Toolset]:
def load_builtin_toolsets(
dal: Optional[SupabaseDal] = None,
allowed_builtin_toolsets: Optional[List[str]] = None,
) -> List[Toolset]:
"""
Load builtin toolsets from YAML files and Python modules.

If allowed_builtin_toolsets is provided, filtering preserves the original
discovery order (loader order) rather than reordering by the allowlist.
This maintains consistent behavior where toolsets appear in the order
they were discovered during loading.
"""
all_toolsets: List[Toolset] = []
logging.debug(f"loading toolsets from {THIS_DIR}")

Expand All @@ -115,15 +126,36 @@ def load_builtin_toolsets(dal: Optional[SupabaseDal] = None) -> List[Toolset]:
toolsets_from_file = load_toolsets_from_file(path, strict_check=True)
all_toolsets.extend(toolsets_from_file)

all_toolsets.extend(load_python_toolsets(dal=dal)) # type: ignore
all_toolsets.extend(load_python_toolsets(dal=dal))

# Apply filter if specified
if allowed_builtin_toolsets is not None:
# Convert to set for O(1) lookup performance
requested_names = set(allowed_builtin_toolsets)
actual_names = {toolset.name for toolset in all_toolsets}
invalid_names = requested_names - actual_names

if invalid_names:
logging.warning(
f"Unknown builtin toolsets specified: {sorted(invalid_names)}"
)

# Filter to only allowed toolsets using set for efficient lookup
# Note: This preserves the original discovery order rather than
# reordering by the allowlist
all_toolsets = [
toolset
for toolset in all_toolsets
if toolset.name in requested_names # O(1) lookup
]

Comment on lines +131 to 151

@coderabbitai coderabbitai Bot Aug 14, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Allowlist still loads/instantiates all toolsets; short-circuit and avoid unnecessary work

Currently, all YAML and Python toolsets are loaded/instantiated and only then filtered. This undermines the goal of speeding up scenarios where most builtins aren’t needed.

Two concrete improvements:

  • Early return for an empty allowlist (fast path).
  • After loading YAML toolsets, if the allowlist is fully satisfied by YAML names, skip constructing Python toolsets altogether.

Apply a minimal fast-path within this hunk:

     # Apply filter if specified
     if allowed_builtin_toolsets is not None:
+        # Fast path: empty allowlist means nothing to load
+        if len(allowed_builtin_toolsets) == 0:
+            return []
         # Convert to set for O(1) lookup performance
         requested_names = set(allowed_builtin_toolsets)

Outside the selected lines, consider this structure to skip Python toolsets when not needed:

def load_builtin_toolsets(
    dal: Optional[SupabaseDal] = None,
    allowed_builtin_toolsets: Optional[List[str]] = None,
) -> List[Toolset]:
    logging.debug(f"loading toolsets from {THIS_DIR}")
    if allowed_builtin_toolsets is not None:
        requested = set(allowed_builtin_toolsets)
        if not requested:
            return []

    # Load YAML toolsets first (sorted for determinism)
    all_toolsets: List[Toolset] = []
    for filename in sorted(os.listdir(THIS_DIR)):
        if not filename.endswith(".yaml"):
            continue
        if filename == "kubernetes_logs.yaml" and not USE_LEGACY_KUBERNETES_LOGS:
            continue
        all_toolsets.extend(load_toolsets_from_file(os.path.join(THIS_DIR, filename), strict_check=True))

    # If an allowlist is provided and all requested names are covered by YAML, skip Python loading
    if allowed_builtin_toolsets is not None:
        yaml_names = {t.name for t in all_toolsets}
        if requested.issubset(yaml_names):
            filtered = [t for t in all_toolsets if t.name in requested]
            for t in filtered:
                t.type = ToolsetType.BUILTIN
                t.path = None
            return filtered

    # Otherwise, load Python toolsets and filter as today
    all_toolsets.extend(load_python_toolsets(dal=dal))
    if allowed_builtin_toolsets is not None:
        requested = set(allowed_builtin_toolsets)
        invalid = requested - {t.name for t in all_toolsets}
        if invalid:
            logging.warning("Unknown builtin toolsets specified: %s", ", ".join(sorted(invalid)))
        all_toolsets = [t for t in all_toolsets if t.name in requested]

    for t in all_toolsets:
        t.type = ToolsetType.BUILTIN
        t.path = None
    return all_toolsets
🤖 Prompt for AI Agents 
In holmes/plugins/toolsets/__init__.py around lines 131 to 151, the allowlist is
applied after all YAML and Python toolsets are loaded which still instantiates
everything; change the logic to (1) early-return [] when
allowed_builtin_toolsets is provided as an empty list, (2) load YAML toolsets
first, build requested = set(allowed_builtin_toolsets) and if requested is a
subset of the loaded YAML names then filter those YAML toolsets, set each
filtered toolset's type to ToolsetType.BUILTIN and path to None, and return the
filtered list without loading Python toolsets, and (3) otherwise proceed to load
Python toolsets and perform the existing allowlist filtering and unknown-name
logging; ensure lookups use sets for O(1) checks and preserve discovery order
when filtering.

# disable built-in toolsets by default, and the user can enable them explicitly in config.
# disable builtin toolsets by default, and the user can enable them explicitly in config.
for toolset in all_toolsets:
toolset.type = ToolsetType.BUILTIN
# dont' expose build-in toolsets path
# don't expose builtin toolsets path
toolset.path = None

return all_toolsets # type: ignore
return all_toolsets


def is_old_toolset_config(
Expand Down
Loading