mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-11-06 05:17:15 +00:00
Teknium
c7993457bc
- Budget prompts for one chat run one at a time and re-read the total, so a parallel call sees what the previous approval booked and identical prompts can't coalesce into one answer that releases several charges. - Totals are never evicted (a forgotten total would renew a spent budget) and move with their policy when a process starts multiplexing. - A native entry rejected for a malformed budget is not revived by a same-name package's looser one.
105 lines
5.3 KiB
Python
105 lines
5.3 KiB
Python
"""Connection-ledger keys for tools.mcp_tool under a profile multiplexer.
|
|
|
|
Every ledger in ``tools.mcp_tool`` (``_servers``, connecting/error/cooldown maps, circuit
|
|
breaker, lazy configs, trust metadata) is keyed by a *connection key*: the bare server name
|
|
outside a multiplexer (single-profile processes are unchanged, byte for byte), and
|
|
``(owner_scope, name)`` under one. Two profiles that both configure ``github`` with their own
|
|
token are two connections; keying by name alone let the first profile's connection shadow the
|
|
second's forever — its ``register_mcp_servers`` saw the name as "already connected", adopted
|
|
nothing (different credentials) and left the profile silently tool-less (#106005, #91654).
|
|
|
|
A profile may still *adopt* another profile's live connection when the route and credentials
|
|
match (``mcp_tool_registration._same_server_route``); ``_server_tool_scopes[key]`` records every
|
|
scope that has done so, and ``_resolve_server_key`` finds that shared connection for a caller
|
|
whose own scope has none.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Optional, Tuple, Union
|
|
|
|
from tools.mcp_tool_common import _core
|
|
|
|
ServerKey = str | tuple[str, str]
|
|
|
|
|
|
def _server_key(name: str, scope: Optional[str] = None, *, current: bool = True) -> ServerKey:
|
|
"""Connection key for *name* owned by *scope* (the current registry scope when *current*).
|
|
``None`` scope (no multiplexer) keeps the bare name."""
|
|
if scope is None and current:
|
|
scope = _core._mcp_registry_scope()
|
|
return name if scope is None else (scope, name)
|
|
|
|
|
|
def _key_name(key: ServerKey) -> str:
|
|
return key[1] if isinstance(key, tuple) else key
|
|
|
|
|
|
def _key_scope(key: ServerKey) -> Optional[str]:
|
|
"""Owning registry scope encoded in *key* (None for a bare, unscoped key)."""
|
|
return key[0] if isinstance(key, tuple) else None
|
|
|
|
|
|
def _key_visible_in_scope(key: ServerKey, scope: Optional[str]) -> bool:
|
|
"""Whether the connection under *key* serves *scope*: owned by it or adopted into it.
|
|
Caller holds ``_core._lock`` or tolerates a racy read (status surfaces)."""
|
|
if scope is None:
|
|
return True
|
|
return _key_scope(key) == scope or scope in _core._server_tool_scopes.get(key, ())
|
|
|
|
|
|
def _resolve_server_key(name: str, scope: Optional[str] = None, *, current: bool = True) -> ServerKey:
|
|
"""The connection key a call to *name* from *scope* must use: the scope's own connection
|
|
(live, connecting or lazily registered) when it has one, else a shared connection it
|
|
adopted, else its own (not yet existing) key so bookkeeping lands under this scope."""
|
|
if scope is None and current:
|
|
scope = _core._mcp_registry_scope()
|
|
own = _server_key(name, scope, current=False)
|
|
if scope is None or own in _core._servers or own in _core._lazy_server_configs:
|
|
return own
|
|
for key, scopes in _core._server_tool_scopes.items():
|
|
if scope in scopes and _key_name(key) == name and key in _core._servers:
|
|
return key
|
|
return own
|
|
|
|
|
|
# Ledgers keyed by connection key (``tools.mcp_tool``).
|
|
_KEYED_LEDGERS = (
|
|
"_servers", "_server_scope_keys", "_server_tool_scopes", "_server_connect_errors",
|
|
"_lazy_server_configs", "_lazy_server_fingerprints", "_lazy_server_tool_names",
|
|
"_server_connect_retry_after", "_server_connect_failures", "_server_error_counts",
|
|
"_server_breaker_opened_at", "_server_errors_all_application", "_server_trust_levels",
|
|
"_tool_read_only_hints", "_server_cost_policies")
|
|
_KEYED_SETS = ("_server_connecting", "_parallel_safe_servers")
|
|
|
|
|
|
def rekey_unscoped_connections(scope: str) -> None:
|
|
"""Move every bare-name connection under ``(scope, name)``. A single-profile process keys its
|
|
connections by bare name; once it starts multiplexing, the launch profile's lookups resolve
|
|
``(launch scope, name)``, so without this its live servers read as disconnected until a restart."""
|
|
from tools.registry import registry
|
|
with _core._lock:
|
|
for attr in _KEYED_LEDGERS:
|
|
ledger = getattr(_core, attr)
|
|
for key in [k for k in ledger if isinstance(k, str)]:
|
|
ledger[(scope, key)] = ledger.pop(key)
|
|
for attr in _KEYED_SETS:
|
|
keys = getattr(_core, attr)
|
|
bare = {k for k in keys if isinstance(k, str)}
|
|
keys.difference_update(bare)
|
|
keys.update((scope, k) for k in bare)
|
|
# The owner and the tools follow the key: status, config pruning and /reload-mcp select by
|
|
# owner == scope, and teardown deregisters from the owner's overlay, not the global map.
|
|
tool_names = []
|
|
for key in [k for k, owner in _core._server_scope_keys.items() if owner is None and _key_scope(k) == scope]:
|
|
_core._server_scope_keys[key] = scope
|
|
tool_names += getattr(_core._servers.get(key), "_registered_tool_names", None) or ()
|
|
for key in [k for k in _core._lazy_server_tool_names if _key_scope(k) == scope]:
|
|
tool_names += _core._lazy_server_tool_names[key]
|
|
for key in set(_core._servers) | set(_core._lazy_server_tool_names):
|
|
if _key_scope(key) == scope:
|
|
_core._server_tool_scopes.setdefault(key, set()).add(scope)
|
|
from tools.mcp_tool_costs import rekey_unscoped
|
|
rekey_unscoped(scope) # chats' spend totals follow their server's policy key
|
|
registry.move_to_scope(tool_names, scope)
|