From 7891f93ddf4b11e10cf98967ed8ca82046fa5864 Mon Sep 17 00:00:00 2001 From: adchaudh Date: Thu, 24 Sep 2026 08:39:11 -0700 Subject: [PATCH 1/2] feat(python): register discovered tools from tool search results --- python/README.md | 29 +++++++ python/copilot/session.py | 67 ++++++++++++++++ python/copilot/tools.py | 2 + python/test_session.py | 163 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 261 insertions(+) diff --git a/python/README.md b/python/README.md index efe960a7f5..fc2c31d037 100644 --- a/python/README.md +++ b/python/README.md @@ -494,6 +494,35 @@ async def lookup_issue(params: LookupParams) -> str: # your logic ``` +#### Discovering Tools During Tool Search + +An override named `tool_search_tool` can return tool definitions that were not +registered when the session started. Return them as `ToolResult(tools=[...])`: + +```python +async def search_tools(invocation: ToolInvocation) -> ToolResult: + discovered = await find_tools(invocation.arguments) + return ToolResult( + text_result_for_llm="Loaded matching tools.", + tools=[ + Tool( + name=item.name, + description=item.description, + parameters=item.schema, + handler=item.handler, + ) + for item in discovered + ], + ) +``` + +Register `search_tools` as a `Tool` named `tool_search_tool` with +`overrides_built_in_tool=True`. Every returned tool needs a handler. The SDK +registers the definitions and handlers before returning their names to the +model as tool references. This uses +the runtime's experimental live-tool registration RPC; callers do not need to +call `session.tools.set` or handle `external_tool.requested` themselves. + ## Auto routing tiers Change the Auto routing preference without changing the selected model. The runtime does not apply the preference immediately: it records the request and commits it only when a later user turn using the `auto` model successfully obtains a usable model from the provider, so a `pending` status confirms acceptance rather than effect. Only the most recent request survives. diff --git a/python/copilot/session.py b/python/copilot/session.py index 6534e4cbae..279903b44d 100644 --- a/python/copilot/session.py +++ b/python/copilot/session.py @@ -45,6 +45,7 @@ LogRequest, MCPOauthHandlePendingRequest, MCPOauthPendingRequestResponse, + MCPServerConfigDeferTools, ModelSwitchAutoTierResult, ModelSwitchToRequest, PermissionDecision, @@ -52,10 +53,12 @@ PermissionDecisionContext, PermissionDecisionRequest, PermissionDecisionUserNotAvailable, + ProtocolExternalToolDefinition, ProviderTokenAcquireRequest, ProviderTokenAcquireResult, SessionLogLevel, SessionRpc, + ToolsSetRequest, UIElicitationRequest, UIElicitationResponse, UIElicitationResponseAction, @@ -1647,7 +1650,9 @@ def __init__( self._event_handlers: set[Callable[[SessionEvent], None]] = set() self._event_handlers_lock = threading.Lock() self._tool_handlers: dict[str, ToolHandler] = {} + self._registered_tools: dict[str, Tool] = {} self._tool_handlers_lock = threading.Lock() + self._tool_catalog_lock = asyncio.Lock() self._pending_external_tools: dict[str, asyncio.Task[None]] = {} self._permission_handler: _PermissionHandlerFn | None = None self._permission_handler_lock = threading.Lock() @@ -2426,6 +2431,16 @@ async def _execute_tool_and_respond( else: tool_result = result # type: ignore[assignment] + if tool_result.tools is not None: + if tool_name != _TOOL_SEARCH_TOOL_NAME or tool_result.result_type != "success": + raise ValueError( + "ToolResult.tools is only valid for successful tool_search_tool calls" + ) + names = await self._register_discovered_tools(tool_result.tools) + tool_result.tool_references = list( + dict.fromkeys([*(tool_result.tool_references or []), *names]) + ) + # Exception-originated failures (from define_tool's exception handler) are # sent via the top-level error param so the CLI formats them with its # standard "Failed to execute..." message. Deliberate user-returned @@ -2879,6 +2894,7 @@ def _register_tools(self, tools: list[Tool] | None) -> None: """ with self._tool_handlers_lock: self._tool_handlers.clear() + self._registered_tools = {tool.name: tool for tool in tools or []} if not tools: return for tool in tools: @@ -2886,6 +2902,56 @@ def _register_tools(self, tools: list[Tool] | None) -> None: continue self._tool_handlers[tool.name] = tool.handler + async def _register_discovered_tools(self, tools: list[Tool]) -> list[str]: + if not tools: + return [] + names = [tool.name for tool in tools] + if any(not name for name in names) or len(set(names)) != len(names): + raise ValueError("discovered tools must have unique, nonempty names") + if any(tool.handler is None for tool in tools): + raise ValueError("discovered tools must have handlers") + + async with self._tool_catalog_lock: + with self._tool_handlers_lock: + previous = self._registered_tools + merged = {**previous, **{tool.name: tool for tool in tools}} + definitions = [ + ProtocolExternalToolDefinition( + name=tool.name, + description=tool.description, + parameters=tool.parameters, + overrides_built_in_tool=tool.overrides_built_in_tool, + skip_permission=tool.skip_permission, + defer=MCPServerConfigDeferTools(tool.defer or "auto"), + metadata=tool.metadata, + is_terminal=tool.is_terminal, + ) + for tool in merged.values() + ] + self._registered_tools = merged + # The CLI may apply tools.set before its reply reaches us. + provisioned = [tool.name for tool in tools if tool.name not in self._tool_handlers] + for tool in tools: + if tool.name in provisioned and tool.handler is not None: + self._tool_handlers[tool.name] = tool.handler + try: + await self.rpc.tools.set(ToolsSetRequest(tools=definitions)) + except JsonRpcError: + # A response error means the CLI rejected the update. Transport + # failures are ambiguous, so retain handlers for tools it may have applied. + with self._tool_handlers_lock: + if not self._destroyed: + self._registered_tools = previous + for name in provisioned: + self._tool_handlers.pop(name, None) + raise + with self._tool_handlers_lock: + if not self._destroyed: + for tool in tools: + if tool.handler is not None: + self._tool_handlers[tool.name] = tool.handler + return names + def _get_tool_handler(self, name: str) -> ToolHandler | None: """ Retrieve a registered tool handler by name. @@ -3287,6 +3353,7 @@ async def disconnect(self) -> None: self._event_handlers.clear() with self._tool_handlers_lock: self._tool_handlers.clear() + self._registered_tools.clear() with self._permission_handler_lock: self._permission_handler = None with self._command_handlers_lock: diff --git a/python/copilot/tools.py b/python/copilot/tools.py index ad0bcb41bd..8c5c37fa88 100644 --- a/python/copilot/tools.py +++ b/python/copilot/tools.py @@ -52,6 +52,8 @@ class ToolResult: session_log: str | None = None tool_telemetry: dict[str, Any] | None = None tool_references: list[str] | None = None + # For tool_search_tool: register these definitions and make them callable now. + tools: list[Tool] | None = field(default=None, kw_only=True) _from_exception: bool = field(default=False, repr=False) diff --git a/python/test_session.py b/python/test_session.py index 5750364bb5..36d148f312 100644 --- a/python/test_session.py +++ b/python/test_session.py @@ -10,6 +10,7 @@ from _session_test_helpers import get_next_event_of_type, wait_for_event from copilot import AgentMessageSource, MessageSource +from copilot._jsonrpc import JsonRpcError from copilot.session import Attachment, CopilotSession from copilot.session_events import ( AssistantMessageData, @@ -495,6 +496,168 @@ async def blocked_tool(_invocation): client.request.assert_not_awaited() +@pytest.mark.asyncio +async def test_tool_search_result_loads_new_tool_before_resuming_the_model(): + requests = [] + replied = asyncio.Event() + + async def request(method, params): + requests.append((method, params)) + if method == "session.tools.getCurrentMetadata": + return {"tools": []} + if method == "session.tools.set": + return {} + if method == "session.tools.handlePendingToolCall": + replied.set() + return {"success": True} + raise AssertionError(method) + + client = Mock(request=AsyncMock(side_effect=request)) + session = CopilotSession("session-1", client) + discovered = Tool( + "get_shipping_eta", + "Look up a shipping ETA", + lambda _invocation: ToolResult(text_result_for_llm="2 business days"), + ) + + async def search(_invocation): + return ToolResult(tools=[discovered]) + + session._register_tools( + [Tool("tool_search_tool", "Find tools", search, overrides_built_in_tool=True)] + ) + session._dispatch_event( + _event( + ExternalToolRequestedData( + request_id="search-request", + session_id="session-1", + tool_call_id="search-call", + tool_name="tool_search_tool", + arguments={"paths": ["shipping"]}, + ), + SessionEventType.EXTERNAL_TOOL_REQUESTED, + ) + ) + await asyncio.wait_for(replied.wait(), timeout=1) + + methods = [method for method, _ in requests] + assert methods.index("session.tools.set") < methods.index("session.tools.handlePendingToolCall") + definitions = next( + params["tools"] for method, params in requests if method == "session.tools.set" + ) + assert {tool["name"] for tool in definitions} == {"tool_search_tool", "get_shipping_eta"} + search_result = next( + params["result"] + for method, params in requests + if method == "session.tools.handlePendingToolCall" + ) + assert search_result["toolReferences"] == ["get_shipping_eta"] + + replied.clear() + session._dispatch_event( + _event( + ExternalToolRequestedData( + request_id="eta-request", + session_id="session-1", + tool_call_id="eta-call", + tool_name="get_shipping_eta", + ), + SessionEventType.EXTERNAL_TOOL_REQUESTED, + ) + ) + await asyncio.wait_for(replied.wait(), timeout=1) + assert requests[-1][1]["result"]["textResultForLlm"] == "2 business days" + + +@pytest.mark.asyncio +async def test_discovered_handler_is_ready_while_registration_rpc_is_pending(): + handled = asyncio.Event() + results = [] + + async def request(method, params): + if method == "session.tools.set": + session._dispatch_event( + _event( + ExternalToolRequestedData( + request_id="independent-request", + session_id="session-1", + tool_call_id="independent-call", + tool_name="get_shipping_eta", + ), + SessionEventType.EXTERNAL_TOOL_REQUESTED, + ) + ) + await asyncio.wait_for(handled.wait(), timeout=1) + return {} + if method == "session.tools.handlePendingToolCall": + results.append(params) + handled.set() + return {"success": True} + raise AssertionError(method) + + session = CopilotSession("session-1", Mock(request=AsyncMock(side_effect=request))) + discovered = Tool("get_shipping_eta", "Look up an ETA", lambda _: ToolResult("2 days")) + + assert await session._register_discovered_tools([discovered]) == ["get_shipping_eta"] + assert results[0]["result"]["textResultForLlm"] == "2 days" + + +@pytest.mark.asyncio +async def test_tool_search_registration_failure_does_not_install_local_handler(): + replied = asyncio.Event() + responses = [] + + async def request(method, params): + if method == "session.tools.getCurrentMetadata": + return {"tools": []} + if method == "session.tools.set": + raise JsonRpcError(-32000, "registration rejected") + if method == "session.tools.handlePendingToolCall": + responses.append(params) + replied.set() + return {"success": True} + raise AssertionError(method) + + session = CopilotSession("session-1", Mock(request=AsyncMock(side_effect=request))) + discovered = Tool("get_shipping_eta", "Look up an ETA", lambda _: ToolResult("2 days")) + session._register_tools( + [ + Tool( + "tool_search_tool", + "Find tools", + lambda _: ToolResult(tools=[discovered]), + overrides_built_in_tool=True, + ) + ] + ) + session._dispatch_event( + _event( + ExternalToolRequestedData( + request_id="search-request", + session_id="session-1", + tool_call_id="search-call", + tool_name="tool_search_tool", + ), + SessionEventType.EXTERNAL_TOOL_REQUESTED, + ) + ) + await asyncio.wait_for(replied.wait(), timeout=1) + + assert session._get_tool_handler("get_shipping_eta") is None + assert "registration rejected" in responses[0]["error"] + + +@pytest.mark.asyncio +async def test_discovered_tool_requires_handler(): + client = Mock(request=AsyncMock()) + session = CopilotSession("session-1", client) + + with pytest.raises(ValueError, match="discovered tools must have handlers"): + await session._register_discovered_tools([Tool("get_shipping_eta", "Look up an ETA")]) + + client.request.assert_not_awaited() + + @pytest.mark.asyncio async def test_disconnect_from_tool_task_does_not_cancel_detach_request(): client = Mock() From ac163cee781bfa92576dbfd84fa85623cd839bf4 Mon Sep 17 00:00:00 2001 From: adchaudh Date: Thu, 24 Sep 2026 11:27:13 -0700 Subject: [PATCH 2/2] fix(python): validate and preserve additive tool discovery state --- python/README.md | 19 ++++++++--- python/copilot/session.py | 26 +++++++++----- python/test_session.py | 72 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 103 insertions(+), 14 deletions(-) diff --git a/python/README.md b/python/README.md index fc2c31d037..ea959e3208 100644 --- a/python/README.md +++ b/python/README.md @@ -517,11 +517,20 @@ async def search_tools(invocation: ToolInvocation) -> ToolResult: ``` Register `search_tools` as a `Tool` named `tool_search_tool` with -`overrides_built_in_tool=True`. Every returned tool needs a handler. The SDK -registers the definitions and handlers before returning their names to the -model as tool references. This uses -the runtime's experimental live-tool registration RPC; callers do not need to -call `session.tools.set` or handle `external_tool.requested` themselves. +`overrides_built_in_tool=True`, and enable tool search for the session. The runtime +must expose that search tool; current runtimes require a deferred tool inventory +before they activate it. Returning definitions does not change that activation rule. + +Every returned tool needs a handler and a new, unique name. The SDK preserves +previously registered tools, installs the new handlers, and registers the expanded +catalog before completing the search. For tools already loaded, return their names +in `tool_references` instead of redeclaring them. + +This feature uses the runtime's experimental `session.tools.set` RPC internally. +Do not mix it with direct calls to that RPC: the SDK owns the complete custom-tool +catalog for this connection to the session. On an explicit RPC rejection it rolls +back the new handlers. On a timeout or connection failure, it retains them because +the runtime may already have applied the registration. ## Auto routing tiers diff --git a/python/copilot/session.py b/python/copilot/session.py index 279903b44d..0c95adfbd3 100644 --- a/python/copilot/session.py +++ b/python/copilot/session.py @@ -11,6 +11,7 @@ import asyncio import functools import inspect +import json import logging import os import pathlib @@ -2432,6 +2433,11 @@ async def _execute_tool_and_respond( tool_result = result # type: ignore[assignment] if tool_result.tools is not None: + if ( + self._destroyed + or self._pending_external_tools.get(request_id) is not asyncio.current_task() + ): + return if tool_name != _TOOL_SEARCH_TOOL_NAME or tool_result.result_type != "success": raise ValueError( "ToolResult.tools is only valid for successful tool_search_tool calls" @@ -2913,7 +2919,13 @@ async def _register_discovered_tools(self, tools: list[Tool]) -> list[str]: async with self._tool_catalog_lock: with self._tool_handlers_lock: + if self._destroyed: + raise RuntimeError("Cannot register discovered tools on a disconnected session") previous = self._registered_tools + if previous.keys() & set(names): + raise ValueError( + "discovered tool names are already registered; use tool_references instead" + ) merged = {**previous, **{tool.name: tool for tool in tools}} definitions = [ ProtocolExternalToolDefinition( @@ -2928,28 +2940,24 @@ async def _register_discovered_tools(self, tools: list[Tool]) -> list[str]: ) for tool in merged.values() ] + request = ToolsSetRequest(tools=definitions) + json.dumps(request.to_dict()) # Validate before publishing handlers. self._registered_tools = merged # The CLI may apply tools.set before its reply reaches us. - provisioned = [tool.name for tool in tools if tool.name not in self._tool_handlers] for tool in tools: - if tool.name in provisioned and tool.handler is not None: + if tool.handler is not None: self._tool_handlers[tool.name] = tool.handler try: - await self.rpc.tools.set(ToolsSetRequest(tools=definitions)) + await self.rpc.tools.set(request) except JsonRpcError: # A response error means the CLI rejected the update. Transport # failures are ambiguous, so retain handlers for tools it may have applied. with self._tool_handlers_lock: if not self._destroyed: self._registered_tools = previous - for name in provisioned: + for name in names: self._tool_handlers.pop(name, None) raise - with self._tool_handlers_lock: - if not self._destroyed: - for tool in tools: - if tool.handler is not None: - self._tool_handlers[tool.name] = tool.handler return names def _get_tool_handler(self, name: str) -> ToolHandler | None: diff --git a/python/test_session.py b/python/test_session.py index 36d148f312..7e202778f3 100644 --- a/python/test_session.py +++ b/python/test_session.py @@ -658,6 +658,78 @@ async def test_discovered_tool_requires_handler(): client.request.assert_not_awaited() +@pytest.mark.asyncio +async def test_discovery_cannot_replace_existing_tools(): + client = Mock(request=AsyncMock()) + session = CopilotSession("session-1", client) + original = Tool("existing", "Existing tool", lambda _: ToolResult("original")) + session._register_tools([original]) + + with pytest.raises(ValueError, match="already registered"): + await session._register_discovered_tools( + [Tool("existing", "Replacement", lambda _: ToolResult("replacement"))] + ) + + client.request.assert_not_awaited() + assert session._get_tool_handler("existing") is original.handler + + +@pytest.mark.asyncio +async def test_concurrent_discoveries_preserve_both_catalogs(): + catalogs = [] + + async def request(method, params): + assert method == "session.tools.set" + catalogs.append({tool["name"] for tool in params["tools"]}) + await asyncio.sleep(0) + return {} + + session = CopilotSession("session-1", Mock(request=AsyncMock(side_effect=request))) + first = Tool("first", "First", lambda _: ToolResult("first")) + second = Tool("second", "Second", lambda _: ToolResult("second")) + await asyncio.gather( + session._register_discovered_tools([first]), + session._register_discovered_tools([second]), + ) + + assert catalogs[-1] == {"first", "second"} + assert session._get_tool_handler("first") is first.handler + assert session._get_tool_handler("second") is second.handler + + +@pytest.mark.asyncio +async def test_ambiguous_registration_failure_retains_callable_tools(): + client = Mock(request=AsyncMock(side_effect=TimeoutError("response lost"))) + session = CopilotSession("session-1", client) + tool = Tool("discovered", "Discovered", lambda _: ToolResult("result")) + + with pytest.raises(TimeoutError): + await session._register_discovered_tools([tool]) + + # The CLI may have applied the update before the response was lost. + assert session._get_tool_handler("discovered") is tool.handler + assert session._registered_tools["discovered"] is tool + + +@pytest.mark.asyncio +async def test_invalid_discovered_schema_does_not_publish_local_state(): + client = Mock(request=AsyncMock()) + session = CopilotSession("session-1", client) + tool = Tool( + "invalid", + "Invalid schema", + lambda _: ToolResult("result"), + parameters={"not_json": object()}, + ) + + with pytest.raises(TypeError): + await session._register_discovered_tools([tool]) + + client.request.assert_not_awaited() + assert session._get_tool_handler("invalid") is None + assert "invalid" not in session._registered_tools + + @pytest.mark.asyncio async def test_disconnect_from_tool_task_does_not_cancel_detach_request(): client = Mock()