mirror of
https://github.com/home-assistant/core.git
synced 2025-07-17 02:07:09 +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.async_on_unload(unsub)
|
||||||
|
|
||||||
entry.runtime_data = coordinator
|
entry.runtime_data = coordinator
|
||||||
entry.async_on_unload(coordinator.close)
|
|
||||||
|
|
||||||
return True
|
return True
|
||||||
|
|
||||||
|
@ -40,6 +40,7 @@ async def mcp_client(url: str) -> AsyncGenerator[ClientSession]:
|
|||||||
await session.initialize()
|
await session.initialize()
|
||||||
yield session
|
yield session
|
||||||
except ExceptionGroup as err:
|
except ExceptionGroup as err:
|
||||||
|
_LOGGER.debug("Error creating MCP client: %s", err)
|
||||||
raise err.exceptions[0] from err
|
raise err.exceptions[0] from err
|
||||||
|
|
||||||
|
|
||||||
@ -51,13 +52,13 @@ class ModelContextProtocolTool(llm.Tool):
|
|||||||
name: str,
|
name: str,
|
||||||
description: str | None,
|
description: str | None,
|
||||||
parameters: vol.Schema,
|
parameters: vol.Schema,
|
||||||
session: ClientSession,
|
server_url: str,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Initialize the tool."""
|
"""Initialize the tool."""
|
||||||
self.name = name
|
self.name = name
|
||||||
self.description = description
|
self.description = description
|
||||||
self.parameters = parameters
|
self.parameters = parameters
|
||||||
self.session = session
|
self.server_url = server_url
|
||||||
|
|
||||||
async def async_call(
|
async def async_call(
|
||||||
self,
|
self,
|
||||||
@ -67,10 +68,16 @@ class ModelContextProtocolTool(llm.Tool):
|
|||||||
) -> JsonObjectType:
|
) -> JsonObjectType:
|
||||||
"""Call the tool."""
|
"""Call the tool."""
|
||||||
try:
|
try:
|
||||||
result = await self.session.call_tool(
|
async with asyncio.timeout(TIMEOUT):
|
||||||
tool_input.tool_name, tool_input.tool_args
|
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:
|
except httpx.HTTPStatusError as error:
|
||||||
|
_LOGGER.debug("Error when calling tool: %s", error)
|
||||||
raise HomeAssistantError(f"Error when calling tool: {error}") from error
|
raise HomeAssistantError(f"Error when calling tool: {error}") from error
|
||||||
return result.model_dump(exclude_unset=True, exclude_none=True)
|
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."""
|
"""Define an object to hold MCP data."""
|
||||||
|
|
||||||
config_entry: ConfigEntry
|
config_entry: ConfigEntry
|
||||||
_session: ClientSession | None = None
|
|
||||||
_setup_error: Exception | None = None
|
|
||||||
|
|
||||||
def __init__(self, hass: HomeAssistant, config_entry: ConfigEntry) -> None:
|
def __init__(self, hass: HomeAssistant, config_entry: ConfigEntry) -> None:
|
||||||
"""Initialize ModelContextProtocolCoordinator."""
|
"""Initialize ModelContextProtocolCoordinator."""
|
||||||
@ -91,52 +96,6 @@ class ModelContextProtocolCoordinator(DataUpdateCoordinator[list[llm.Tool]]):
|
|||||||
config_entry=config_entry,
|
config_entry=config_entry,
|
||||||
update_interval=UPDATE_INTERVAL,
|
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]:
|
async def _async_update_data(self) -> list[llm.Tool]:
|
||||||
"""Fetch data from API endpoint.
|
"""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
|
This is the place to pre-process the data to lookup tables
|
||||||
so entities can quickly look up their data.
|
so entities can quickly look up their data.
|
||||||
"""
|
"""
|
||||||
if self._session is None:
|
|
||||||
raise UpdateFailed("No session available")
|
|
||||||
try:
|
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:
|
except httpx.HTTPError as err:
|
||||||
|
_LOGGER.debug("Error communicating with API: %s", err)
|
||||||
raise UpdateFailed(f"Error communicating with API: {err}") from err
|
raise UpdateFailed(f"Error communicating with API: {err}") from err
|
||||||
|
|
||||||
_LOGGER.debug("Received tools: %s", result.tools)
|
_LOGGER.debug("Received tools: %s", result.tools)
|
||||||
@ -165,7 +128,7 @@ class ModelContextProtocolCoordinator(DataUpdateCoordinator[list[llm.Tool]]):
|
|||||||
tool.name,
|
tool.name,
|
||||||
tool.description,
|
tool.description,
|
||||||
parameters,
|
parameters,
|
||||||
self._session,
|
self.config_entry.data[CONF_URL],
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
return tools
|
return tools
|
||||||
|
Loading…
x
Reference in New Issue
Block a user