mirror of
https://github.com/home-assistant/core.git
synced 2025-04-24 01:08:12 +00:00
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:
parent
c451518959
commit
28ef0a33ad
@ -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
|
||||
|
||||
|
@ -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
|
||||
|
Loading…
x
Reference in New Issue
Block a user