fix(embedded): support Probatio schema conversion (#2286)

* fix(embedded): support Probatio schema conversion

* test(embedded): cover nested converter import failures

Author: kingpanther13

custom_components/ha_mcp_tools/llm_api.py

@@ -49,17 +49,17 @@
 import asyncio
 import importlib
 import logging
-from collections.abc import AsyncIterator, Iterable
+from collections.abc import AsyncIterator, Callable, Iterable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
+from functools import cache
 from typing import TYPE_CHECKING, Any, cast
 
 import voluptuous as vol
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import llm
 from homeassistant.helpers.httpx_client import get_async_client
-from voluptuous_openapi import convert_to_voluptuous
 
 from .const import (
     DATA_LLM_API_UNSUB,
@@ -128,6 +128,25 @@
 _CALL_TOOL_NAME = "ha_call_tool"
 _SEARCH_RESULT_LIMIT = 8
 
+
+@cache
+def _schema_converter() -> Callable[[Any], Any]:
+    """Resolve the Core-provided schema converter once, off the event loop."""
+    try:
+        legacy = importlib.import_module("voluptuous_openapi")
+    except ModuleNotFoundError as err:
+        if err.name != "voluptuous_openapi":
+            raise
+        probatio = importlib.import_module("probatio")
+        return cast(Callable[[Any], Any], probatio.from_openapi)
+    return cast(Callable[[Any], Any], legacy.convert_to_voluptuous)
+
+
+def convert_to_voluptuous(schema: Any) -> vol.Schema:
+    """Convert an OpenAPI schema on stable and Probatio-based HA Core."""
+    return cast(vol.Schema, _schema_converter()(schema))
+
+
 # Used when the server's initialize result carries no instructions (it always
 # should — ha-mcp ships server-level instructions — but never render an empty
 # prompt if a build does not).
@@ -199,24 +218,25 @@ def _is_transport_failure(err: BaseException) -> bool:
 
 
 def _import_mcp_sdk() -> None:
-    """Import the mcp client SDK modules (blocking; run on the executor).
+    """Import lazy LLM dependencies (blocking; run on the executor).
 
-    Raises ImportError when the SDK is not importable — the caller decides
-    whether that skips registration (SDK missing entirely) or surfaces as a
-    conversation error.
+    Raises ImportError when the MCP SDK or Core's schema converter is not
+    importable — the caller decides whether that skips registration or
+    surfaces as a conversation error.
     """
     importlib.import_module("mcp.client.session")
     importlib.import_module("mcp.client.streamable_http")
+    _schema_converter()
 
 
 async def async_probe_mcp_sdk(hass: HomeAssistant) -> bool:
-    """Return True when the mcp client SDK imports (first import off-loop)."""
+    """Return True when lazy LLM dependencies import (first import off-loop)."""
     try:
         await hass.async_add_executor_job(_import_mcp_sdk)
     except ImportError as err:
         _LOGGER.warning(
-            "The installed server package provides no importable 'mcp' client "
-            "SDK (%s); the conversation-agent LLM API will not be available",
+            "A required LLM dependency is not importable (%s); the "
+            "conversation-agent LLM API will not be available",
             err,
         )
         return False
@@ -538,9 +558,7 @@ async def async_get_api_instance(
     def _convert_parameters(self, tool: Any) -> vol.Schema | None:
         """Convert one tool's JSON schema, or None (logged) when it fails."""
         try:
-            # cast: voluptuous_openapi is an untyped (ignored) import, so the
-            # call returns Any; its documented return type is vol.Schema.
-            return cast(vol.Schema, convert_to_voluptuous(tool.inputSchema))
+            return convert_to_voluptuous(tool.inputSchema)
         except Exception:
             # One unconvertible schema must not take down the whole
             # toolset for the conversation — skip that tool, loudly.

tests/src/unit/_embedded_stubs.py

@@ -613,13 +613,18 @@ def setmod(name: str, **attrs: Any) -> ModuleType:
         async_register_api=_llm_async_register_api,
     )
     _pin_llm_on_helpers()
-    # voluptuous_openapi (an HA-core runtime dependency, not installed in this
-    # test environment). Pass-through conversion: the unit tests assert wiring,
-    # not the real JSON-schema -> voluptuous translation.
+    # HA Core's schema conversion dependency differs across supported releases:
+    # stable uses voluptuous_openapi while 2026.9+ provides Probatio. Pass-through
+    # conversion keeps these unit tests focused on wiring; compatibility selection
+    # itself is covered in test_llm_api.
     setmod(
         "voluptuous_openapi",
         convert_to_voluptuous=lambda schema: {"_converted": schema},
     )
+    setmod(
+        "probatio",
+        from_openapi=lambda schema: {"_converted": schema},
+    )
     # aiohttp_client + event helpers for the periodic auto-update check
     # (embedded_setup fetches PyPI; embedded_entry registers the interval).
     setmod(

tests/src/unit/test_llm_api.py

@@ -14,6 +14,7 @@
 
 import asyncio
 import logging
+import threading
 from contextlib import asynccontextmanager
 from types import SimpleNamespace
 from typing import Any
@@ -260,7 +261,95 @@ def _boom() -> None:
 
         assert fake_llm_apis(hass) == {}
         assert DATA_LLM_API_UNSUB not in hass.data.get(DOMAIN, {})
-        assert "no importable 'mcp' client SDK" in caplog.text
+        assert "required LLM dependency is not importable" in caplog.text
+
+
+class TestSchemaConversionCompatibility:
+    @pytest.fixture(autouse=True)
+    def _clear_schema_converter_cache(self):
+        llm_api._schema_converter.cache_clear()
+        yield
+        llm_api._schema_converter.cache_clear()
+
+    def test_prefers_stable_core_converter_when_available(self, monkeypatch):
+        schema = {"type": "object"}
+        legacy = SimpleNamespace(
+            convert_to_voluptuous=lambda value: {"voluptuous_openapi": value}
+        )
+
+        def _import_module(name):
+            assert name == "voluptuous_openapi"
+            return legacy
+
+        monkeypatch.setattr(llm_api.importlib, "import_module", _import_module)
+
+        assert llm_api.convert_to_voluptuous(schema) == {"voluptuous_openapi": schema}
+
+    def test_falls_back_to_probatio_on_newer_core(self, monkeypatch):
+        schema = {"type": "object"}
+        probatio = SimpleNamespace(from_openapi=lambda value: {"probatio": value})
+
+        def _import_module(name):
+            if name == "voluptuous_openapi":
+                raise ModuleNotFoundError(
+                    "No module named 'voluptuous_openapi'",
+                    name="voluptuous_openapi",
+                )
+            assert name == "probatio"
+            return probatio
+
+        monkeypatch.setattr(llm_api.importlib, "import_module", _import_module)
+
+        assert llm_api.convert_to_voluptuous(schema) == {"probatio": schema}
+
+    def test_reraises_nested_module_not_found(self, monkeypatch):
+        def _import_module(name):
+            assert name == "voluptuous_openapi"
+            raise ModuleNotFoundError(
+                "No module named 'legacy_dependency'",
+                name="legacy_dependency",
+            )
+
+        monkeypatch.setattr(llm_api.importlib, "import_module", _import_module)
+
+        with pytest.raises(ModuleNotFoundError, match="legacy_dependency"):
+            llm_api.convert_to_voluptuous({"type": "object"})
+
+    async def test_converter_import_is_warmed_once_off_event_loop(self, monkeypatch):
+        main_thread = threading.get_ident()
+        imports: list[tuple[str, int]] = []
+        legacy = SimpleNamespace(
+            convert_to_voluptuous=lambda value: {"voluptuous_openapi": value}
+        )
+
+        def _import_module(name):
+            imports.append((name, threading.get_ident()))
+            if name.startswith("mcp."):
+                return SimpleNamespace()
+            if name == "voluptuous_openapi":
+                return legacy
+            if name == "probatio":
+                raise ModuleNotFoundError("No module named 'probatio'", name="probatio")
+            raise AssertionError(name)
+
+        async def _executor(func, *args):
+            return await asyncio.to_thread(func, *args)
+
+        hass = _make_hass()
+        hass.async_add_executor_job = AsyncMock(side_effect=_executor)
+        monkeypatch.setattr(llm_api.importlib, "import_module", _import_module)
+
+        assert await llm_api.async_probe_mcp_sdk(hass)
+        schema = {"type": "object"}
+        assert llm_api.convert_to_voluptuous(schema) == {"voluptuous_openapi": schema}
+        assert llm_api.convert_to_voluptuous(schema) == {"voluptuous_openapi": schema}
+
+        assert [name for name, _ in imports] == [
+            "mcp.client.session",
+            "mcp.client.streamable_http",
+            "voluptuous_openapi",
+        ]
+        assert all(thread_id != main_thread for _, thread_id in imports)
 
 
 class TestFullModeInstance: