butler/mcp/deferred.py
Rama main · réplica de lectura del repositorio original.
"""MCP deferred tool discovery — compact catalog + on-demand schema load (PR-X4)."""
from __future__ import annotations
import json
import logging
import re
import threading
from pathlib import Path
from typing import Any, cast
from butler.core.harness_flags import mcp_deferred_tools_enabled
from butler.mcp.bridge import refs_to_openai_definitions
from butler.mcp.config import mcp_enabled
from butler.mcp.naming import is_mcp_registered_name
logger = logging.getLogger(__name__)
_lock = threading.RLock()
_promoted: dict[str, set[str]] = {}
def _session_key(session_key: str = "") -> str:
from butler.mcp.deferred_ops import resolve_mcp_session_key_safe
return cast(str, resolve_mcp_session_key_safe(session_key))
def get_promoted_tools(session_key: str = "") -> set[str]:
sk = _session_key(session_key)
with _lock:
return set(_promoted.get(sk) or ())
def promote_tools(names: list[str], *, session_key: str = "") -> list[str]:
sk = _session_key(session_key)
added: list[str] = []
with _lock:
bucket = _promoted.setdefault(sk, set())
for raw in names:
name = str(raw or "").strip()
if not name or not is_mcp_registered_name(name):
continue
if name not in bucket:
bucket.add(name)
added.append(name)
return added
def _tool_def_name(defn: dict[str, Any]) -> str:
raw_fn = defn.get("function")
fn = raw_fn if isinstance(raw_fn, dict) else {}
return str(fn.get("name") or defn.get("name") or "").strip()
def promote_experience_mcp_tools(
names: list[str],
*,
session_key: str = "",
) -> tuple[list[str], list[dict[str, str]]]:
"""Validate MCP refs exist, then promote. Returns ``(added, rejected)``."""
sk = _session_key(session_key)
cleaned = [str(n or "").strip() for n in (names or []) if str(n or "").strip()]
if not cleaned:
return [], []
if not mcp_enabled():
return [], [{"name": n, "reason": "mcp_disabled"} for n in cleaned]
rejected: list[dict[str, str]] = []
to_promote: list[str] = []
from butler.mcp.deferred_ops import warm_mcp_connection_safe
if not warm_mcp_connection_safe(sk, _all_mcp_refs):
return [], [{"name": n, "reason": "connect_failed"} for n in cleaned]
from butler.mcp.manager import get_manager
mgr = get_manager()
for name in cleaned:
if not is_mcp_registered_name(name):
rejected.append({"name": name, "reason": "invalid_registered_name"})
continue
if mgr.get_tool_ref(sk, name) is None:
rejected.append({"name": name, "reason": "tool_not_found"})
continue
to_promote.append(name)
added = promote_tools(to_promote, session_key=sk)
return added, rejected
def merge_deferred_mcp_into_turn_tools(
tools: list[dict[str, Any]],
*,
session_key: str = "",
) -> list[dict[str, Any]]:
"""Append promoted MCP OpenAI schemas to the per-turn tool list."""
mcp_defs = get_deferred_mcp_definitions(session_key)
if not mcp_defs:
return tools
seen = {_tool_def_name(t) for t in tools if _tool_def_name(t)}
out = list(tools)
for defn in mcp_defs:
name = _tool_def_name(defn)
if name and name not in seen:
out.append(defn)
seen.add(name)
return out
def clear_promoted(session_key: str = "") -> None:
sk = _session_key(session_key)
with _lock:
_promoted.pop(sk, None)
def _resolve_workspace() -> Path | None:
from butler.mcp.deferred_ops import resolve_mcp_workspace_safe
return cast(Path | None, resolve_mcp_workspace_safe())
def _all_mcp_refs(session_key: str = "") -> list[Any]:
from butler.mcp.manager import get_manager
sk = _session_key(session_key)
return cast(
list[Any],
get_manager().ensure_connected(sk, workspace=_resolve_workspace()),
)
def search_mcp_tools(
query: str,
*,
limit: int = 12,
session_key: str = "",
) -> list[dict[str, Any]]:
if not mcp_enabled():
return []
q = str(query or "").strip().lower()
refs = _all_mcp_refs(session_key)
scored: list[tuple[int, Any]] = []
for ref in refs:
hay = f"{ref.registered_name} {ref.original_name} {ref.description}".lower()
if not q:
score = 0
elif q in hay:
score = 2 if q in ref.registered_name.lower() else 1
else:
tokens = [t for t in re.split(r"\W+", q) if len(t) >= 2]
score = sum(1 for t in tokens if t in hay)
if q and score <= 0:
continue
scored.append((score, ref))
scored.sort(key=lambda x: (-x[0], x[1].registered_name))
out: list[dict[str, Any]] = []
for _, ref in scored[: max(1, min(50, int(limit)))]:
out.append({
"registered_name": ref.registered_name,
"server_id": ref.server_id,
"original_name": ref.original_name,
"classification": ref.classification,
"description": (ref.description or "")[:240],
})
return out
def get_deferred_mcp_definitions(session_key: str = "") -> list[dict[str, Any]]:
"""Full OpenAI schemas only for session-promoted MCP tools."""
if not mcp_enabled() or not mcp_deferred_tools_enabled():
return []
sk = _session_key(session_key)
promoted = get_promoted_tools(sk)
if not promoted:
return []
refs = [r for r in _all_mcp_refs(sk) if r.registered_name in promoted]
return cast(list[dict[str, Any]], refs_to_openai_definitions(refs))
def tool_search_handler(query: str, limit: int = 12, promote: bool = False) -> str:
matches = search_mcp_tools(query, limit=limit)
promoted: list[str] = []
if promote and matches:
promoted = promote_tools(
[str(m.get("registered_name") or "") for m in matches],
session_key=_session_key(),
)
return json.dumps(
{
"ok": True,
"code": "MCP_TOOL_SEARCH",
"query": query,
"count": len(matches),
"tools": matches,
"promoted": promoted,
},
ensure_ascii=False,
)
def load_mcp_tools_handler(tool_names: list[str]) -> str:
names = [str(n or "").strip() for n in (tool_names or []) if str(n or "").strip()]
added = promote_tools(names, session_key=_session_key())
return json.dumps(
{
"ok": True,
"code": "MCP_LOAD_TOOLS",
"promoted": added,
"session_promoted": sorted(get_promoted_tools(_session_key())),
},
ensure_ascii=False,
)
__all__ = [
"clear_promoted",
"get_deferred_mcp_definitions",
"get_promoted_tools",
"load_mcp_tools_handler",
"merge_deferred_mcp_into_turn_tools",
"promote_experience_mcp_tools",
"promote_tools",
"search_mcp_tools",
"tool_search_handler",
]