feat: add FastMCP 3.1 Code Mode transform

.env.example

@@ -25,4 +25,11 @@ LOG_LEVEL=INFO
 
 # Optional: MCP Server Configuration
 MCP_SERVER_NAME=ha-mcp
-# MCP_SERVER_VERSION defaults to the package version (e.g. 6.7.2)
+# MCP_SERVER_VERSION defaults to the package version (e.g. 6.7.2)
+
+# Code Mode (experimental) — replaces all tools with meta-tools
+# (search, get_schema, execute) to reduce context bloat.
+# LLM discovers tools on demand and chains calls in a Python sandbox.
+ENABLE_CODE_MODE=false
+# Also expose ListTools for full catalog listing (useful for small catalogs)
+ENABLE_CODE_MODE_LIST_TOOLS=false

homeassistant-addon-dev/Dockerfile

@@ -14,12 +14,12 @@ ENV UV_COMPILE_BYTECODE=1 UV_LINK_MODE=copy
 # Install dependencies first (cached separately from source changes)
 COPY pyproject.toml uv.lock ./
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv sync --locked --no-install-project --no-dev
+    uv sync --no-install-project --no-dev
 
 # Copy source and install the project itself
 COPY src/ ./src/
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv sync --locked --no-dev --no-editable
+    uv sync --no-dev --no-editable
 
 # --- Runtime stage: clean image without uv ---
 FROM python:3.13-slim@sha256:3de9a8d7aedbb7984dc18f2dff178a7850f16c1ae7c34ba9d7ecc23d0755e35f

homeassistant-addon-dev/config.yaml

@@ -19,10 +19,14 @@ options:
   backup_hint: "normal"
   enable_skills: true
   enable_skills_as_tools: false
+  enable_code_mode: false
+  enable_code_mode_list_tools: false
 schema:
   backup_hint: list(strong|normal|weak|auto)
   secret_path: str?
   enable_skills: bool?
   enable_skills_as_tools: bool?
+  enable_code_mode: bool?
+  enable_code_mode_list_tools: bool?
 ports:
   9583/tcp: 9583

homeassistant-addon-dev/translations/en.yaml

@@ -18,3 +18,11 @@ configuration:
     description: >-
       Expose skills via list_resources/read_resource tools for MCP clients
       that don't support resources natively. Adds 3 extra tools.
+  enable_code_mode:
+    name: Code Mode (experimental)
+    description: >-
+      Replaces all tools with meta-tools (search, get_schema, execute).
+  enable_code_mode_list_tools:
+    name: Code Mode ListTools
+    description: >-
+      Also expose ListTools in Code Mode for full catalog discovery.

homeassistant-addon/start.py

@@ -90,6 +90,8 @@ def main() -> int:
     custom_secret_path = ""  # default
     enable_skills = True  # default
     enable_skills_as_tools = False  # default
+    enable_code_mode = False  # default
+    enable_code_mode_list_tools = False  # default
 
     if config_file.exists():
         try:
@@ -101,6 +103,10 @@ def main() -> int:
             enable_skills = raw_skills if isinstance(raw_skills, bool) else True
             raw_skills_as_tools = config.get("enable_skills_as_tools", False)
             enable_skills_as_tools = raw_skills_as_tools if isinstance(raw_skills_as_tools, bool) else False
+            raw_code_mode = config.get("enable_code_mode", False)
+            enable_code_mode = raw_code_mode if isinstance(raw_code_mode, bool) else False
+            raw_code_mode_list = config.get("enable_code_mode_list_tools", False)
+            enable_code_mode_list_tools = raw_code_mode_list if isinstance(raw_code_mode_list, bool) else False
         except Exception as e:
             log_error(f"Failed to read config: {e}, using defaults")
 
@@ -114,6 +120,8 @@ def main() -> int:
     os.environ["BACKUP_HINT"] = backup_hint
     os.environ["ENABLE_SKILLS"] = str(enable_skills).lower()
     os.environ["ENABLE_SKILLS_AS_TOOLS"] = str(enable_skills_as_tools).lower()
+    os.environ["ENABLE_CODE_MODE"] = str(enable_code_mode).lower()
+    os.environ["ENABLE_CODE_MODE_LIST_TOOLS"] = str(enable_code_mode_list_tools).lower()
 
     # Validate Supervisor token
     supervisor_token = os.environ.get("SUPERVISOR_TOKEN")

pyproject.toml

@@ -24,7 +24,7 @@ classifiers = [
 ]
 
 dependencies = [
-    "fastmcp==3.1.0",
+    "fastmcp[code-mode]==3.1.0",
     "httpx[socks]==0.28.1",
     'jq==1.11.0; sys_platform != "win32"',
     "pydantic==2.12.5",

src/ha_mcp/config.py

@@ -85,6 +85,15 @@ class Settings(BaseSettings):
     # Disable when using clients with programmatic tool use (future).
     enable_dashboard_partial_tools: bool = Field(True, alias="ENABLE_DASHBOARD_PARTIAL_TOOLS")
 
+    # Code Mode configuration (FastMCP 3.1 experimental)
+    # Replaces all tools with meta-tools (search, get_schema, execute) to reduce
+    # context bloat and allow the LLM to chain multiple tool calls in one execution.
+    enable_code_mode: bool = Field(False, alias="ENABLE_CODE_MODE")
+
+    # When code mode is enabled, also expose ListTools for full catalog discovery.
+    # Useful for smaller tool catalogs; wasteful for large ones.
+    enable_code_mode_list_tools: bool = Field(False, alias="ENABLE_CODE_MODE_LIST_TOOLS")
+
     # Skills configuration
     # Serve bundled HA best-practice skills as MCP resources (skill:// URIs).
     # Resources are not auto-injected — clients must explicitly request them.

src/ha_mcp/server.py

@@ -41,6 +41,97 @@
 ]
 
 
