From 778b7de9cb590926d7778826c4a0fd84106d8211 Mon Sep 17 00:00:00 2001 From: Omar Abdelwahab Date: Wed, 12 Nov 2025 15:08:00 -0800 Subject: [PATCH] fix: add authorization parameter to ToolRuntimeRouter and routing table The auto-routing layer was missing the authorization parameter: - ToolRuntimeRouter.invoke_tool() now accepts and passes authorization - ToolRuntimeRouter.list_runtime_tools() now accepts and passes authorization - ToolGroupsRoutingTable.list_tools() now accepts and forwards authorization - ToolGroupsRoutingTable._index_tools() now accepts and uses authorization This fixes the '__autorouted__' provider signature mismatch error in CI. --- src/llama_stack/core/routers/tool_runtime.py | 7 ++++--- src/llama_stack/core/routing_tables/toolgroups.py | 10 ++++++---- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/src/llama_stack/core/routers/tool_runtime.py b/src/llama_stack/core/routers/tool_runtime.py index fb13d94a4..fe170eeb7 100644 --- a/src/llama_stack/core/routers/tool_runtime.py +++ b/src/llama_stack/core/routers/tool_runtime.py @@ -36,16 +36,17 @@ class ToolRuntimeRouter(ToolRuntime): logger.debug("ToolRuntimeRouter.shutdown") pass - async def invoke_tool(self, tool_name: str, kwargs: dict[str, Any]) -> Any: + async def invoke_tool(self, tool_name: str, kwargs: dict[str, Any], authorization: str | None = None) -> Any: logger.debug(f"ToolRuntimeRouter.invoke_tool: {tool_name}") provider = await self.routing_table.get_provider_impl(tool_name) return await provider.invoke_tool( tool_name=tool_name, kwargs=kwargs, + authorization=authorization, ) async def list_runtime_tools( - self, tool_group_id: str | None = None, mcp_endpoint: URL | None = None + self, tool_group_id: str | None = None, mcp_endpoint: URL | None = None, authorization: str | None = None ) -> ListToolDefsResponse: logger.debug(f"ToolRuntimeRouter.list_runtime_tools: {tool_group_id}") - return await self.routing_table.list_tools(tool_group_id) + return await self.routing_table.list_tools(tool_group_id, authorization=authorization) diff --git a/src/llama_stack/core/routing_tables/toolgroups.py b/src/llama_stack/core/routing_tables/toolgroups.py index 2d47bbb17..0761c5582 100644 --- a/src/llama_stack/core/routing_tables/toolgroups.py +++ b/src/llama_stack/core/routing_tables/toolgroups.py @@ -43,7 +43,7 @@ class ToolGroupsRoutingTable(CommonRoutingTableImpl, ToolGroups): routing_key = self.tool_to_toolgroup[routing_key] return await super().get_provider_impl(routing_key, provider_id) - async def list_tools(self, toolgroup_id: str | None = None) -> ListToolDefsResponse: + async def list_tools(self, toolgroup_id: str | None = None, authorization: str | None = None) -> ListToolDefsResponse: if toolgroup_id: if group_id := parse_toolgroup_from_toolgroup_name_pair(toolgroup_id): toolgroup_id = group_id @@ -55,7 +55,7 @@ class ToolGroupsRoutingTable(CommonRoutingTableImpl, ToolGroups): for toolgroup in toolgroups: if toolgroup.identifier not in self.toolgroups_to_tools: try: - await self._index_tools(toolgroup) + await self._index_tools(toolgroup, authorization=authorization) except AuthenticationRequiredError: # Send authentication errors back to the client so it knows # that it needs to supply credentials for remote MCP servers. @@ -70,9 +70,11 @@ class ToolGroupsRoutingTable(CommonRoutingTableImpl, ToolGroups): return ListToolDefsResponse(data=all_tools) - async def _index_tools(self, toolgroup: ToolGroup): + async def _index_tools(self, toolgroup: ToolGroup, authorization: str | None = None): provider_impl = await super().get_provider_impl(toolgroup.identifier, toolgroup.provider_id) - tooldefs_response = await provider_impl.list_runtime_tools(toolgroup.identifier, toolgroup.mcp_endpoint) + tooldefs_response = await provider_impl.list_runtime_tools( + toolgroup.identifier, toolgroup.mcp_endpoint, authorization=authorization + ) tooldefs = tooldefs_response.data for t in tooldefs: