Update MCP to reconnect to the server on demand (#141215)

* Reconnect to the MCP client on deman

* Remove debug log

* Update log messages
This commit is contained in:
Allen Porter 2025-03-23 09:11:40 -07:00 committed by GitHub
parent c451518959
commit 28ef0a33ad
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 20 additions and 58 deletions

View File

@ -39,7 +39,6 @@ async def async_setup_entry(
entry.async_on_unload(unsub)
entry.runtime_data = coordinator
entry.async_on_unload(coordinator.close)
return True

View File

@ -40,6 +40,7 @@ async def mcp_client(url: str) -> AsyncGenerator[ClientSession]:
await session.initialize()
yield session
except ExceptionGroup as err:
_LOGGER.debug("Error creating MCP client: %s", err)
raise err.exceptions[0] from err
@ -51,13 +52,13 @@ class ModelContextProtocolTool(llm.Tool):
name: str,
description: str | None,
parameters: vol.Schema,
session: ClientSession,
server_url: str,
) -> None:
"""Initialize the tool."""
self.name = name
self.description = description
self.parameters = parameters
self.session = session
self.server_url = server_url
async def async_call(
self,
@ -67,10 +68,16 @@ class ModelContextProtocolTool(llm.Tool):
) -> JsonObjectType:
"""Call the tool."""
try:
result = await self.session.call_tool(
tool_input.tool_name, tool_input.tool_args
)
async with asyncio.timeout(TIMEOUT):
async with mcp_client(self.server_url) as session:
result = await session.call_tool(
tool_input.tool_name, tool_input.tool_args
)
except TimeoutError as error:
_LOGGER.debug("Timeout when calling tool: %s", error)
raise HomeAssistantError(f"Timeout when calling tool: {error}") from error
except httpx.HTTPStatusError as error:
_LOGGER.debug("Error when calling tool: %s", error)
raise HomeAssistantError(f"Error when calling tool: {error}") from error
return result.model_dump(exclude_unset=True, exclude_none=True)
@ -79,8 +86,6 @@ class ModelContextProtocolCoordinator(DataUpdateCoordinator[list[llm.Tool]]):
"""Define an object to hold MCP data."""
config_entry: ConfigEntry
_session: ClientSession | None = None
_setup_error: Exception | None = None
def __init__(self, hass: HomeAssistant, config_entry: ConfigEntry) -> None:
"""Initialize ModelContextProtocolCoordinator."""
@ -91,52 +96,6 @@ class ModelContextProtocolCoordinator(DataUpdateCoordinator[list[llm.Tool]]):
config_entry=config_entry,
update_interval=UPDATE_INTERVAL,
)
self._stop = asyncio.Event()
async def _async_setup(self) -> None:
"""Set up the client connection."""
connected = asyncio.Event()
stop = asyncio.Event()
self.config_entry.async_create_background_task(
self.hass, self._connect(connected, stop), "mcp-client"
)
try:
async with asyncio.timeout(TIMEOUT):
await connected.wait()
self._stop = stop
finally:
if self._setup_error is not None:
raise self._setup_error
async def _connect(self, connected: asyncio.Event, stop: asyncio.Event) -> None:
"""Create a server-sent event MCP client."""
url = self.config_entry.data[CONF_URL]
try:
async with (
sse_client(url=url) as streams,
ClientSession(*streams) as session,
):
await session.initialize()
self._session = session
connected.set()
await stop.wait()
except httpx.HTTPStatusError as err:
self._setup_error = err
_LOGGER.debug("Error connecting to MCP server: %s", err)
raise UpdateFailed(f"Error connecting to MCP server: {err}") from err
except ExceptionGroup as err:
self._setup_error = err.exceptions[0]
_LOGGER.debug("Error connecting to MCP server: %s", err)
raise UpdateFailed(
"Error connecting to MCP server: {err.exceptions[0]}"
) from err.exceptions[0]
finally:
self._session = None
async def close(self) -> None:
"""Close the client connection."""
if self._stop is not None:
self._stop.set()
async def _async_update_data(self) -> list[llm.Tool]:
"""Fetch data from API endpoint.
@ -144,11 +103,15 @@ class ModelContextProtocolCoordinator(DataUpdateCoordinator[list[llm.Tool]]):
This is the place to pre-process the data to lookup tables
so entities can quickly look up their data.
"""
if self._session is None:
raise UpdateFailed("No session available")
try:
result = await self._session.list_tools()
async with asyncio.timeout(TIMEOUT):
async with mcp_client(self.config_entry.data[CONF_URL]) as session:
result = await session.list_tools()
except TimeoutError as error:
_LOGGER.debug("Timeout when listing tools: %s", error)
raise UpdateFailed(f"Timeout when listing tools: {error}") from error
except httpx.HTTPError as err:
_LOGGER.debug("Error communicating with API: %s", err)
raise UpdateFailed(f"Error communicating with API: {err}") from err
_LOGGER.debug("Received tools: %s", result.tools)
@ -165,7 +128,7 @@ class ModelContextProtocolCoordinator(DataUpdateCoordinator[list[llm.Tool]]):
tool.name,
tool.description,
parameters,
self._session,
self.config_entry.data[CONF_URL],
)
)
return tools