+# BM25 keyword boosts — appended to tool descriptions so BM25 ranks them
+# higher for common queries (from PR #727 search quality tuning).
+_SEARCH_KEYWORDS: dict[str, str] = {
+    "ha_search_entities": (
+        "find entities lookup discover search lights sensors switches "
+        "covers climate fans media_player binary_sensor"
+    ),
+    "ha_config_get_automation": (
+        "read inspect fetch view existing automation config triggers "
+        "conditions actions get show detail"
+    ),
+    "ha_config_set_helper": (
+        "create new add helper input_boolean input_number input_text "
+        "counter timer input_datetime input_select"
+    ),
+    "ha_config_get_script": (
+        "read inspect fetch view existing script config sequence "
+        "actions get show detail"
+    ),
+    "ha_get_entity": (
+        "get entity state attributes details single specific entity_id"
+    ),
+}
+
+# Description overrides — narrows the description so BM25 ranks the tool
+# LOWER for broad queries (prevents ha_deep_search from dominating).
+_SEARCH_OVERRIDES: dict[str, str] = {
+    "ha_deep_search": (
+        "Search INSIDE automation, script, and helper YAML configurations. "
+        "Use ONLY when you need to find where a specific service call, "
+        "entity reference, or config field appears within existing "
+        "automation/script/helper definitions. "
+        "NOT for finding entities or discovering tools."
+    ),
+}
+
+
+class _PatchedMontySandboxProvider:
+    """Patched sandbox provider that fixes FastMCP 3.1.0 bug.
+
+    FastMCP 3.1.0's MontySandboxProvider passes ``external_functions`` to the
+    ``Monty()`` constructor, but pydantic-monty doesn't accept it there.
+    Fixed on FastMCP main branch but not released in 3.1.0.
+    Remove this class once FastMCP >= 3.2.
+    """
+
+    def __init__(self, *, limits: Any = None):
+        self.limits = limits
+
+    async def run(
+        self,
+        code: str,
+        *,
+        inputs: dict[str, Any] | None = None,
+        external_functions: dict[str, Any] | None = None,
+    ) -> Any:
+        import asyncio
+        import importlib
+
+        pydantic_monty = importlib.import_module("pydantic_monty")
+
+        inputs = inputs or {}
+        async_functions = {}
+        for key, value in (external_functions or {}).items():
+            if asyncio.iscoroutinefunction(value):
+                async_functions[key] = value
+            else:
+                async_functions[key] = _make_async_wrapper(value)
+
+        # Fixed: don't pass external_functions to Monty() constructor
+        monty = pydantic_monty.Monty(
+            code,
+            inputs=list(inputs.keys()),
+        )
+        run_kwargs: dict[str, Any] = {"external_functions": async_functions}
+        if inputs:
+            run_kwargs["inputs"] = inputs
+        if self.limits is not None:
+            run_kwargs["limits"] = self.limits
+        return await pydantic_monty.run_monty_async(monty, **run_kwargs)
+
+
+def _make_async_wrapper(fn: Any) -> Any:
+    """Wrap a sync callable as async."""
+
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
+        return fn(*args, **kwargs)
+
+    return wrapper
+
+
 class HomeAssistantSmartMCPServer(EnhancedToolsMixin):
     """Home Assistant MCP Server with smart tools and fuzzy search.
 
@@ -137,6 +228,101 @@ def _initialize_server(self) -> None:
         # Register bundled skills as MCP resources
         self._register_skills()
 
+        # Apply Code Mode transform if enabled (must be after all tools are registered)
+        self._apply_code_mode()
+
+    def _apply_code_mode(self) -> None:
+        """Apply CodeMode transform if enabled via settings.
+
+        CodeMode (FastMCP 3.1 experimental) replaces all registered tools with
+        meta-tools (search, get_schema, execute) so the LLM discovers tools on
+        demand and chains calls in a Python sandbox, reducing context bloat.
+
+        Default 3-stage flow: search -> get_schema -> execute.
+        When ENABLE_CODE_MODE_LIST_TOOLS is also set, adds ListTools for full
+        catalog discovery (useful for smaller catalogs).
+
+        Includes a workaround for FastMCP 3.1.0 bug where MontySandboxProvider
+        passes `external_functions` to Monty() constructor which doesn't accept it.
+        """
+        if not self.settings.enable_code_mode:
+            return
+
+        try:
+            from fastmcp.experimental.transforms.code_mode import (
+                CodeMode,
+                GetSchemas,
+                ListTools,
+                Search,
+            )
+        except ImportError:
+            logger.warning(
+                "CodeMode not available — install fastmcp[code-mode] to enable. "
+                "Falling back to standard tool mode."
+            )
+            return
+
+        # Build discovery tools list — tuned for token efficiency:
+        # - Search returns "brief" (names + one-line descriptions only)
+        # - Limit search to top 10 results to avoid bloating context
+        # - LLM calls GetSchemas only for the specific tools it needs
+        discovery_tools: list = [
+            Search(default_limit=5),
+            GetSchemas(),
+        ]
+        if self.settings.enable_code_mode_list_tools:
+            discovery_tools.insert(0, ListTools())
+
+        # Use patched sandbox provider to work around FastMCP 3.1.0 bug:
+        # MontySandboxProvider passes external_functions to Monty() constructor
+        # but pydantic-monty doesn't accept that kwarg in __init__.
+        sandbox = _PatchedMontySandboxProvider(
+            limits={"max_duration_secs": 30},
+        )
+
+        try:
+            # Apply BM25 keyword tuning before CodeMode indexes tools.
+            # Boosts relevant tools and narrows ha_deep_search so it doesn't
+            # dominate broad queries (from PR #727 search quality work).
+            self._apply_search_keywords()
+
+            code_mode = CodeMode(
+                discovery_tools=discovery_tools,
+                sandbox_provider=sandbox,
+            )
+            self.mcp.add_transform(code_mode)
+            tool_names = [type(t).__name__ for t in discovery_tools]
+            logger.info(
+                "Code Mode enabled (discovery: %s + execute)",
+                ", ".join(tool_names),
+            )
+        except Exception:
+            logger.exception(
+                "Failed to apply CodeMode transform, falling back to standard mode"
+            )
+
+    def _apply_search_keywords(self) -> None:
+        """Modify tool descriptions to improve BM25 search ranking."""
+        from collections.abc import Sequence
+
+        from fastmcp.server.transforms import Transform
+
+        class _KeywordBoostTransform(Transform):
+            async def list_tools(self, tools: Sequence) -> Sequence:
+                result = []
+                for tool in tools:
+                    override = _SEARCH_OVERRIDES.get(tool.name)
+                    if override is not None:
+                        result.append(tool.model_copy(update={"description": override}))
+                    elif tool.name in _SEARCH_KEYWORDS:
+                        desc = f"{tool.description}\n{_SEARCH_KEYWORDS[tool.name]}"
+                        result.append(tool.model_copy(update={"description": desc}))
+                    else:
+                        result.append(tool)
+                return result
+
+        self.mcp.add_transform(_KeywordBoostTransform())
+
     def _get_skills_dir(self) -> Path | None:
         """Return the bundled skills directory if it exists.