"""ToolProjection read-model and ToolProjectionRegistry. Provides a central, tenant-aware catalog of all backend tools across all mcp_servers. The registry is built from the discovery layer (tool schemas already stored in the domain model) and cached; the cache is invalidated on config reload so that status changes (active / withdrawn) propagate without a process restart. Thread-safe: uses RLock throughout. """ from __future__ import annotations import logging import threading from collections.abc import Mapping from dataclasses import dataclass, field from typing import TYPE_CHECKING, Literal from ...domain.services.digest_computation import compute_tool_digest from ...domain.value_objects.tool_digest import DigestEnforcement, ToolDigest if TYPE_CHECKING: from ...domain.model.tool_catalog import ToolSchema logger = logging.getLogger(__name__) class _AllTenants: """Sentinel type: a withdrawal entry that applies to ALL tenants.""" # Sentinel value: marks a withdrawal that applies to ALL tenants. # An overlay entry is therefore `set[str] | _AllTenants`; narrow with # `isinstance(entry, set)` to operate on the per-tenant set. _ALL_TENANTS = _AllTenants() # --------------------------------------------------------------------------- # Read-model # --------------------------------------------------------------------------- @dataclass(frozen=True) class ToolProjection: """Immutable read-model for a single backend tool. Attributes: mcp_server: Owning mcp_server identifier. tool: Tool name as reported by the backend. schema: Full JSON-Schema dict (name, description, inputSchema, …). digest: SEP-1766 SHA-256 fingerprint of the canonical schema. status: Base status – "active" (default) or "withdrawn". tenant_overrides: Per-tenant status overrides keyed by tenant_id. """ mcp_server: str tool: str schema: dict digest: ToolDigest status: Literal["active", "withdrawn"] = "active" tenant_overrides: Mapping[str, str] = field(default_factory=dict) def effective_status(self, tenant_id: str | None = None) -> str: """Return the status that applies for *tenant_id*. If *tenant_id* is ``None`` or has no override, the base ``status`` is returned. """ if tenant_id is not None and tenant_id in self.tenant_overrides: return self.tenant_overrides[tenant_id] return self.status def is_withdrawn_for(self, tenant_id: str | None = None) -> bool: """Return ``True`` when this tool is withdrawn for *tenant_id*.""" return self.effective_status(tenant_id) == "withdrawn" # --------------------------------------------------------------------------- # Registry # --------------------------------------------------------------------------- class ToolProjectionRegistry: """Central, cached, tenant-aware catalog of all backend tools. Built from the domain repository (tool schemas already discovered); cached for the lifetime of a config epoch; invalidated on config reload via :meth:`invalidate`. Thread-safe: uses RLock. """ def __init__(self) -> None: self._lock = threading.RLock() # Primary store: (mcp_server, tool) -> ToolProjection self._projections: dict[tuple[str, str], ToolProjection] = {} # Tracks whether the registry has been populated at least once self._built: bool = False # Every server whose catalogue has been projected at least once in this # process, an empty catalogue included (#1446). Never shrinks: a stop # does not remove a projection, so an idle server still counts. self._projected: set[str] = set() # Config-withdrawal overlay: (mcp_server, kind, name) -> set of tenant_ids # or _ALL_TENANTS sentinel. Populated at config-load time; re-applied on # every reload. _ALL_TENANTS means withdrawn for every tenant. # # `kind` is "tool", "prompt" or "resource" (#1028): a prompt and a # resource are withdrawable exactly the way a tool is, and the overlay # they share is the same one the reload/runtime split is already correct # for. Only tools carry a discovered projection (schema + digest) -- # prompts and resources are relayed live -- so `_projections` below stays # tool-keyed and the kinds meet again at `is_withdrawn`. self._config_withdrawals: dict[tuple[str, str, str], set[str] | _AllTenants] = {} # Runtime-withdrawal overlay: survives config reloads (clear_config_withdrawals does NOT touch this). # Same shape as _config_withdrawals. self._runtime_withdrawals: dict[tuple[str, str, str], set[str] | _AllTenants] = {} # Config-pin overlay: (mcp_server, tool) -> {tenant_id -> pinned ToolDigest}. # Populated at config-load time; re-applied on every reload (#233). self._config_pins: dict[tuple[str, str], dict[str, ToolDigest]] = {} # All-tenants config-pin overlay: (mcp_server, tool) -> pinned ToolDigest. # The counterpart of the `_ALL_TENANTS` withdrawal sentinel, and the only # pin an unauthenticated caller can be held to: `resolve_pin` keys the # per-tenant map by a tenant id, so with auth off -- where the caller is # anonymous and `tenant_id` is always None -- every per-tenant pin missed # and the digest capability enforced nothing at all (#902). self._config_pins_all_tenants: dict[tuple[str, str], ToolDigest] = {} # Per-mcp_server digest-enforcement mode for pin mismatches; an unset # server defaults to the strictest (block). Scoped per server so one # server's `audit` cannot downgrade another server's pins (#278). self._digest_enforcement: dict[str, DigestEnforcement] = {} # ------------------------------------------------------------------ # Population (called by bootstrap / config-reload) # ------------------------------------------------------------------ def build_from_tools( self, mcp_server: str, tools: list[ToolSchema], *, status_overrides: Mapping[str, Literal["active", "withdrawn"]] | None = None, tenant_overrides: Mapping[str, Mapping[str, str]] | None = None, ) -> None: """Populate projections for *mcp_server* from its discovered tools. Existing projections for *mcp_server* are replaced atomically. Args: mcp_server: Owning mcp_server identifier. tools: Discovered :class:`~domain.model.tool_catalog.ToolSchema` list. status_overrides: Optional per-tool base-status overrides, keyed by tool name. Values must be "active" or "withdrawn". tenant_overrides: Optional per-tool tenant overrides, keyed by tool name then by tenant_id. """ status_overrides = status_overrides or {} tenant_overrides = tenant_overrides or {} new_projections: dict[tuple[str, str], ToolProjection] = {} for tool_schema in tools: tool_dict = tool_schema.to_dict() digest = compute_tool_digest(tool_dict) base_status: Literal["active", "withdrawn"] = status_overrides.get(tool_schema.name, "active") per_tenant = dict(tenant_overrides.get(tool_schema.name, {})) projection = ToolProjection( mcp_server=mcp_server, tool=tool_schema.name, schema=tool_dict, digest=digest, status=base_status, tenant_overrides=per_tenant, ) new_projections[(mcp_server, tool_schema.name)] = projection with self._lock: # Remove stale entries for this server, add new ones stale_keys = [k for k in self._projections if k[0] == mcp_server] for k in stale_keys: del self._projections[k] self._projections.update(new_projections) self._built = True self._projected.add(mcp_server) logger.debug( "tool_projection_registry_built", extra={ "mcp_server": mcp_server, "tool_count": len(new_projections), }, ) # ------------------------------------------------------------------ # Config-withdrawal overlay (populated at config-load time) # ------------------------------------------------------------------ def set_config_withdrawal( self, mcp_server: str, tool: str, tenant_id: str | None = None, *, kind: str = "tool", ) -> None: """Mark a tool, prompt or resource as withdrawn via config. Args: mcp_server: Owning mcp_server identifier. tool: Tool name, prompt name, or upstream resource URI. tenant_id: If ``None``, it is withdrawn for ALL tenants. Otherwise only for the given tenant. kind: What *tool* names -- "tool" (default), "prompt" or "resource". """ key = (mcp_server, kind, tool) with self._lock: current = self._config_withdrawals.get(key) if tenant_id is None: # ALL-tenants sentinel — overrides any per-tenant set. self._config_withdrawals[key] = _ALL_TENANTS elif isinstance(current, set): current.add(tenant_id) elif current is None: self._config_withdrawals[key] = {tenant_id} # else: current is _ALL_TENANTS — already the broader rule, keep it. logger.debug( "config_withdrawal_set", extra={"mcp_server": mcp_server, "tool": tool, "tenant_id": tenant_id}, ) def clear_config_withdrawals(self) -> None: """Remove all config-declared withdrawals. Called before re-applying config on reload so that removing a withdrawal from the config file actually restores the tool. IMPORTANT: Does NOT touch ``_runtime_withdrawals`` — runtime withdrawals are intentionally reload-safe (issue #235). """ with self._lock: self._config_withdrawals.clear() logger.debug("config_withdrawals_cleared") def _is_config_withdrawn_for(self, mcp_server: str, tool: str, tenant_id: str | None, kind: str = "tool") -> bool: """Return True if (mcp_server, kind, tool) is config-withdrawn for tenant_id.""" entry = self._config_withdrawals.get((mcp_server, kind, tool)) if entry is None: return False if entry is _ALL_TENANTS: return True # entry is a set[str] return tenant_id is not None and tenant_id in entry # type: ignore[operator] # ------------------------------------------------------------------ # Config-pin overlay (populated at config-load time) # ------------------------------------------------------------------ def set_config_pin( self, mcp_server: str, tool: str, tenant_id: str | None, digest: ToolDigest, ) -> None: """Pin a tool to a specific digest via config. Args: mcp_server: Owning mcp_server identifier. tool: Tool name. tenant_id: Tenant the pin applies to. ``None`` pins the tool for ALL tenants, including callers that carry no tenant identity -- the same meaning ``None`` has in :meth:`set_config_withdrawal`. digest: The :class:`~domain.value_objects.tool_digest.ToolDigest` the tool is expected to match. """ with self._lock: if tenant_id is None: self._config_pins_all_tenants[(mcp_server, tool)] = digest else: self._config_pins.setdefault((mcp_server, tool), {})[tenant_id] = digest logger.debug( "config_pin_set", extra={"mcp_server": mcp_server, "tool": tool, "tenant_id": tenant_id}, ) def set_digest_enforcement(self, mcp_server: str, mode: DigestEnforcement) -> None: """Set the digest-enforcement mode applied to *mcp_server*'s pin mismatches.""" with self._lock: self._digest_enforcement[mcp_server] = mode logger.debug("digest_enforcement_set", extra={"mcp_server": mcp_server, "mode": mode.value}) def resolve_pin(self, mcp_server: str, tool: str, tenant_id: str | None) -> ToolDigest | None: """Return the pinned digest for *(mcp_server, tool)* and *tenant_id*. Resolution is narrowest-first, the same order the tool-access policies use: a pin declared for this tenant wins over one declared for all tenants. An all-tenants pin applies to a caller with no tenant identity -- that is the whole point of it, and it is why this no longer returns ``None`` on sight of a ``None`` tenant (#902). A caller with no identity still matches no per-tenant pin, which is correct: those pins name a tenant this caller has not been shown to be. Returns ``None`` only when no pin covers the tool for this caller. """ key = (mcp_server, tool) with self._lock: if tenant_id is not None: tenant_pin = self._config_pins.get(key, {}).get(tenant_id) if tenant_pin is not None: return tenant_pin return self._config_pins_all_tenants.get(key) def digest_enforcement(self, mcp_server: str) -> DigestEnforcement: """Return the digest-enforcement mode for *mcp_server* (block if unset).""" with self._lock: return self._digest_enforcement.get(mcp_server, DigestEnforcement.BLOCK) def clear_config_pins(self) -> None: """Remove all config-declared pins and reset enforcement to block. Called before re-applying config on reload so that removing a pin (or the ``digest_enforcement`` setting) from the config file actually reverts to the strict default (#233). """ with self._lock: self._config_pins.clear() self._config_pins_all_tenants.clear() self._digest_enforcement.clear() logger.debug("config_pins_cleared") def adopt_config_overlays(self, staged: ToolProjectionRegistry, *, replace: bool) -> None: """Take the withdrawals, pins and enforcement modes a configuration registered on *staged*. The configuration registers them on a fresh registry, and they are taken from it here under this registry's lock. A concurrent resolve sees the previous overlays or the new ones, never the empty ones a reload used to leave between clearing and re-registering (#1424). With *replace*, which is what a reload asks for, the config overlays become exactly *staged*'s, so deleting a withdrawal or a pin from the file restores the tool and the strict default (#233). Without it they are added to what is there, which is what a first load does. Runtime withdrawals are never touched either way (#235). """ with staged._lock: withdrawals = { key: set(entry) if isinstance(entry, set) else entry for key, entry in staged._config_withdrawals.items() } pins = {key: dict(by_tenant) for key, by_tenant in staged._config_pins.items()} pins_all_tenants = dict(staged._config_pins_all_tenants) enforcement = dict(staged._digest_enforcement) with self._lock: if replace: self._config_withdrawals = withdrawals self._config_pins = pins self._config_pins_all_tenants = pins_all_tenants self._digest_enforcement = enforcement return # Added through the setters, so a merge means what registering the # same entries one by one has always meant. for (mcp_server, kind, name), entry in withdrawals.items(): for tenant_id in sorted(entry) if isinstance(entry, set) else [None]: self.set_config_withdrawal(mcp_server, name, tenant_id, kind=kind) for (mcp_server, tool), by_tenant in pins.items(): for tenant_id, digest in by_tenant.items(): self.set_config_pin(mcp_server, tool, tenant_id, digest) for (mcp_server, tool), digest in pins_all_tenants.items(): self.set_config_pin(mcp_server, tool, None, digest) self._digest_enforcement.update(enforcement) # ------------------------------------------------------------------ # Runtime-withdrawal overlay (survives config reloads) # ------------------------------------------------------------------ def withdraw( self, mcp_server: str, tool: str, tenant_id: str | None = None, *, kind: str = "tool", ) -> None: """Mark a tool, prompt or resource as withdrawn at runtime (survives config reload). Args: mcp_server: Owning mcp_server identifier. tool: Tool name, prompt name, or upstream resource URI. tenant_id: If ``None``, it is withdrawn for ALL tenants. Otherwise only for the given tenant. kind: What *tool* names -- "tool" (default), "prompt" or "resource". """ key = (mcp_server, kind, tool) with self._lock: current = self._runtime_withdrawals.get(key) if tenant_id is None: self._runtime_withdrawals[key] = _ALL_TENANTS elif isinstance(current, set): current.add(tenant_id) elif current is None: self._runtime_withdrawals[key] = {tenant_id} # else: current is _ALL_TENANTS — already covers all tenants, keep it. logger.debug( "runtime_withdrawal_set", extra={"mcp_server": mcp_server, "tool": tool, "tenant_id": tenant_id}, ) def restore( self, mcp_server: str, tool: str, tenant_id: str | None = None, *, kind: str = "tool", ) -> None: """Remove a runtime withdrawal for a tool, prompt or resource. Affects ONLY the runtime overlay; a config-declared withdrawal independently persists (effective = config OR runtime). Args: mcp_server: Owning mcp_server identifier. tool: Tool name, prompt name, or upstream resource URI. tenant_id: If ``None``, removes the runtime withdrawal for ALL tenants (clears the entire key). Otherwise removes the given tenant from the per-tenant set. kind: What *tool* names -- "tool" (default), "prompt" or "resource". """ key = (mcp_server, kind, tool) with self._lock: current = self._runtime_withdrawals.get(key) if current is None: return # Nothing to restore. if tenant_id is None: # Remove the entire entry → no runtime withdrawal remains. del self._runtime_withdrawals[key] elif isinstance(current, set): current.discard(tenant_id) if not current: del self._runtime_withdrawals[key] # else: current is _ALL_TENANTS — can't partially remove from an # ALL-tenants entry; do nothing to avoid re-enabling other tenants. logger.debug( "runtime_withdrawal_restored", extra={"mcp_server": mcp_server, "tool": tool, "tenant_id": tenant_id}, ) def is_withdrawn_for_all_tenants(self, mcp_server: str, name: str, *, kind: str = "tool") -> bool: """Is there a runtime withdrawal of *name* that covers every tenant? The one entry a per-tenant :meth:`restore` leaves in place. Config withdrawals are not consulted, because restore never touches them either. """ with self._lock: return self._runtime_withdrawals.get((mcp_server, kind, name)) is _ALL_TENANTS def _is_runtime_withdrawn_for(self, mcp_server: str, tool: str, tenant_id: str | None, kind: str = "tool") -> bool: """Return True if (mcp_server, kind, tool) is runtime-withdrawn for tenant_id.""" entry = self._runtime_withdrawals.get((mcp_server, kind, tool)) if entry is None: return False if entry is _ALL_TENANTS: return True return tenant_id is not None and tenant_id in entry # type: ignore[operator] def _is_withdrawn_for(self, mcp_server: str, tool: str, tenant_id: str | None, kind: str = "tool") -> bool: """Return True if config OR runtime withdraws (mcp_server, kind, tool) for tenant_id.""" return self._is_config_withdrawn_for(mcp_server, tool, tenant_id, kind) or self._is_runtime_withdrawn_for( mcp_server, tool, tenant_id, kind ) def is_withdrawn( self, mcp_server: str, name: str, *, kind: str = "tool", tenant_id: str | None = None, ) -> bool: """Is *name* withdrawn for *tenant_id* by config or at runtime? The public form of the overlay, for the surfaces that have no projection to read a status off: a prompt and a resource are relayed live, so :meth:`resolve` -- which returns a schema and a digest -- has nothing to say about them. Tools keep going through :meth:`resolve`, whose answer also folds in the base status set at discovery time. """ with self._lock: return self._is_withdrawn_for(mcp_server, name, tenant_id, kind) # ------------------------------------------------------------------ # Query API (read-only) # ------------------------------------------------------------------ def resolve( self, mcp_server: str, tool: str, tenant_id: str | None = None, ) -> ToolProjection | None: """Return the :class:`ToolProjection` for *(mcp_server, tool)*. Consults the config-withdrawal overlay first. If the tool is config-withdrawn for *tenant_id* (or for ALL tenants), a projection marked ``withdrawn`` is returned even when the tool has not yet been discovered (no ``build_from_tools`` call). A placeholder digest with all-zero hex is used for undiscovered tools — it is valid per the :class:`~domain.value_objects.tool_digest.ToolDigest` schema (64 hex chars) and carries no semantic meaning. Returns ``None`` only when the tool is completely unknown (not in the discovered store AND not config-withdrawn for this tenant). Args: mcp_server: Owning mcp_server identifier. tool: Tool name. tenant_id: Optional tenant identifier (informational — the full projection is returned regardless so callers can re-check for other tenants without a second lookup). Returns: The matching :class:`ToolProjection`, or ``None``. """ with self._lock: discovered = self._projections.get((mcp_server, tool)) withdrawn = self._is_withdrawn_for(mcp_server, tool, tenant_id) if not withdrawn: # Neither config nor runtime overlay applies — return discovered as-is. return discovered # At least one overlay applies: build a withdrawn projection. # Collect ALL tenants withdrawn by either overlay for per-tenant synthesis. config_entry = self._config_withdrawals.get((mcp_server, "tool", tool)) runtime_entry = self._runtime_withdrawals.get((mcp_server, "tool", tool)) # Is it a blanket (ALL-tenants) withdrawal from either source? all_tenants_withdrawn = config_entry is _ALL_TENANTS or runtime_entry is _ALL_TENANTS if discovered is not None: # Discovered projection exists — augment it so is_withdrawn_for() fires. if all_tenants_withdrawn: return ToolProjection( mcp_server=discovered.mcp_server, tool=discovered.tool, schema=discovered.schema, digest=discovered.digest, status="withdrawn", tenant_overrides=discovered.tenant_overrides, ) else: # Merge per-tenant sets from both overlays. merged_overrides = dict(discovered.tenant_overrides) for entry in (config_entry, runtime_entry): if isinstance(entry, set): for tid in entry: merged_overrides[tid] = "withdrawn" return ToolProjection( mcp_server=discovered.mcp_server, tool=discovered.tool, schema=discovered.schema, digest=discovered.digest, status=discovered.status, tenant_overrides=merged_overrides, ) # Tool not yet discovered — synthesize a minimal withdrawn projection. # Placeholder digest: 64 zeros (valid hex, no semantic meaning). placeholder_digest = ToolDigest(tool_name=tool, sha256="0" * 64) if all_tenants_withdrawn: return ToolProjection( mcp_server=mcp_server, tool=tool, schema={}, digest=placeholder_digest, status="withdrawn", ) else: # Collect union of per-tenant entries from both overlays. merged_tenants: set[str] = set() for entry in (config_entry, runtime_entry): if entry is not None and entry is not _ALL_TENANTS: merged_tenants.update(entry) # type: ignore[arg-type] overrides = dict.fromkeys(merged_tenants, "withdrawn") return ToolProjection( mcp_server=mcp_server, tool=tool, schema={}, digest=placeholder_digest, status="active", tenant_overrides=overrides, ) def list_for_server(self, mcp_server: str) -> list[ToolProjection]: """Return all projections for *mcp_server* (snapshot).""" with self._lock: return [p for (s, _), p in self._projections.items() if s == mcp_server] def all(self) -> list[ToolProjection]: """Return a snapshot of all projections across all servers.""" with self._lock: return list(self._projections.values()) # ------------------------------------------------------------------ # Cache invalidation # ------------------------------------------------------------------ def invalidate(self) -> None: """Discard all cached projections, overlays, and pin state. Full reset — intended for testing only. In production, config reload calls :meth:`clear_config_withdrawals` (which preserves runtime withdrawals) and :meth:`clear_config_pins` rather than this method. """ with self._lock: self._projections.clear() self._config_withdrawals.clear() self._runtime_withdrawals.clear() self._config_pins.clear() self._config_pins_all_tenants.clear() self._digest_enforcement.clear() self._built = False self._projected.clear() logger.debug("tool_projection_registry_invalidated") def was_projected(self, mcp_server: str) -> bool: """Whether *mcp_server*'s catalogue has been projected at least once on this replica. True from its first ``McpServerStarted`` handled here, even when it discovered no tools, and it stays true when the server stops: this is what a front-door replica's readiness waits for (#1446), and what the catalogue retry never restarts a server after. A peer's start does not count, since the handler that projects is a local view (#922). """ with self._lock: return mcp_server in self._projected @property def is_built(self) -> bool: """``True`` if the registry has been populated at least once.""" with self._lock: return self._built # --------------------------------------------------------------------------- # Singleton # --------------------------------------------------------------------------- _registry: ToolProjectionRegistry | None = None _registry_lock = threading.Lock() def get_tool_projection_registry() -> ToolProjectionRegistry: """Return the process-global :class:`ToolProjectionRegistry` singleton.""" global _registry if _registry is None: with _registry_lock: if _registry is None: _registry = ToolProjectionRegistry() return _registry def reset_tool_projection_registry() -> None: """Reset the singleton (useful for testing).""" global _registry with _registry_lock: if _registry is not None: _registry.invalidate() _registry = None