diff --git a/homeassistant/components/mcp/config_flow.py b/homeassistant/components/mcp/config_flow.py index d40307b8017c..c518fa00ef68 100644 --- a/homeassistant/components/mcp/config_flow.py +++ b/homeassistant/components/mcp/config_flow.py @@ -117,8 +117,15 @@ async def validate_input( except vol.Invalid as error: raise InvalidUrl from error try: - async with mcp_client(hass, url, token_manager=token_manager) as session: - response = await session.initialize() + async with mcp_client(hass, url, token_manager=token_manager) as ( + _session, + response, + ): + if not response.capabilities.tools: + raise MissingCapabilities( + f"MCP Server {url} does not support 'Tools' capability" + ) + return {"title": response.serverInfo.name} except httpx.TimeoutException as error: _LOGGER.info("Timeout connecting to MCP server: %s", error) raise TimeoutConnectError from error @@ -132,13 +139,6 @@ async def validate_input( _LOGGER.info("Cannot connect to MCP server: %s", error) raise CannotConnect from error - if not response.capabilities.tools: - raise MissingCapabilities( - f"MCP Server {url} does not support 'Tools' capability" - ) - - return {"title": response.serverInfo.name} - class ModelContextProtocolConfigFlow(AbstractOAuth2FlowHandler, domain=DOMAIN): """Handle a config flow for Model Context Protocol.""" diff --git a/homeassistant/components/mcp/coordinator.py b/homeassistant/components/mcp/coordinator.py index 52c4ed275934..b11c02639aa9 100644 --- a/homeassistant/components/mcp/coordinator.py +++ b/homeassistant/components/mcp/coordinator.py @@ -12,6 +12,7 @@ from mcp import McpError from mcp.client.session import ClientSession from mcp.client.sse import sse_client from mcp.client.streamable_http import streamable_http_client +from mcp.types import InitializeResult from probatio import from_openapi import voluptuous as vol @@ -65,7 +66,7 @@ async def mcp_client( hass: HomeAssistant, url: str, token_manager: TokenManager | None = None, -) -> AsyncGenerator[ClientSession]: +) -> AsyncGenerator[tuple[ClientSession, InitializeResult]]: """Create an MCP client. This is an asynccontext manager that exists to wrap other async context managers @@ -84,8 +85,8 @@ async def mcp_client( ) as (read_stream, write_stream, _), ClientSession(read_stream, write_stream) as session, ): - await session.initialize() - yield session + result = await session.initialize() + yield session, result except ExceptionGroup as streamable_err: main_error = streamable_err.exceptions[0] # Method not Allowed likely means this is not a streamable HTTP server, @@ -109,8 +110,8 @@ async def mcp_client( ) as streams, ClientSession(*streams) as session, ): - await session.initialize() - yield session + result = await session.initialize() + yield session, result except ExceptionGroup as sse_err: _LOGGER.debug("Error creating SSE MCP client: %s", sse_err) raise sse_err.exceptions[0] from sse_err @@ -149,9 +150,10 @@ class ModelContextProtocolTool(llm.Tool): """Call the tool.""" try: async with asyncio.timeout(TIMEOUT): - async with mcp_client( - hass, self.server_url, self.token_manager - ) as session: + async with mcp_client(hass, self.server_url, self.token_manager) as ( + session, + _, + ): result = await session.call_tool( tool_input.tool_name, tool_input.tool_args ) @@ -219,7 +221,7 @@ class ModelContextProtocolCoordinator(DataUpdateCoordinator[list[llm.Tool]]): async with asyncio.timeout(TIMEOUT): async with mcp_client( self.hass, self.config_entry.data[CONF_URL], self.token_manager - ) as session: + ) as (session, _): result = await session.list_tools() except TimeoutError as error: _LOGGER.debug("Timeout when listing tools: %s", error) diff --git a/tests/components/mcp/test_config_flow.py b/tests/components/mcp/test_config_flow.py index ad8aa2baf80c..59e794fd5e2a 100644 --- a/tests/components/mcp/test_config_flow.py +++ b/tests/components/mcp/test_config_flow.py @@ -119,6 +119,7 @@ async def test_form( assert result["result"].unique_id is None assert len(mock_setup_entry.mock_calls) == 1 + mock_mcp_client.return_value.initialize.assert_called_once() @pytest.mark.parametrize(