feat(mcp): cache MCP tool vectors with lazy embedding in ToolManager

agent/tools/tool_manager.py

@@ -71,6 +71,22 @@ class ToolManager:
         if not hasattr(self, '_mcp_active_configs'):
             # server_name -> normalized config dict, for diff-based reload.
             self._mcp_active_configs: dict = {}
+        if not hasattr(self, '_mcp_tool_vectors'):
+            # mcp_tool_name -> embedding vector, used by on-demand tool
+            # retrieval. Populated lazily on first retrieval so users who
+            # never enable the feature pay zero embedding cost.
+            self._mcp_tool_vectors: dict = {}
+        if not hasattr(self, '_mcp_vector_lock'):
+            # Guards incremental index builds so concurrent turns don't
+            # double-embed the same newly-loaded MCP tools.
+            self._mcp_vector_lock = threading.Lock()
+        if not hasattr(self, '_embedding_provider_initialized'):
+            # The embedding provider is created once, lazily, and reused for
+            # both tool-index and per-query embeddings. None means keyword-only
+            # mode (no provider configured) — retrieval then falls back to full
+            # injection at the caller.
+            self._embedding_provider_initialized = False
+            self._embedding_provider = None
 
     def load_tools(self, tools_dir: str = "", config_dict=None):
         """
@@ -574,6 +590,91 @@ class ToolManager:
 
         return (sorted(added), sorted(removed))
 
+    # ------------------------------------------------------------------
+    # On-demand MCP tool retrieval support
+    #
+    # The vector index and the embedding provider are owned here (singleton,
+    # process-wide, aligned with the MCP tool lifecycle). The context-aware
+    # selection itself lives in agent.tools.mcp.tool_retrieval, driven by the
+    # executor which is the only place that knows the conversation context.
+    # ------------------------------------------------------------------
+
+    def count_mcp_tools(self) -> int:
+        """Return the number of currently loaded MCP tools."""
+        return len(self._mcp_tool_instances)
+
+    def get_mcp_tool_vectors(self) -> dict:
+        """Return ``{mcp_tool_name: vector}`` for currently loaded MCP tools.
+
+        Lazily embeds any MCP tools not yet in the cache (MCP servers load
+        asynchronously, so tools may appear over time). Returns an empty dict
+        when no embedding provider is available or embedding fails — the caller
+        then falls back to full injection. Never raises.
+        """
+        try:
+            self._ensure_mcp_tool_vectors()
+        except Exception as e:
+            logger.debug(f"[ToolManager] MCP tool vector build skipped: {e}")
+        return dict(self._mcp_tool_vectors)
+
+    def embed_query(self, text: str):
+        """Embed a retrieval query with the shared provider.
+
+        Returns the embedding vector, or None if no provider is available or
+        the call fails (caller falls back to full injection). Never raises.
+        """
+        if not text:
+            return None
+        provider = self._get_embedding_provider()
+        if provider is None:
+            return None
+        try:
+            return provider.embed_query(text)
+        except Exception as e:
+            logger.debug(f"[ToolManager] query embedding failed: {e}")
+            return None
+
+    def _ensure_mcp_tool_vectors(self) -> None:
+        """Incrementally embed MCP tools that are not yet cached."""
+        # Snapshot to avoid concurrent-mutation while the async loader runs.
+        current = dict(self._mcp_tool_instances)
+        missing = [name for name in current if name not in self._mcp_tool_vectors]
+        if not missing:
+            return
+
+        provider = self._get_embedding_provider()
+        if provider is None:
+            return
+
+        with self._mcp_vector_lock:
+            # Re-check under lock: another thread may have filled these in.
+            missing = [name for name in current if name not in self._mcp_tool_vectors]
+            if not missing:
+                return
+            texts = [self._mcp_tool_embed_text(current[name]) for name in missing]
+            vectors = provider.embed_batch(texts)
+            for name, vec in zip(missing, vectors):
+                self._mcp_tool_vectors[name] = vec
+
+    @staticmethod
+    def _mcp_tool_embed_text(tool) -> str:
+        """Build the text that represents an MCP tool for embedding."""
+        name = getattr(tool, "name", "") or ""
+        description = getattr(tool, "description", "") or ""
+        return f"{name}: {description}".strip()
+
+    def _get_embedding_provider(self):
+        """Lazily create and cache the shared embedding provider (or None)."""
+        if not self._embedding_provider_initialized:
+            try:
+                from agent.memory.embedding import create_default_embedding_provider
+                self._embedding_provider = create_default_embedding_provider()
+            except Exception as e:
+                logger.warning(f"[ToolManager] embedding provider init failed: {e}")
+                self._embedding_provider = None
+            self._embedding_provider_initialized = True
+        return self._embedding_provider
+
     def create_tool(self, name: str) -> BaseTool:
         """
         Get a new instance of a tool by name.