diff --git a/docs/openapi.json b/docs/openapi.json index 441612854..eb249cb2b 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -1566,6 +1566,21 @@ "mode": { "$ref": "#/components/schemas/QueryMode", "default": "ask" + }, + "available_tool_ui_ids": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array", + "maxItems": 256 + }, + { + "type": "null" + } + ], + "title": "Available Tool Ui Ids" } }, "additionalProperties": false, diff --git a/ols/app/endpoints/ols.py b/ols/app/endpoints/ols.py index fe61eeec8..5b26f5f1a 100644 --- a/ols/app/endpoints/ols.py +++ b/ols/app/endpoints/ols.py @@ -540,6 +540,7 @@ def generate_response( client_headers=client_headers, streaming=streaming, audit_ctx=audit_ctx, + available_tool_ui_ids=llm_request.available_tool_ui_ids, ) if streaming: return docs_summarizer.generate_response( diff --git a/ols/app/models/models.py b/ols/app/models/models.py index 3754c0f75..4e7c96876 100644 --- a/ols/app/models/models.py +++ b/ols/app/models/models.py @@ -92,6 +92,7 @@ class LLMRequest(BaseModel): media_type: Optional[str] = MEDIA_TYPE_TEXT mcp_headers: Optional[dict[str, dict[str, str]]] = None mode: QueryMode = QueryMode.ASK + available_tool_ui_ids: Optional[list[str]] = Field(default=None, max_length=256) # provides examples for /docs endpoint model_config = { diff --git a/ols/src/query_helpers/docs_summarizer.py b/ols/src/query_helpers/docs_summarizer.py index b0365d037..b5d10115f 100644 --- a/ols/src/query_helpers/docs_summarizer.py +++ b/ols/src/query_helpers/docs_summarizer.py @@ -31,7 +31,12 @@ from ols.src.skills.skills_rag import create_skill_support_tool from ols.src.tools.offloaded_content import OffloadManager from ols.utils.audit_logger import AuditContext -from ols.utils.mcp_utils import ClientHeaders, build_mcp_config, get_mcp_tools +from ols.utils.mcp_utils import ( + ClientHeaders, + build_mcp_config, + filter_tools_by_available_ui, + get_mcp_tools, +) from ols.utils.token_handler import ( PromptTooLongError, TokenBudgetTracker, @@ -63,6 +68,7 @@ def __init__( client_headers: ClientHeaders | None = None, streaming: bool = False, audit_ctx: Optional[AuditContext] = None, + available_tool_ui_ids: Optional[list[str]] = None, **kwargs: object, ) -> None: """Initialize the DocsSummarizer. @@ -72,6 +78,7 @@ def __init__( client_headers: Optional client-provided MCP headers for authentication streaming: Whether this summarizer is used for the streaming endpoint audit_ctx: Audit context for structured event logging + available_tool_ui_ids: Tool UI extension IDs available on the client *args: Additional positional arguments passed to the parent class **kwargs: Additional keyword arguments passed to the parent class """ @@ -80,6 +87,7 @@ def __init__( self._prepare_llm() self.verbose = config.ols_config.logging_config.app_log_level == logging.DEBUG self.streaming = streaming + self.available_tool_ui_ids = available_tool_ui_ids self._cluster_version = ( K8sClientSingleton.get_cluster_version() if self._mode == constants.QueryMode.TROUBLESHOOTING @@ -423,6 +431,9 @@ async def generate_response( # noqa: C901 # pylint: disable=too-many-branches, all_mcp_tools = await get_mcp_tools( mcp_tools_query, self.user_token, self.client_headers ) + all_mcp_tools = filter_tools_by_available_ui( + all_mcp_tools, self.available_tool_ui_ids + ) if skill is not None and skill_content is not None and has_support_files: all_mcp_tools.append(create_skill_support_tool(skill)) tool_definitions_text = self._serialized_tool_definitions_text(all_mcp_tools) diff --git a/ols/utils/mcp_utils.py b/ols/utils/mcp_utils.py index f5f16d0f2..cff7f2191 100644 --- a/ols/utils/mcp_utils.py +++ b/ols/utils/mcp_utils.py @@ -176,6 +176,49 @@ def _normalize_tool_schema(tool: StructuredTool) -> None: schema.setdefault("required", []) +def filter_tools_by_available_ui( + tools: list[StructuredTool], + available_ui_ids: list[str] | None, +) -> list[StructuredTool]: + """Filter out tools whose UI extension is not available on the client. + + Tools that declare an ``olsUi.id`` in their ``_meta`` metadata are only + useful when the corresponding console extension is installed. When the + client provides the list of available extension IDs, tools whose + ``olsUi.id`` is not in that list are excluded so the LLM never calls them. + + Args: + tools: All resolved MCP tools. + available_ui_ids: Extension IDs the client can render, or ``None`` + to skip filtering entirely (backward compatibility). + + Returns: + Filtered list of tools. + """ + if available_ui_ids is None: + return tools + + available_set = set(available_ui_ids) + filtered: list[StructuredTool] = [] + for tool in tools: + meta = ( + (tool.metadata or {}).get("_meta") + if isinstance(tool.metadata, dict) + else None + ) + ols_ui = meta.get("olsUi") if isinstance(meta, dict) else None + ols_ui_id = ols_ui.get("id") if isinstance(ols_ui, dict) else None + if ols_ui_id is not None and ols_ui_id not in available_set: + logger.info( + "Excluding tool '%s': olsUi.id '%s' not in available_tool_ui_ids", + tool.name, + ols_ui_id, + ) + continue + filtered.append(tool) + return filtered + + async def gather_mcp_tools( mcp_servers: MCPServersDict, allowed_tool_names: Optional[set[str]] = None ) -> list[StructuredTool]: diff --git a/tests/unit/app/models/test_models.py b/tests/unit/app/models/test_models.py index 6dfe523ea..8aed49e02 100644 --- a/tests/unit/app/models/test_models.py +++ b/tests/unit/app/models/test_models.py @@ -65,6 +65,18 @@ def test_llm_request_optional_inputs(): assert llm_request.provider == provider assert llm_request.model == model + @staticmethod + def test_llm_request_available_tool_ui_ids(): + """Test available_tool_ui_ids field on LLMRequest.""" + request = LLMRequest(query="test") + assert request.available_tool_ui_ids is None + + request = LLMRequest(query="test", available_tool_ui_ids=["perses-dashboard"]) + assert request.available_tool_ui_ids == ["perses-dashboard"] + + request = LLMRequest(query="test", available_tool_ui_ids=[]) + assert request.available_tool_ui_ids == [] + @staticmethod def test_llm_request_provider_and_model(): """Test the LLMRequest model with provider and model.""" diff --git a/tests/unit/utils/test_mcp_utils.py b/tests/unit/utils/test_mcp_utils.py index 9b912f3c5..513766693 100644 --- a/tests/unit/utils/test_mcp_utils.py +++ b/tests/unit/utils/test_mcp_utils.py @@ -9,6 +9,7 @@ from ols.app.models.config import MCPServerConfig from ols.utils.mcp_utils import ( build_mcp_config, + filter_tools_by_available_ui, gather_mcp_tools, get_mcp_tools, get_servers_requiring_client_headers, @@ -647,3 +648,76 @@ async def test_no_servers_configured(self): result = await get_mcp_tools("test query") assert result == [] + + +class TestFilterToolsByAvailableUI: + """Tests for filter_tools_by_available_ui function.""" + + @staticmethod + def _make_tool(name: str, metadata: dict | None = None) -> MagicMock: + tool = MagicMock(spec=StructuredTool) + tool.name = name + tool.metadata = metadata + return tool + + def test_none_available_ids_returns_all(self): + """When available_ui_ids is None, all tools pass through.""" + tool_with_ui = self._make_tool( + "perses", {"_meta": {"olsUi": {"id": "perses-dashboard"}}} + ) + tool_without_ui = self._make_tool("kubectl", {"mcp_server": "k8s"}) + result = filter_tools_by_available_ui([tool_with_ui, tool_without_ui], None) + assert len(result) == 2 + + def test_empty_list_excludes_tools_with_ui_id(self): + """Empty available list excludes tools that declare olsUi.id.""" + tool_with_ui = self._make_tool( + "perses", {"_meta": {"olsUi": {"id": "perses-dashboard"}}} + ) + tool_without_ui = self._make_tool("kubectl", {"mcp_server": "k8s"}) + result = filter_tools_by_available_ui([tool_with_ui, tool_without_ui], []) + assert len(result) == 1 + assert result[0].name == "kubectl" + + def test_matching_id_passes(self): + """Tool with matching olsUi.id passes the filter.""" + tool = self._make_tool( + "perses", {"_meta": {"olsUi": {"id": "perses-dashboard"}}} + ) + result = filter_tools_by_available_ui([tool], ["perses-dashboard"]) + assert len(result) == 1 + assert result[0].name == "perses" + + def test_non_matching_id_excluded(self): + """Tool with non-matching olsUi.id is excluded.""" + tool = self._make_tool( + "perses", {"_meta": {"olsUi": {"id": "perses-dashboard"}}} + ) + result = filter_tools_by_available_ui([tool], ["other-ui"]) + assert len(result) == 0 + + def test_tool_without_meta_passes(self): + """Tool without _meta always passes.""" + tool = self._make_tool("kubectl", {"mcp_server": "k8s"}) + result = filter_tools_by_available_ui([tool], []) + assert len(result) == 1 + + def test_tool_with_none_metadata_passes(self): + """Tool with None metadata always passes.""" + tool = self._make_tool("kubectl", None) + result = filter_tools_by_available_ui([tool], []) + assert len(result) == 1 + + def test_tool_with_meta_but_no_ols_ui_passes(self): + """Tool with _meta but no olsUi field always passes.""" + tool = self._make_tool( + "kubectl", {"_meta": {"ui": {"resourceUri": "ui://foo"}}} + ) + result = filter_tools_by_available_ui([tool], []) + assert len(result) == 1 + + def test_tool_with_none_ols_ui_passes(self): + """Tool with olsUi set to None always passes.""" + tool = self._make_tool("kubectl", {"_meta": {"olsUi": None}}) + result = filter_tools_by_available_ui([tool], []) + assert len(result) == 1