324 lines
17 KiB
Python
324 lines
17 KiB
Python
"""Live-agent tool-list maintenance after MCP (re)discovery: refreshing an
|
|
AIAgent's tools/tool names, preserving the cached tools[] prefix across rebuilds,
|
|
and re-injecting post-build tools."""
|
|
|
|
import logging
|
|
import json
|
|
import threading
|
|
from typing import Optional
|
|
from tools.mcp_tool_common import _core
|
|
|
|
logger = logging.getLogger("tools.mcp_tool")
|
|
|
|
# Serializes in-place swaps of ``agent.tools`` / ``agent.valid_tool_names`` by
|
|
# the reload RPC, gateway reload and late-binding refresh thread; the run loop
|
|
# reads them during tool iteration and must never see a half-updated pair.
|
|
_agent_tools_lock = threading.Lock()
|
|
|
|
|
|
def _def_name(tool_def: dict) -> str:
|
|
return (tool_def.get("function") or {}).get("name", "")
|
|
|
|
|
|
def _agent_tool_defs(agent) -> list:
|
|
return list(getattr(agent, "tools", None) or [])
|
|
|
|
|
|
def agent_tool_names(agent) -> list:
|
|
"""Names of ``agent.tools`` in wire order (unnamed entries skipped)."""
|
|
return [name for name in map(_def_name, _agent_tool_defs(agent)) if name]
|
|
|
|
|
|
def _resolve_refresh_toolsets(agent, enabled_override, disabled_override):
|
|
"""Explicit reloads pass freshly-resolved toolsets (so a server just ENABLED in config is
|
|
picked up) and the agent's selection is updated to match; automatic paths pass nothing
|
|
and reuse the build-time selection."""
|
|
enabled = getattr(agent, "enabled_toolsets", None)
|
|
disabled = getattr(agent, "disabled_toolsets", None)
|
|
if enabled_override is not None or disabled_override is not None:
|
|
enabled = enabled_override if enabled_override is not None else enabled
|
|
if disabled_override is not None:
|
|
from toolsets import session_disabled_toolsets
|
|
disabled = session_disabled_toolsets(disabled_override, getattr(agent, "platform", None))
|
|
agent.enabled_toolsets, agent.disabled_toolsets = enabled, disabled
|
|
return enabled, disabled
|
|
|
|
|
|
def _tool_defs_content_changed(agent, new_defs: list) -> bool:
|
|
"""Byte-level diff of the serialized tool arrays (dynamic schemas change CONTENT under
|
|
stable names); False if either side fails to serialize."""
|
|
try:
|
|
dump = lambda defs: json.dumps(defs, sort_keys=True, separators=(",", ":"), default=str)
|
|
return dump(_agent_tool_defs(agent)) != dump(new_defs)
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def _drop_session_tools(agent, new_defs: list, new_names: set) -> tuple:
|
|
from toolsets import agent_tool_drops
|
|
drops = agent_tool_drops(agent)
|
|
if not drops:
|
|
return new_defs, new_names
|
|
return [entry for entry in new_defs if _def_name(entry) not in drops], new_names - drops
|
|
|
|
|
|
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 published (unchanged, or a newer snapshot won)."""
|
|
with _agent_tools_lock:
|
|
# Tolerate an agent that never set the generation (or a non-int mock).
|
|
published_gen = getattr(agent, "_tool_snapshot_generation", -1)
|
|
published_gen = published_gen if isinstance(published_gen, int) else -1
|
|
if snapshot_generation < published_gen:
|
|
return None # a newer snapshot already won
|
|
current_defs = _agent_tool_defs(agent)
|
|
current = {_def_name(t) for t in current_defs}
|
|
if prefix_registered is not None:
|
|
new_defs, new_names = _merge_preserving_prefix(current_defs, new_defs, prefix_registered)
|
|
new_defs, new_names = _drop_session_tools(agent, new_defs, new_names)
|
|
# Record the generation even when unchanged so an in-flight older caller can't clobber.
|
|
agent._tool_snapshot_generation = max(published_gen, snapshot_generation)
|
|
# Same NAME set: no change for MCP-reload callers. Content-aware callers
|
|
# (compaction boundary) also diff serialized bytes.
|
|
if new_names == current and not (content_aware and _tool_defs_content_changed(agent, new_defs)):
|
|
return None
|
|
agent.tools = new_defs
|
|
agent.valid_tool_names = new_names
|
|
# Publish context-engine routing names atomically with the snapshot.
|
|
engine_names = getattr(agent, "_context_engine_tool_names", None)
|
|
if isinstance(engine_names, set):
|
|
engine_names.clear()
|
|
engine_names.update(staged_engine_names)
|
|
return new_names - current
|
|
|
|
|
|
def refresh_agent_mcp_tools(
|
|
agent, *, enabled_override=None, disabled_override=None, quiet_mode: bool = True,
|
|
content_aware: bool = False, preserve_prefix: bool = False) -> set:
|
|
"""Re-derive an already-built agent's tool snapshot from the live registry; returns the
|
|
newly-added tool names (empty when unchanged). The agent snapshots ``agent.tools`` at build
|
|
time, so servers that connect later (slow OAuth, ``/reload-mcp``) are invisible until
|
|
rebuilt. Shared by the TUI RPC, gateway reload, late-binding thread and between-turns
|
|
refresh: respects the toolset filter, diffs by tool NAME (a count compare misses an
|
|
equal-size swap), re-injects the memory-provider / context-engine tools ``agent_init``
|
|
appends after ``get_tool_definitions`` plus guarded session capabilities, and publishes
|
|
``(tools, valid_tool_names)`` together.
|
|
|
|
``preserve_prefix``: for rebuilds inside a live conversation the tool array is a cached
|
|
request prefix and any moved byte re-prefills the whole history — existing tools keep their
|
|
slot (schemas still refresh), a still-registered tool whose ``check_fn`` merely flapped is
|
|
carried forward (``check_fn`` gates exposure, never invocation), a deregistered tool is
|
|
dropped, new tools append at the tail. The caller owns the prompt-cache contract."""
|
|
from model_tools import get_tool_definitions
|
|
from tools.registry import registry
|
|
enabled, disabled = _resolve_refresh_toolsets(agent, enabled_override, disabled_override)
|
|
# Generation captured BEFORE the slow get_tool_definitions call (a slower caller holding an
|
|
# OLDER set must not clobber a newer one); definitions computed OUTSIDE the lock.
|
|
snapshot_generation = registry._generation
|
|
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}
|
|
# Post-build families re-appended on LOCALS only; live attributes untouched until publish.
|
|
staged_engine_names = _reinject_post_build_tools(agent, new_defs, new_names)
|
|
_reinject_authorized_dynamic_tools(agent, new_defs, new_names)
|
|
# Registry membership is read OUTSIDE ``_agent_tools_lock``: taking ``registry._lock``
|
|
# under the tools lock would be the first nesting of the two.
|
|
prefix_registered: Optional[set] = None
|
|
if preserve_prefix:
|
|
try:
|
|
prefix_registered = {entry.name for entry in registry.get_all_entries()}
|
|
except Exception:
|
|
pass # fail open to the plain rebuild
|
|
added = _publish_tool_snapshot(
|
|
agent, new_defs, new_names, snapshot_generation=snapshot_generation,
|
|
staged_engine_names=staged_engine_names, content_aware=content_aware, prefix_registered=prefix_registered)
|
|
if added is None:
|
|
return set()
|
|
persist_agent_tool_names(agent) # re-pin so a rebuild after agent-cache eviction restores this order
|
|
return added
|
|
|
|
|
|
def reprobe_tool_availability() -> None:
|
|
"""Explicit ``/reload-mcp`` hatch out of the tools[] freeze: drop the ``check_fn`` verdict
|
|
cache AND the ``get_tool_definitions`` memo (keyed on registry generation, so it would
|
|
otherwise replay the stale verdicts)."""
|
|
from model_tools import _clear_tool_defs_cache
|
|
from tools.registry import invalidate_check_fn_cache
|
|
invalidate_check_fn_cache()
|
|
_clear_tool_defs_cache()
|
|
|
|
|
|
def tool_pin_version() -> str:
|
|
"""The code identity a tools[] pin was built by (checkout/build sha, else the release version).
|
|
Cached per process: an updated checkout only reaches a process through a restart."""
|
|
from hermes_cli import __version__
|
|
from hermes_cli.build_info import get_code_identity
|
|
identity = get_code_identity()
|
|
return identity.get("sha") or identity.get("version") or __version__
|
|
|
|
|
|
def persist_agent_tool_names(agent) -> None:
|
|
"""Best-effort: write ``agent.tools`` to the session row (freeze pin). The full definitions,
|
|
keyed by the code that built them: another process or surface derives different bytes."""
|
|
db = getattr(agent, "_session_db", None)
|
|
session_id = getattr(agent, "session_id", None)
|
|
if not db or not session_id:
|
|
return
|
|
try:
|
|
db.update_session_tool_names(session_id, {"version": tool_pin_version(), "tools": _agent_tool_defs(agent)})
|
|
except Exception:
|
|
logger.debug("tool_names persist skipped", exc_info=True)
|
|
|
|
|
|
def _config_permitted_names(agent) -> set:
|
|
"""Tool names this agent's toolset selection allows before ``check_fn``: all a pin may carry
|
|
forward. A client-surface toolset counts as allowed (only its client can add it, so its absence
|
|
here is no config choice); ``disabled_toolsets`` still strips it."""
|
|
from model_tools import _select_tool_names
|
|
from toolsets import CLIENT_SURFACE_TOOLSETS
|
|
enabled = getattr(agent, "enabled_toolsets", None)
|
|
if enabled is not None:
|
|
enabled = [*enabled, *CLIENT_SURFACE_TOOLSETS]
|
|
return _select_tool_names(enabled, getattr(agent, "disabled_toolsets", None), True)
|
|
|
|
|
|
def _drop_gated_carried_tools(merged: list, carried: set) -> list:
|
|
"""A carried tool also passes the session-level schema gates the fresh build applied
|
|
(``browser_exec`` needs ``terminal`` in the same array), judged on the merged array."""
|
|
from model_tools import _DYNAMIC_SCHEMA_REWRITERS
|
|
available = {_def_name(t) for t in merged}
|
|
return [t for t in merged if _def_name(t) not in carried or _def_name(t) not in _DYNAMIC_SCHEMA_REWRITERS
|
|
or _DYNAMIC_SCHEMA_REWRITERS[_def_name(t)](t, available) is not None]
|
|
|
|
|
|
def restore_agent_tool_prefix(agent, saved) -> bool:
|
|
"""Fold a freshly built agent's ``tools`` onto the session's pin; True if changed.
|
|
A fresh AIAgent (gateway cache eviction, ``--resume`` in a new process, a surface hop) has no
|
|
predecessor to preserve, so the pin stands in. Pinned by the SAME code, a tool still available
|
|
here keeps its pinned BYTES, whatever this process derives for it (tool_search's per-surface
|
|
catalog, per-surface dynamic parameters, the ``-q`` footprint): tools[] heads every request.
|
|
Pinned by other code (``hermes update``, a legacy name list) a tool's contract may have moved,
|
|
so each takes its current definition. A pinned tool this process did not build is carried
|
|
only while its toolset config allows it here; deregistered tools drop, new tools append."""
|
|
pinned, version = (saved.get("tools") or [], saved.get("version")) if isinstance(saved, dict) else (saved, None)
|
|
if not pinned:
|
|
return False
|
|
from tools.registry import registry
|
|
fresh_defs = _agent_tool_defs(agent)
|
|
fresh = {_def_name(t): t for t in fresh_defs}
|
|
registered_names = {entry.name for entry in registry.get_all_entries()}
|
|
same_code = version is not None and version == tool_pin_version()
|
|
|
|
def _pinned_def(item):
|
|
name = item if isinstance(item, str) else _def_name(item)
|
|
if isinstance(item, dict) and same_code:
|
|
return item
|
|
if name in fresh:
|
|
return fresh[name]
|
|
entry = registry.get_entry(name)
|
|
return None if entry is None else {"type": "function", "function": {**entry.schema, "name": entry.name}}
|
|
|
|
pinned_defs = [d for d in map(_pinned_def, pinned) if d]
|
|
carried = {_def_name(d) for d in pinned_defs if _def_name(d) not in fresh and _def_name(d) in registered_names}
|
|
carried &= _config_permitted_names(agent) if carried else set()
|
|
merged = [d for d in pinned_defs if _def_name(d) in fresh or _def_name(d) in carried]
|
|
pinned_names = {_def_name(d) for d in merged}
|
|
merged.extend(t for t in fresh_defs if _def_name(t) not in pinned_names)
|
|
merged = _drop_gated_carried_tools(merged, carried)
|
|
merged_names = {_def_name(t) for t in merged}
|
|
_reinject_authorized_dynamic_tools(agent, merged, merged_names)
|
|
merged, merged_names = _drop_session_tools(agent, merged, merged_names)
|
|
changed = merged != fresh_defs
|
|
if changed:
|
|
with _agent_tools_lock:
|
|
agent.tools = merged
|
|
agent.valid_tool_names = merged_names
|
|
if not same_code or merged != list(pinned):
|
|
persist_agent_tool_names(agent)
|
|
return changed
|
|
|
|
|
|
def _merge_preserving_prefix(current_defs: list, new_defs: list, registered_names: set) -> tuple[list, set]:
|
|
"""Fold a fresh tool snapshot into a live one without moving existing bytes. Ordered by
|
|
``current_defs`` (the cached request prefix): a name in both keeps its slot but takes the
|
|
fresh schema; a name only in the live list is kept if still registered (``check_fn``
|
|
flapped), else dropped; a name only in the fresh list is appended at the tail.
|
|
|
|
The bridge tools keep their BUILT entry, not the fresh one: ``tool_search``'s description
|
|
is derived from the session (deferred count, listing, whether ``manage_connections`` was
|
|
present), so a late MCP server or a ``check_fn`` flap would rewrite it every turn. Search
|
|
reads the live catalog at dispatch, so the stale count costs nothing."""
|
|
from tools.tool_search_catalog import BRIDGE_TOOL_NAMES
|
|
fresh = {_def_name(entry): entry for entry in new_defs if _def_name(entry)}
|
|
merged = []
|
|
for entry in current_defs:
|
|
name = _def_name(entry)
|
|
replacement = fresh.pop(name, None)
|
|
if name in BRIDGE_TOOL_NAMES:
|
|
merged.append(entry)
|
|
elif replacement is not None:
|
|
merged.append(replacement)
|
|
elif name and name in registered_names:
|
|
merged.append(entry)
|
|
merged.extend(fresh.values())
|
|
return merged, {_def_name(t) for t in merged}
|
|
|
|
|
|
def _reinject_authorized_dynamic_tools(agent, tools_list: list, name_set: set) -> None:
|
|
"""``message_agent`` is injected by an auth gate, never registered, so a registry-derived
|
|
rebuild drops it. Scrub any stale copy from the STAGED pair and re-add it only when the live
|
|
gate re-authorizes, so the publisher exposes a coherent ``(tools, valid_tool_names)``."""
|
|
from tools.bot_mode_dm import MESSAGE_AGENT_TOOL_NAME, message_agent_authorized, message_agent_tool_schema
|
|
|
|
tools_list[:] = [entry for entry in tools_list if _def_name(entry) != MESSAGE_AGENT_TOOL_NAME]
|
|
name_set.discard(MESSAGE_AGENT_TOOL_NAME)
|
|
if message_agent_authorized(agent):
|
|
tools_list.append(message_agent_tool_schema())
|
|
name_set.add(MESSAGE_AGENT_TOOL_NAME)
|
|
|
|
|
|
def _reinject_post_build_tools(agent, tools_list: list, name_set: set) -> set:
|
|
"""Append memory-provider and context-engine tools onto the caller's staged ``tools_list``
|
|
/ ``name_set`` (never the live agent attributes), mirroring ``agent_init``'s post-build
|
|
injection. Idempotent and fail-soft. Returns the context-engine routing names THIS rebuild
|
|
appended: a name already owned by a registry/plugin tool is not claimed, matching agent_init."""
|
|
def _add(schema) -> bool:
|
|
name = schema.get("name", "") if isinstance(schema, dict) else ""
|
|
if not name or name in name_set:
|
|
return False
|
|
tools_list.append({"type": "function", "function": schema})
|
|
name_set.add(name)
|
|
return True
|
|
|
|
def _schema_getter(attr: str, method: str):
|
|
getter = getattr(getattr(agent, attr, None) or None, method, None)
|
|
return getter if callable(getter) else None
|
|
|
|
enabled = getattr(agent, "enabled_toolsets", None)
|
|
try:
|
|
get_mem_schemas = _schema_getter("_memory_manager", "get_all_tool_schemas")
|
|
if get_mem_schemas is not None:
|
|
from agent.memory_manager import memory_provider_tools_enabled # same gate inject_memory_provider_tools uses
|
|
if memory_provider_tools_enabled(
|
|
enabled, getattr(agent, "disabled_toolsets", None), memory_tool_present="memory" in name_set):
|
|
for schema in get_mem_schemas():
|
|
_add(schema)
|
|
except Exception:
|
|
logger.debug("Memory-provider tool re-injection skipped", exc_info=True)
|
|
|
|
# The `context_engine` toolset is intentionally empty, so lcm_* tools exist only via this
|
|
# append. Honor the enabled_toolsets gate agent_init uses, or a restricted-toolset platform
|
|
# would re-leak tools the build excluded.
|
|
# See #5544.
|
|
staged_engine_names: set = set()
|
|
try:
|
|
get_schemas = _schema_getter("context_compressor", "get_tool_schemas")
|
|
if (enabled is None or "context_engine" in enabled) and get_schemas is not None:
|
|
# Claim the routing name only when WE appended the schema.
|
|
staged_engine_names.update(s["name"] for s in get_schemas() if _add(s))
|
|
except Exception:
|
|
logger.debug("Context-engine tool re-injection skipped", exc_info=True)
|
|
return staged_engine_names
|