diff --git a/python/packages/claude/agent_framework_claude/_agent.py b/python/packages/claude/agent_framework_claude/_agent.py index e34f507c07f..ecdf65bf3cb 100644 --- a/python/packages/claude/agent_framework_claude/_agent.py +++ b/python/packages/claude/agent_framework_claude/_agent.py @@ -8,6 +8,7 @@ import logging import sys from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, MutableMapping, Sequence +from contextlib import AsyncExitStack from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, cast, overload @@ -29,6 +30,7 @@ normalize_messages, normalize_tools, ) +from agent_framework._mcp import _expand_mcp_tools # pyright: ignore[reportPrivateUsage] from agent_framework._telemetry import mark_feature_used from agent_framework.exceptions import AgentException, AgentInvalidRequestException from agent_framework.observability import AgentTelemetryLayer @@ -385,6 +387,9 @@ def __init__( # Separate built-in tools (strings) from custom tools (callables/FunctionTool) self._builtin_tools: list[str] = [] self._custom_tools: list[ToolTypes] = [] + self._mcp_exit_stack = AsyncExitStack() + self._mcp_lock = asyncio.Lock() + self._warned_about_injected_client_tools = False self._normalize_tools(tools) self._default_options = opts @@ -455,6 +460,11 @@ async def stop(self) -> None: with contextlib.suppress(Exception): await self._client.disconnect() + async with self._mcp_lock: + with contextlib.suppress(Exception): + await self._mcp_exit_stack.aclose() + self._mcp_exit_stack = AsyncExitStack() + self._started = False async def _connect_injected_client(self) -> None: @@ -512,10 +522,26 @@ async def _acquire_client(self, session: AgentSession) -> tuple[ClaudeSDKClient, "different conversation. Omit `client=` so each run gets its own isolated " "client, or use a separate ClaudeAgent per session." ) + configured_tool_count = len(self._builtin_tools) + len(self._custom_tools) + if configured_tool_count and not self._warned_about_injected_client_tools: + self._warned_about_injected_client_tools = True + logger.warning( + "Ignoring %d tool(s) configured on this agent: an injected ClaudeSDKClient " + "carries its own options, so the framework cannot register tools on it. " + "Omit `client=` to let the agent build the client with these tools.", + configured_tool_count, + ) await self._connect_injected_client() return self._client, False - opts = self._prepare_client_options(resume_session_id=self._get_chat_conversation_id(session)) + # Runs are deliberately not serialized here, so guard expansion: two concurrent runs + # would otherwise each open their own connection to the same MCP server. + async with self._mcp_lock: + custom_tools = await _expand_mcp_tools(self._custom_tools, self._mcp_exit_stack) + opts = self._prepare_client_options( + resume_session_id=self._get_chat_conversation_id(session), + custom_tools=custom_tools, + ) client = ClaudeSDKClient(options=opts) try: await client.connect() @@ -548,11 +574,16 @@ def _injected_session_matches(self, session: AgentSession) -> bool: return bound.service_session_id == session.service_session_id return bound.session_id == session.session_id - def _prepare_client_options(self, resume_session_id: str | None = None) -> SDKOptions: + def _prepare_client_options( + self, + resume_session_id: str | None = None, + custom_tools: Sequence[ToolTypes] | None = None, + ) -> SDKOptions: """Prepare SDK options for client initialization. Args: resume_session_id: Optional session ID to resume. + custom_tools: Tools to expose, already expanded. Defaults to the configured tools. Returns: SDKOptions instance configured for the client. @@ -587,9 +618,8 @@ def _prepare_client_options(self, resume_session_id: str | None = None) -> SDKOp opts["tools"] = self._builtin_tools # Prepare custom tools (FunctionTool instances) - custom_tools_server, custom_tool_names = ( - self._prepare_tools(self._custom_tools) if self._custom_tools else (None, []) - ) + tools_to_expose = self._custom_tools if custom_tools is None else custom_tools + custom_tools_server, custom_tool_names = self._prepare_tools(tools_to_expose) if tools_to_expose else (None, []) # MCP servers - merge user-provided servers with custom tools server mcp_servers = dict(self._mcp_servers) if self._mcp_servers else {} diff --git a/python/packages/claude/tests/test_claude_agent.py b/python/packages/claude/tests/test_claude_agent.py index 38359f9d152..c334bbe25ee 100644 --- a/python/packages/claude/tests/test_claude_agent.py +++ b/python/packages/claude/tests/test_claude_agent.py @@ -1,10 +1,13 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio +import logging from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest from agent_framework import AgentResponseUpdate, AgentSession, Content, Message, tool +from agent_framework._mcp import MCPTool from agent_framework._settings import load_settings from agent_framework.exceptions import AgentInvalidRequestException @@ -601,7 +604,7 @@ async def test_acquire_client_forwards_resume_id(self) -> None: _, owns_client = await agent._acquire_client(session) # type: ignore[reportPrivateUsage] assert owns_client is True - prepare.assert_called_once_with(resume_session_id="provider-session-1") + assert prepare.call_args.kwargs["resume_session_id"] == "provider-session-1" async def test_acquire_client_reuses_injected_client_for_same_session(self) -> None: """An injected client is reused across runs of one session and never owned by the run.""" @@ -1472,3 +1475,99 @@ async def test_telemetry_uses_correct_provider_name(self, monkeypatch: pytest.Mo call_kwargs = mock_get_span.call_args[1] assert call_kwargs["attributes"]["gen_ai.provider.name"] == "anthropic.claude" + + +class _StubMCPTool(MCPTool): + """MCPTool whose connection is faked, so tool expansion can be tested without a server.""" + + def __init__(self, name: str, functions: list[Any]) -> None: + super().__init__(name=name) + self._stub_functions = functions + self.connect_calls = 0 + + async def connect(self, *, reset: bool = False) -> None: # type: ignore[override] # pyrefly: ignore[bad-override] # ty: ignore[invalid-method-override] + self.connect_calls += 1 + await asyncio.sleep(0) + self.is_connected = True + self._functions = self._stub_functions + + def get_mcp_client(self) -> Any: + raise AssertionError("stub must not open a transport") + + +class TestClaudeAgentMCPTools: + """Tests for exposing MCP server tools to the Claude SDK.""" + + async def test_mcp_tools_are_expanded_when_the_client_is_acquired(self) -> None: + """An MCPTool must reach the SDK as the tools it exposes, not be dropped.""" + + @tool + def remote_search(query: str) -> str: + """Search.""" + return query + + mock_client = MagicMock() + mock_client.connect = AsyncMock() + mcp_tool = _StubMCPTool(name="server", functions=[remote_search]) + with patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client) as sdk_client: + agent = ClaudeAgent(tools=[mcp_tool]) + await agent._acquire_client(agent.create_session()) # type: ignore[reportPrivateUsage] + + options = sdk_client.call_args.kwargs["options"] + assert options.allowed_tools == [f"mcp__{TOOLS_MCP_SERVER_NAME}__remote_search"] + assert TOOLS_MCP_SERVER_NAME in options.mcp_servers + + # The configured server itself is kept, so a later run can reconnect after stop(). + assert agent._custom_tools == [mcp_tool] # type: ignore[reportPrivateUsage] + + await agent.stop() + + def test_unexpanded_mcp_tool_is_not_an_sdk_tool(self) -> None: + """Without expansion the SDK sees nothing, which is what this feature fixes.""" + agent = ClaudeAgent(tools=[_StubMCPTool(name="server", functions=[])]) + + server, tool_names = agent._prepare_tools(agent._custom_tools) # type: ignore[reportPrivateUsage] + + assert server is None + assert tool_names == [] + + async def test_injected_client_warns_that_configured_tools_are_ignored( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """An injected client carries its own options, so configured tools cannot be registered.""" + injected = MagicMock() + injected.connect = AsyncMock() + agent = ClaudeAgent(client=injected, tools=[_StubMCPTool(name="server", functions=[])]) + + with caplog.at_level(logging.WARNING, logger="agent_framework"): + await agent._acquire_client(agent.create_session()) # type: ignore[reportPrivateUsage] + + assert "cannot register tools" in caplog.text + + async def test_injected_client_warning_counts_builtin_tools(self, caplog: pytest.LogCaptureFixture) -> None: + """Built-in tools are configured through the same options, so they are ignored too.""" + injected = MagicMock() + injected.connect = AsyncMock() + agent = ClaudeAgent(client=injected, tools=["Read", "Write"]) + + with caplog.at_level(logging.WARNING, logger="agent_framework"): + await agent._acquire_client(agent.create_session()) # type: ignore[reportPrivateUsage] + + assert "Ignoring 2 tool(s)" in caplog.text + + async def test_concurrent_runs_share_one_mcp_connection(self) -> None: + """Runs are not serialized, so two of them must not each connect to the same server.""" + mcp_tool = _StubMCPTool(name="server", functions=[]) + mock_client = MagicMock() + mock_client.connect = AsyncMock() + + with patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client): + agent = ClaudeAgent(tools=[mcp_tool]) + await asyncio.gather( + agent._acquire_client(agent.create_session()), # type: ignore[reportPrivateUsage] + agent._acquire_client(agent.create_session()), # type: ignore[reportPrivateUsage] + ) + + assert mcp_tool.connect_calls == 1 + + await agent.stop() diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 3d1d5717e37..556293868ff 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -30,7 +30,10 @@ experimental, ) from ._telemetry import FeatureIndex, mark_feature_used -from ._tools import FunctionTool +from ._tools import ( + FunctionTool, + _append_unique_tools, # pyright: ignore[reportPrivateUsage] +) from ._types import ( ChatOptions, Content, @@ -2725,6 +2728,48 @@ async def __aexit__( # region: MCP Plugin Implementations +async def _expand_mcp_tools( # pyright: ignore[reportUnusedFunction] + tools: Sequence[Any], exit_stack: AsyncExitStack +) -> list[Any]: + """Return a tool list with every MCPTool replaced by the remote tools it exposes. + + An MCPTool is a server connection rather than a callable tool, so SDKs that only accept + function tools cannot consume one directly. ``ChatAgent`` already expands them this way; + provider agents that talk to an external SDK need the same step. + + The input is never modified: callers keep their configured MCPTool instances so a later + run can reconnect after the agent is stopped. + + Args: + tools: The tools to expand; entries that are not MCPTool instances are kept as-is. + exit_stack: Stack that owns the lifetime of connections opened here. + + Returns: + A new list with every MCPTool replaced by its exposed functions. + + Raises: + ToolException: If a server uses progressive disclosure, whose loader tools only work + inside a framework function-calling run. + ValueError: If two tools share a name. + """ + duplicate_message = "Tool names must be unique. Consider setting `tool_name_prefix` on the MCPTool." + expanded: list[Any] = [] + for tool in tools: + if not isinstance(tool, MCPTool): + _append_unique_tools(expanded, [tool], duplicate_error_message=duplicate_message) + continue + if tool.use_progressive_disclosure: + raise ToolException( + f"MCP server '{tool.name}' uses progressive disclosure, which exposes loader tools " + "that only work inside a framework function-calling run. Construct it without " + "`use_progressive_disclosure` to use it with this agent." + ) + if not tool.is_connected: + await exit_stack.enter_async_context(tool) + _append_unique_tools(expanded, tool.functions, duplicate_error_message=duplicate_message) + return expanded + + class MCPStdioTool(MCPTool): """MCP tool for connecting to stdio-based MCP servers. diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 669c5d931a6..9814822bb5d 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -7,7 +7,7 @@ import os import sys import warnings -from contextlib import _AsyncGeneratorContextManager # type: ignore +from contextlib import AsyncExitStack, _AsyncGeneratorContextManager # type: ignore from contextvars import ContextVar from datetime import timedelta from typing import Any, cast @@ -28,11 +28,13 @@ MCPStreamableHTTPTool, MCPWebsocketTool, Message, + tool, ) from agent_framework._feature_stage import _WARNED_FEATURES, ExperimentalFeature, ExperimentalWarning from agent_framework._mcp import ( MCPTool, _build_prefixed_mcp_name, + _expand_mcp_tools, _get_input_model_from_mcp_prompt, _normalize_additional_tool_argument_names, _normalize_mcp_name, @@ -8028,3 +8030,109 @@ def provider(_kwargs: dict[str, Any]) -> dict[str, str]: # endregion + + +class _StubMCPTool(MCPTool): + """MCPTool whose connection is faked, so expansion can be tested without a server.""" + + def __init__(self, name: str, functions: list[FunctionTool]) -> None: + super().__init__(name=name) + self._stub_functions = functions + self.connect_calls = 0 + + async def connect(self, *, reset: bool = False) -> None: # type: ignore[override] # pyrefly: ignore[bad-override] # ty: ignore[invalid-method-override] + self.connect_calls += 1 + self.is_connected = True + self._functions = self._stub_functions + + def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: + raise AssertionError("stub must not open a transport") + + +async def test_expand_mcp_tools_replaces_servers_with_their_functions() -> None: + """An MCP server is a connection, so SDK-facing tool lists need its functions instead.""" + + @tool + def remote_search(query: str) -> str: + """Search.""" + return query + + @tool + def local_echo(text: str) -> str: + """Echo.""" + return text + + mcp_tool = _StubMCPTool(name="server", functions=[remote_search]) + + async with AsyncExitStack() as stack: + expanded = await _expand_mcp_tools([local_echo, mcp_tool], stack) + + assert expanded == [local_echo, remote_search] + assert mcp_tool.connect_calls == 1 + + +async def test_expand_mcp_tools_reuses_a_connected_server() -> None: + """A server that is already connected must not be connected a second time.""" + + @tool + def remote_search(query: str) -> str: + """Search.""" + return query + + mcp_tool = _StubMCPTool(name="server", functions=[remote_search]) + + async with AsyncExitStack() as stack: + await mcp_tool.connect() + expanded = await _expand_mcp_tools([mcp_tool], stack) + + assert expanded == [remote_search] + assert mcp_tool.connect_calls == 1 + + +async def test_expand_mcp_tools_leaves_the_caller_list_untouched() -> None: + """Expansion must not consume the configuration, or a later run cannot reconnect.""" + + @tool + def remote_search(query: str) -> str: + """Search.""" + return query + + mcp_tool = _StubMCPTool(name="server", functions=[remote_search]) + configured = [mcp_tool] + + async with AsyncExitStack() as stack: + first = await _expand_mcp_tools(configured, stack) + second = await _expand_mcp_tools(configured, stack) + + assert configured == [mcp_tool] + assert first == second == [remote_search] + + +async def test_expand_mcp_tools_rejects_progressive_disclosure() -> None: + """Loader tools only work inside a framework run, so an external SDK must not receive them.""" + mcp_tool = _StubMCPTool(name="server", functions=[]) + mcp_tool.use_progressive_disclosure = True + + async with AsyncExitStack() as stack: + with pytest.raises(ToolException, match="progressive disclosure"): + await _expand_mcp_tools([mcp_tool], stack) + + +async def test_expand_mcp_tools_rejects_duplicate_tool_names() -> None: + """Two tools with one name would silently shadow each other inside the target SDK.""" + + @tool(name="search") + def local_search(query: str) -> str: + """Search locally.""" + return query + + @tool(name="search") + def remote_search(query: str) -> str: + """Search remotely.""" + return query + + mcp_tool = _StubMCPTool(name="server", functions=[remote_search]) + + async with AsyncExitStack() as stack: + with pytest.raises(ValueError, match="tool_name_prefix"): + await _expand_mcp_tools([local_search, mcp_tool], stack) diff --git a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py index 2f7b963cb75..008c97c9d74 100644 --- a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py +++ b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py @@ -9,7 +9,8 @@ import logging import sys import warnings -from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, MutableMapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Mapping, MutableMapping, Sequence +from contextlib import AsyncExitStack from typing import Any, ClassVar, Generic, Literal, TypedDict, cast, overload from urllib.parse import urlparse @@ -30,9 +31,14 @@ add_usage_details, normalize_messages, ) +from agent_framework._mcp import MCPTool, _expand_mcp_tools # pyright: ignore[reportPrivateUsage] from agent_framework._settings import load_settings from agent_framework._telemetry import mark_feature_used -from agent_framework._tools import FunctionTool, ToolTypes +from agent_framework._tools import ( + FunctionTool, + ToolTypes, + _append_unique_tools, # pyright: ignore[reportPrivateUsage] +) from agent_framework._types import ( AgentRunInputs, _get_data_bytes_as_str, # pyright: ignore[reportPrivateUsage] @@ -641,6 +647,9 @@ def __init__( ) self._tools = normalize_tools(tools) + self._mcp_exit_stack = AsyncExitStack() + self._mcp_lock = asyncio.Lock() + self._runtime_mcp_locks: dict[MCPTool, tuple[asyncio.Lock, int]] = {} self._permission_handler = on_permission_request self._on_pre_tool_use: PreToolUseHandler | None = on_pre_tool_use self._function_approval_handler: FunctionApprovalCallback | None = on_function_approval @@ -708,6 +717,11 @@ async def stop(self) -> None: with contextlib.suppress(Exception): await self._client.stop() + async with self._mcp_lock: + with contextlib.suppress(Exception): + await self._mcp_exit_stack.aclose() + self._mcp_exit_stack = AsyncExitStack() + self._started = False @property @@ -872,7 +886,6 @@ async def _run_impl( existing = list(opts.get("tools") or []) opts["tools"] = existing + list(session_context.tools) - copilot_session = await self._get_or_create_session(session, streaming=False, runtime_options=opts) usage_details: UsageDetails | None = None finish_reason: str | None = None model: str | None = None @@ -906,14 +919,21 @@ def usage_event_handler(event: SessionEvent) -> None: prompt = "\n".join(session_context.instructions) + "\n" + prompt attachments = self._prepare_attachments_for_copilot(context_messages) - unsubscribe = copilot_session.on(usage_event_handler) - try: - mark_feature_used(FeatureIndex.GITHUB_COPILOT) - response_event = await copilot_session.send_and_wait(prompt, attachments=attachments, timeout=timeout) - except Exception as ex: - raise AgentException(f"GitHub Copilot request failed: {ex}") from ex - finally: - unsubscribe() + async with self._expand_runtime_mcp_tools(opts) as runtime_tools: + copilot_session = await self._get_or_create_session( + session, + streaming=False, + runtime_options=opts, + runtime_tools=runtime_tools, + ) + unsubscribe = copilot_session.on(usage_event_handler) + try: + mark_feature_used(FeatureIndex.GITHUB_COPILOT) + response_event = await copilot_session.send_and_wait(prompt, attachments=attachments, timeout=timeout) + except Exception as ex: + raise AgentException(f"GitHub Copilot request failed: {ex}") from ex + finally: + unsubscribe() response_messages: list[Message] = [] response_id: str | None = None @@ -1001,8 +1021,6 @@ async def _stream_updates( existing = list(opts.get("tools") or []) opts["tools"] = existing + list(session_context.tools) - copilot_session = await self._get_or_create_session(session, streaming=True, runtime_options=opts) - if _ctx_holder is not None: _ctx_holder["session_context"] = session_context _ctx_holder["session"] = session @@ -1092,18 +1110,25 @@ def event_handler(event: SessionEvent) -> None: error_msg = error_data.message or "Unknown error" queue.put_nowait(AgentException(f"GitHub Copilot session error: {error_msg}")) - unsubscribe = copilot_session.on(event_handler) + async with self._expand_runtime_mcp_tools(opts) as runtime_tools: + copilot_session = await self._get_or_create_session( + session, + streaming=True, + runtime_options=opts, + runtime_tools=runtime_tools, + ) + unsubscribe = copilot_session.on(event_handler) - try: - mark_feature_used(FeatureIndex.GITHUB_COPILOT) - await copilot_session.send(prompt, attachments=attachments) + try: + mark_feature_used(FeatureIndex.GITHUB_COPILOT) + await copilot_session.send(prompt, attachments=attachments) - while (item := await queue.get()) is not None: - if isinstance(item, Exception): - raise item - yield item - finally: - unsubscribe() + while (item := await queue.get()) is not None: + if isinstance(item, Exception): + raise item + yield item + finally: + unsubscribe() async def _run_before_providers( self, @@ -1396,6 +1421,7 @@ async def _get_or_create_session( agent_session: AgentSession, streaming: bool = False, runtime_options: dict[str, Any] | None = None, + runtime_tools: Sequence[Any] | None = None, ) -> CopilotSession: """Get an existing session or create a new one for the session. @@ -1403,6 +1429,7 @@ async def _get_or_create_session( agent_session: The conversation session. streaming: Whether to enable streaming for the session. runtime_options: Runtime options from run that take precedence. + runtime_tools: Per-run tools after MCP expansion. Returns: A CopilotSession instance. @@ -1420,9 +1447,9 @@ async def _get_or_create_session( raise AgentException( "GitHubCopilotAgent expects a string service_session_id for session resumption." ) - return await self._resume_session(service_session_id, streaming, runtime_options) + return await self._resume_session(service_session_id, streaming, runtime_options, runtime_tools) - session = await self._create_session(streaming, runtime_options) + session = await self._create_session(streaming, runtime_options, runtime_tools) agent_session.service_session_id = session.session_id return session except Exception as ex: @@ -1432,6 +1459,7 @@ def _build_session_kwargs( self, streaming: bool, runtime_options: dict[str, Any] | None, + tools: Sequence[Any] | None = None, ) -> dict[str, Any]: """Assemble keyword arguments for ``create_session`` / ``resume_session``. @@ -1448,6 +1476,8 @@ def _build_session_kwargs( Args: streaming: Whether to enable streaming for the session. runtime_options: Runtime options that take precedence over default_options. + tools: Tools to expose, already expanded. Defaults to merging the agent's tools with + any caller-supplied ones. Returns: The keyword arguments to splat into the SDK session factory. @@ -1459,7 +1489,7 @@ def _build_session_kwargs( # Merge agent-level tools with any caller-supplied tools (from default_options # or per-run options, the latter winning) and convert to SDK tools. - all_tools = list(self._tools or []) + list(kwargs.get("tools") or []) + all_tools = list(tools) if tools is not None else list(self._tools or []) + list(kwargs.get("tools") or []) kwargs["tools"] = self._prepare_tools(all_tools) if all_tools else None kwargs["streaming"] = streaming @@ -1488,27 +1518,100 @@ def _build_session_kwargs( return kwargs + @contextlib.asynccontextmanager + async def _expand_runtime_mcp_tools( + self, runtime_options: dict[str, Any] | None + ) -> AsyncGenerator[list[Any] | None]: + """Expand per-run MCP tools and release their connections after the run.""" + if runtime_options is None or "tools" not in runtime_options: + yield None + return + + configured = list(runtime_options.get("tools") or []) + mcp_tools = sorted({tool for tool in configured if isinstance(tool, MCPTool)}, key=id) + reservations: list[tuple[MCPTool, asyncio.Lock]] = [] + async with self._mcp_lock: + for tool in mcp_tools: + lock, users = self._runtime_mcp_locks.get(tool, (asyncio.Lock(), 0)) + self._runtime_mcp_locks[tool] = (lock, users + 1) + reservations.append((tool, lock)) + + acquired_locks: list[asyncio.Lock] = [] + try: + # Stable ordering prevents deadlocks when concurrent runs share multiple MCPTool instances. + for _, lock in reservations: + await lock.acquire() + acquired_locks.append(lock) + + async with AsyncExitStack() as exit_stack: + yield await _expand_mcp_tools(configured, exit_stack) + finally: + for lock in reversed(acquired_locks): + lock.release() + async with self._mcp_lock: + for tool, lock in reservations: + registered_lock, users = self._runtime_mcp_locks[tool] + if registered_lock is not lock: + raise RuntimeError("Runtime MCP lock changed while it was reserved.") + if users == 1: + del self._runtime_mcp_locks[tool] + else: + self._runtime_mcp_locks[tool] = (lock, users - 1) + + async def _expanded_tools( + self, + runtime_options: dict[str, Any] | None, + runtime_tools: Sequence[Any] | None = None, + ) -> list[Any]: + """Connect any MCP tools so the SDK receives the tools they expose. + + Covers all three places a tool can come from: the agent, ``default_options`` and the + per-run options. The configured lists are left untouched so a later run can reconnect. + """ + has_runtime_tools = runtime_options is not None and "tools" in runtime_options + option_tools = [] if has_runtime_tools else list(self._default_options.get("tools") or []) + configured = list(self._tools or []) + option_tools + async with self._mcp_lock: + if runtime_options is not None and "tools" in runtime_options and runtime_tools is None: + configured.extend(list(runtime_options.get("tools") or [])) + return await _expand_mcp_tools(configured, self._mcp_exit_stack) + expanded = await _expand_mcp_tools(configured, self._mcp_exit_stack) + + if runtime_tools is not None: + _append_unique_tools( + expanded, + list(runtime_tools), + duplicate_error_message=( + "Tool names must be unique. Consider setting `tool_name_prefix` on the MCPTool." + ), + ) + return expanded + async def _create_session( self, streaming: bool, runtime_options: dict[str, Any] | None = None, + runtime_tools: Sequence[Any] | None = None, ) -> CopilotSession: """Create a new Copilot session. Args: streaming: Whether to enable streaming for the session. runtime_options: Runtime options that take precedence over default_options. + runtime_tools: Per-run tools after MCP expansion. """ if not self._client: raise RuntimeError("GitHub Copilot client not initialized. Call start() first.") - return await self._client.create_session(**self._build_session_kwargs(streaming, runtime_options)) + tools = await self._expanded_tools(runtime_options, runtime_tools) + return await self._client.create_session(**self._build_session_kwargs(streaming, runtime_options, tools)) async def _resume_session( self, session_id: str, streaming: bool, runtime_options: dict[str, Any] | None = None, + runtime_tools: Sequence[Any] | None = None, ) -> CopilotSession: """Resume an existing Copilot session by ID. @@ -1516,11 +1619,15 @@ async def _resume_session( session_id: The session ID to resume. streaming: Whether to enable streaming for the session. runtime_options: Runtime options that take precedence over default_options. + runtime_tools: Per-run tools after MCP expansion. """ if not self._client: raise RuntimeError("GitHub Copilot client not initialized. Call start() first.") - return await self._client.resume_session(session_id, **self._build_session_kwargs(streaming, runtime_options)) + tools = await self._expanded_tools(runtime_options, runtime_tools) + return await self._client.resume_session( + session_id, **self._build_session_kwargs(streaming, runtime_options, tools) + ) class GitHubCopilotAgent( # type: ignore[misc] diff --git a/python/packages/github_copilot/tests/test_github_copilot_agent.py b/python/packages/github_copilot/tests/test_github_copilot_agent.py index 6e17db2f733..b044e5760f1 100644 --- a/python/packages/github_copilot/tests/test_github_copilot_agent.py +++ b/python/packages/github_copilot/tests/test_github_copilot_agent.py @@ -2,6 +2,7 @@ # ruff: noqa: E402 +import asyncio import base64 import inspect import json @@ -27,6 +28,7 @@ Message, tool, ) +from agent_framework._mcp import MCPTool from agent_framework.exceptions import AgentException from copilot.session import PermissionHandler, PreToolUseHookInput from copilot.session_events import ( @@ -4110,3 +4112,259 @@ async def test_integration_run_with_shell_permissions_executes_command() -> None if isinstance(session.service_session_id, str) and agent._client: await agent._client.delete_session(session.service_session_id) + + +class _StubMCPTool(MCPTool): + """MCPTool whose connection is faked, so tool expansion can be tested without a server.""" + + def __init__(self, name: str, functions: list[Any]) -> None: + super().__init__(name=name) + self._stub_functions = functions + self.connect_count = 0 + self.close_count = 0 + + async def connect(self, *, reset: bool = False) -> None: # type: ignore[override] # pyrefly: ignore[bad-override] # ty: ignore[invalid-method-override] + self.connect_count += 1 + self.is_connected = True + self._functions = self._stub_functions + + async def close(self) -> None: + self.close_count += 1 + self.is_connected = False + self._functions = [] + + def get_mcp_client(self) -> Any: + raise AssertionError("stub must not open a transport") + + +class TestGitHubCopilotAgentMCPTools: + """Tests for exposing MCP server tools to the Copilot SDK.""" + + async def test_agent_level_mcp_tool_reaches_the_sdk( + self, + mock_client: MagicMock, + mock_session: MagicMock, + ) -> None: + """An MCPTool configured on the agent must arrive as the tools it exposes.""" + + @tool + def remote_search(query: str) -> str: + """Search.""" + return query + + mcp_tool = _StubMCPTool(name="server", functions=[remote_search]) + agent = GitHubCopilotAgent( + client=mock_client, + tools=[mcp_tool], + ) + await agent.start() + + await agent._get_or_create_session(AgentSession()) # type: ignore + + sdk_tools = mock_client.create_session.call_args.kwargs["tools"] + assert [sdk_tool.name for sdk_tool in sdk_tools] == ["remote_search"] + assert mcp_tool.is_connected + assert mcp_tool.close_count == 0 + + await agent.stop() + assert not mcp_tool.is_connected + assert mcp_tool.close_count == 1 + + async def test_runtime_mcp_tool_reaches_the_sdk( + self, + mock_client: MagicMock, + mock_session: MagicMock, + assistant_message_event: SessionEvent, + ) -> None: + """An MCPTool passed per run must be expanded too, not only agent-level ones.""" + + @tool + def remote_search(query: str) -> str: + """Search.""" + return query + + mock_session.send_and_wait.return_value = assistant_message_event + mcp_tool = _StubMCPTool(name="server", functions=[remote_search]) + agent = GitHubCopilotAgent(client=mock_client) + + await agent.run("search", options=cast(Any, {"tools": [mcp_tool]})) + + sdk_tools = mock_client.create_session.call_args.kwargs["tools"] + assert [sdk_tool.name for sdk_tool in sdk_tools] == ["remote_search"] + assert not mcp_tool.is_connected + assert mcp_tool.close_count == 1 + + await agent.stop() + + async def test_runtime_mcp_tool_closes_after_failed_run( + self, + mock_client: MagicMock, + mock_session: MagicMock, + ) -> None: + """A failed request must release MCP connections supplied for that run.""" + + @tool + def remote_search(query: str) -> str: + """Search.""" + return query + + mock_session.send_and_wait.side_effect = Exception("request failed") + mcp_tool = _StubMCPTool(name="server", functions=[remote_search]) + agent = GitHubCopilotAgent(client=mock_client) + + with pytest.raises(AgentException, match="request failed"): + await agent.run("search", options=cast(Any, {"tools": [mcp_tool]})) + + assert not mcp_tool.is_connected + assert mcp_tool.close_count == 1 + + await agent.stop() + + async def test_runtime_mcp_tool_closes_after_streaming_run( + self, + mock_client: MagicMock, + mock_session: MagicMock, + session_idle_event: SessionEvent, + ) -> None: + """A completed stream must release MCP connections supplied for that run.""" + + @tool + def remote_search(query: str) -> str: + """Search.""" + return query + + mock_session.on = lambda handler: (handler(session_idle_event), lambda: None)[1] + mcp_tool = _StubMCPTool(name="server", functions=[remote_search]) + agent = GitHubCopilotAgent(client=mock_client) + + async for _ in agent.run("search", stream=True, options=cast(Any, {"tools": [mcp_tool]})): + pass + + assert not mcp_tool.is_connected + assert mcp_tool.close_count == 1 + + await agent.stop() + + async def test_shared_runtime_mcp_tool_is_not_closed_during_concurrent_run( + self, + mock_client: MagicMock, + mock_session: MagicMock, + assistant_message_event: SessionEvent, + ) -> None: + """Concurrent runs sharing one MCPTool must serialize its connection ownership.""" + + @tool + def remote_search(query: str) -> str: + """Search.""" + return query + + first_request_started = asyncio.Event() + finish_first_request = asyncio.Event() + request_count = 0 + + async def send_and_wait(*args: Any, **kwargs: Any) -> SessionEvent: + nonlocal request_count + request_count += 1 + if request_count == 1: + first_request_started.set() + await finish_first_request.wait() + return assistant_message_event + + mock_session.send_and_wait.side_effect = send_and_wait + mcp_tool = _StubMCPTool(name="server", functions=[remote_search]) + agent = GitHubCopilotAgent(client=mock_client) + options = cast(Any, {"tools": [mcp_tool]}) + + first_run = asyncio.ensure_future(agent.run("first", options=options)) + await first_request_started.wait() + second_run = asyncio.ensure_future(agent.run("second", options=options)) + for _ in range(10): + if agent._runtime_mcp_locks[mcp_tool][1] == 2: # type: ignore[reportPrivateUsage] + break + await asyncio.sleep(0) + + assert agent._runtime_mcp_locks[mcp_tool][1] == 2 # type: ignore[reportPrivateUsage] + assert mcp_tool.connect_count == 1 + assert mcp_tool.close_count == 0 + + finish_first_request.set() + await asyncio.gather(first_run, second_run) + + assert mcp_tool.connect_count == 2 + assert mcp_tool.close_count == 2 + assert mcp_tool not in agent._runtime_mcp_locks # type: ignore[reportPrivateUsage] + + await agent.stop() + + async def test_default_mcp_tool_reuses_connection_until_stop( + self, + mock_client: MagicMock, + mock_session: MagicMock, + assistant_message_event: SessionEvent, + ) -> None: + """An MCPTool in default options remains connected across runs.""" + + @tool + def remote_search(query: str) -> str: + """Search.""" + return query + + mock_session.send_and_wait.return_value = assistant_message_event + mcp_tool = _StubMCPTool(name="server", functions=[remote_search]) + agent = GitHubCopilotAgent(client=mock_client, default_options=cast(Any, {"tools": [mcp_tool]})) + session = AgentSession() + + await agent.run("first", session=session) + await agent.run("second", session=session) + + assert mcp_tool.connect_count == 1 + assert mcp_tool.close_count == 0 + + await agent.stop() + assert mcp_tool.close_count == 1 + + async def test_runtime_empty_tools_clears_default_options_tools( + self, + mock_client: MagicMock, + mock_session: MagicMock, + ) -> None: + """A per-run ``tools=[]`` must override default_options, not be read as 'not provided'.""" + + @tool(name="default_tool") + def default_tool(query: str) -> str: + """Default.""" + return query + + agent = GitHubCopilotAgent(client=mock_client, default_options=cast(Any, {"tools": [default_tool]})) + await agent.start() + + await agent._get_or_create_session(AgentSession(), runtime_options={"tools": []}) # type: ignore + + assert mock_client.create_session.call_args.kwargs["tools"] is None + + await agent.stop() + + async def test_resumed_session_also_receives_mcp_tools( + self, + mock_client: MagicMock, + mock_session: MagicMock, + ) -> None: + """Resuming a session must expand MCP tools too, not only creating one.""" + + @tool + def remote_search(query: str) -> str: + """Search.""" + return query + + agent = GitHubCopilotAgent( + client=mock_client, + tools=[_StubMCPTool(name="server", functions=[remote_search])], + ) + await agent.start() + + await agent._resume_session("session-1", streaming=False) # type: ignore + + sdk_tools = mock_client.resume_session.call_args.kwargs["tools"] + assert [sdk_tool.name for sdk_tool in sdk_tools] == ["remote_search"] + + await agent.stop()