refactor(mcp): wrap overlong lines in reg/sampling/schema/agent

This commit is contained in:
Teknium
2026-09-02 16:01:24 -07:00
parent a3a9cacf93
commit d9f6d9eda0
4 changed files with 54 additions and 21 deletions

View File

@@ -48,8 +48,10 @@ def _tool_defs_content_changed(agent, new_defs: list) -> bool:
return False
def _publish_tool_snapshot(agent, new_defs: list, new_names: set, *, snapshot_generation: int,
staged_engine_names: set, content_aware: bool, prefix_registered: Optional[set]) -> Optional[set]:
def _publish_tool_snapshot(
agent, new_defs: list, new_names: set, *, snapshot_generation: int,
staged_engine_names: set, content_aware: bool, prefix_registered: Optional[set],
) -> Optional[set]:
"""Single atomic read-diff-publish under ``_agent_tools_lock`` so ``added``
matches what was published and a stale (older-generation) rebuild can't
overwrite a newer one. Returns the added names, or None when nothing was
@@ -120,7 +122,9 @@ def refresh_agent_mcp_tools(
snapshot_generation = registry._generation
# Computed OUTSIDE the lock (can be slow); diff + publish happen together in
# one critical section so concurrent callers can't torn-publish.
new_defs = list(get_tool_definitions(enabled_toolsets=enabled, disabled_toolsets=disabled, quiet_mode=quiet_mode) or [])
new_defs = list(
get_tool_definitions(enabled_toolsets=enabled, disabled_toolsets=disabled, quiet_mode=quiet_mode) or []
)
new_names = {_def_name(t) for t in new_defs}
# Re-append the post-build families on LOCALS only; live agent attributes
# are untouched until the single atomic publish.

View File

@@ -95,7 +95,9 @@ def _select_utility_schemas(server_name: str, server: "MCPServerTask", config: d
if not enabled[family]:
return f"{family} disabled"
if advertised is not None:
return None if getattr(advertised, family, None) is not None else f"server does not advertise '{family}' capability"
if getattr(advertised, family, None) is None:
return f"server does not advertise '{family}' capability"
return None
method = _UTILITY_CAPABILITY_METHODS[handler_key]
return None if hasattr(server.session, method) else f"session lacks {method}"
@@ -163,7 +165,10 @@ class _CachedMCPTool:
for raw in raws:
if isinstance(raw, dict) and raw.get("name"):
schema = raw.get("inputSchema")
out.append(cls(raw["name"], raw.get("description") or "", schema if isinstance(schema, dict) else {}, raw.get("annotations")))
out.append(cls(
raw["name"], raw.get("description") or "",
schema if isinstance(schema, dict) else {}, raw.get("annotations"),
))
return out
@@ -182,7 +187,9 @@ class _Candidate:
return self.origin.startswith(_UTILITY_ORIGIN_PREFIX)
def _tool_candidates(name: str, tools: Iterable[Any], should_register: Callable[[str], bool], tool_timeout) -> List[_Candidate]:
def _tool_candidates(
name: str, tools: Iterable[Any], should_register: Callable[[str], bool], tool_timeout,
) -> List[_Candidate]:
"""Native tools (live SDK objects or ``_CachedMCPTool``) -> candidates. The
injection scan runs on BOTH paths: the cache file is user-writable JSON."""
out: List[_Candidate] = []
@@ -192,7 +199,8 @@ def _tool_candidates(name: str, tools: Iterable[Any], should_register: Callable[
continue
_core._scan_mcp_description(name, t.name, t.description or "")
schema = _core._convert_mcp_schema(name, t)
out.append(_Candidate(schema["name"], f"tool {t.name!r}", schema, _core._make_tool_handler(name, t.name, tool_timeout)))
handler = _core._make_tool_handler(name, t.name, tool_timeout)
out.append(_Candidate(schema["name"], f"tool {t.name!r}", schema, handler))
return out
@@ -204,7 +212,8 @@ def _utility_candidates(name: str, entries: Iterable[Any], tool_timeout) -> List
continue
schema, key = raw.get("schema"), raw.get("handler_key")
if isinstance(schema, dict) and key in _UTILITY_HANDLER_FACTORIES and schema.get("name"):
out.append(_Candidate(schema["name"], f"{_UTILITY_ORIGIN_PREFIX}{key!r}", schema, _UTILITY_HANDLER_FACTORIES[key](name, tool_timeout)))
handler = _UTILITY_HANDLER_FACTORIES[key](name, tool_timeout)
out.append(_Candidate(schema["name"], f"{_UTILITY_ORIGIN_PREFIX}{key!r}", schema, handler))
return out
@@ -220,7 +229,10 @@ def _resolve_name_collisions(name: str, candidates: List[_Candidate]) -> List[_C
origins_by_name: Dict[str, set[str]] = {}
for c in candidates:
if (c.registry_name, c.origin) in seen:
logger.debug("MCP server '%s': duplicate registration candidate %s for '%s'; keeping one", name, c.origin, c.registry_name)
logger.debug(
"MCP server '%s': duplicate registration candidate %s for '%s'; keeping one",
name, c.origin, c.registry_name,
)
continue
seen.add((c.registry_name, c.origin))
unique.append(c)
@@ -256,17 +268,26 @@ def _log_foreign_owner(name: str, c: _Candidate, existing_toolset: str, lazy: bo
"""Diagnostics for a name already owned by another toolset (skipped to preserve the owner)."""
if lazy:
if not c.is_utility:
logger.warning("MCP server '%s' (lazy): cached tool '%s' collides with toolset '%s' — skipping", name, c.registry_name, existing_toolset)
logger.warning(
"MCP server '%s' (lazy): cached tool '%s' collides with toolset '%s' — skipping",
name, c.registry_name, existing_toolset,
)
return
log, fmt = (
(logger.error, "MCP server '%s': %s normalizes to '%s', already owned by MCP toolset '%s' — skipping to preserve the existing owner")
if existing_toolset.startswith("mcp-") else
(logger.warning, "MCP server '%s': %s (→ '%s') collides with built-in tool in toolset '%s' — skipping to preserve built-in")
)
if existing_toolset.startswith("mcp-"):
log, fmt = logger.error, (
"MCP server '%s': %s normalizes to '%s', already owned by MCP toolset '%s' "
"— skipping to preserve the existing owner"
)
else:
log, fmt = logger.warning, (
"MCP server '%s': %s (→ '%s') collides with built-in tool in toolset '%s' — skipping to preserve built-in"
)
log(fmt, name, c.origin, c.registry_name, existing_toolset)
def _register_candidates(name: str, candidates: List[_Candidate], *, check_fn: Callable, scope: Callable[[], Optional[str]], lazy: bool) -> List[str]:
def _register_candidates(
name: str, candidates: List[_Candidate], *, check_fn: Callable, scope: Callable[[], Optional[str]], lazy: bool,
) -> List[str]:
"""Register candidates under toolset ``mcp-{name}``; returns the names that
landed. The ownership pre-check is advisory only — servers connect in
parallel, so ``ToolRegistry.register()`` is the atomic ownership gate and
@@ -287,7 +308,8 @@ def _register_candidates(name: str, candidates: List[_Candidate], *, check_fn: C
if registry.get_toolset_for_tool(c.registry_name) != toolset_name:
if not lazy:
logger.error(
"MCP server '%s': registration of %s as '%s' was rejected by the registry; skipping provenance/count updates",
"MCP server '%s': registration of %s as '%s' was rejected by the registry; "
"skipping provenance/count updates",
name, c.origin, c.registry_name,
)
continue
@@ -315,7 +337,9 @@ def _write_schema_cache(name: str, server: "MCPServerTask", config: dict, should
# Persisted so the lazy path trust-gates identically next startup.
"annotations": {"readOnlyHint": _annotation_read_only_hint(t)},
})
utility_payload = [{"schema": e["schema"], "handler_key": e["handler_key"]} for e in _select_utility_schemas(name, server, config)]
utility_payload = [
{"schema": e["schema"], "handler_key": e["handler_key"]} for e in _select_utility_schemas(name, server, config)
]
cache_meta = getattr(server, "_list_cache_meta", None) or {}
write_cache_entry(
name, config_fingerprint(config), tools=tools_payload, utility_tools=utility_payload,
@@ -360,7 +384,9 @@ def _register_from_cache_sync(name: str, config: dict, entry: dict) -> List[str]
_record_tool_trust_metadata(name, config, cached_tools)
candidates = _tool_candidates(name, cached_tools, _make_tool_filter(name, config), tool_timeout)
candidates += _utility_candidates(name, utility_tools_from_cache_entry(entry), tool_timeout)
registered = _register_candidates(name, candidates, check_fn=_make_check_fn(name), scope=_core._mcp_registry_scope, lazy=True)
registered = _register_candidates(
name, candidates, check_fn=_make_check_fn(name), scope=_core._mcp_registry_scope, lazy=True,
)
if registered:
with _core._lock:
_core._lazy_server_configs[name] = dict(config)

View File

@@ -342,7 +342,9 @@ class ElicitationHandler:
def _result(self, action: str, metric: str):
"""Count *metric* and return ``ElicitResult(action)`` (accept carries empty content)."""
self.metrics[metric] += 1
return _core.ElicitResult(action=action, content={}) if action == "accept" else _core.ElicitResult(action=action)
if action == "accept":
return _core.ElicitResult(action="accept", content={})
return _core.ElicitResult(action=action)
def _consent_thunk(self, message: str, description: str) -> Callable[[], str]:
"""Sync consent call, replaying the agent's contextvars snapshot when the

View File

@@ -132,7 +132,8 @@ _MCP_NAME_DELIM = "__"
def mcp_prefixed_tool_name(server_name: str, tool_name: str) -> str:
"""Registry/wire name: ``mcp__<sanitizedServer>__<sanitizedTool>``."""
return f"{MCP_TOOL_NAME_PREFIX}{sanitize_mcp_name_component(server_name)}{_MCP_NAME_DELIM}{sanitize_mcp_name_component(tool_name)}"
safe_server, safe_tool = sanitize_mcp_name_component(server_name), sanitize_mcp_name_component(tool_name)
return f"{MCP_TOOL_NAME_PREFIX}{safe_server}{_MCP_NAME_DELIM}{safe_tool}"
def _convert_mcp_schema(server_name: str, mcp_tool) -> dict: