"""Small fail-closed adapter for the optional :mod:`distinct_tools` package. Per-run limits travel in ``ToolSelection.config``. The server writes them; this module is the only thing that reads them, and it treats them as a *ceiling the run may lower*, never as a permission the run may raise. Every value is clamped against :data:`HARD_TOOL_LIMITS`, which the worker owns and no server can widen. The clamped, effective values are published on :attr:`ToolBrokerAdapter.limits` so a run log can state what was actually enforced instead of what was merely requested. """ from __future__ import annotations import importlib import math import os import re from collections.abc import Mapping from dataclasses import dataclass, field from types import MappingProxyType from typing import Any from distinct_protocol import JobSpec, ToolSelection _TOOL_REF_RE = re.compile(r"^([A-Za-z0-9][A-Za-z0-9_.-]{0,127})@([A-Za-z0-9][A-Za-z0-9_.-]{0,31})$") #: Worker-owned ceilings. A per-run config may lower any of these and may #: never raise one. ``max_calls`` matches ``StructuredToolHarness`` so the two #: budgets cannot drift apart silently. #: #: ``max_calls`` WAS 3, AND 3 WAS THE REAL CEILING WHEREVER ELSE IT WAS RAISED. #: The harness takes the lower of its own budget and the broker's, so a request #: needing six tools -- "plan the recipe book, convert the imperial quantities, #: scale one recipe, write the document, build the ingredient spreadsheet" -- #: could not complete at any level of model competence, on any machine, no #: matter what the job asked for. It is still a ceiling and still worker-owned; #: it is now high enough that a genuinely multi-step request can finish inside #: it. HARD_TOOL_LIMITS: Mapping[str, Any] = MappingProxyType( { "max_calls": 10, "timeout_seconds": 15.0, "max_input_bytes": 16_384, "max_output_bytes": 131_072, } ) _INTEGER_LIMITS = ("max_calls", "max_input_bytes", "max_output_bytes") class ToolBrokerUnavailable(RuntimeError): pass class ToolNotAllowed(PermissionError): pass class ToolConfigError(ToolBrokerUnavailable): """A per-run tool config was malformed, unknown, or tried to widen a limit.""" @dataclass(frozen=True) class ToolBrokerAdapter: """Validate the job allowlist before delegating to ``distinct_tools``.""" broker: Any allowed: Mapping[str, ToolSelection] limits: Mapping[str, Any] = field(default_factory=lambda: MappingProxyType({})) def tool_manifest(self) -> tuple[Mapping[str, Any], ...]: """Return only the exact, immutable tool schemas granted to the job.""" manifest = self.broker.tool_manifest() if not isinstance(manifest, tuple): manifest = tuple(manifest) return tuple(dict(item) for item in manifest) @property def max_calls(self) -> int: """The enforced job-wide call budget, after clamping.""" value = self.limits.get("max_calls", HARD_TOOL_LIMITS["max_calls"]) return int(value) def invoke( self, tool_id: str, arguments: Mapping[str, Any], *, version: str | None = None, ) -> Any: base_id, requested_version = _split_ref(tool_id) if version is not None: requested_version = version selection = self.allowed.get(base_id) if selection is None: raise ToolNotAllowed(f"tool {base_id!r} was not allowed for this job") if requested_version is not None and requested_version != selection.version: raise ToolNotAllowed( f"tool {base_id!r} version {requested_version!r} was not allowed" ) if not isinstance(arguments, Mapping): raise TypeError("tool arguments must be a mapping") return self.broker.invoke(base_id, selection.version, dict(arguments)) def _limit_error(message: str) -> ToolConfigError: return ToolConfigError(message) def _clamped_int(key: str, value: Any) -> int: ceiling = int(HARD_TOOL_LIMITS[key]) if isinstance(value, bool) or not isinstance(value, int): raise _limit_error(f"per-run tool limit {key!r} must be an integer") if value < 1: raise _limit_error(f"per-run tool limit {key!r} must be at least 1") return min(value, ceiling) def _clamped_timeout(value: Any) -> float: ceiling = float(HARD_TOOL_LIMITS["timeout_seconds"]) if isinstance(value, bool) or not isinstance(value, int | float): raise _limit_error("per-run tool limit 'timeout_seconds' must be a number") timeout = float(value) if not math.isfinite(timeout) or timeout <= 0: raise _limit_error("per-run tool limit 'timeout_seconds' must be positive and finite") return min(timeout, ceiling) def effective_tool_limits( selections: tuple[ToolSelection, ...], ) -> tuple[dict[str, Any], dict[str, int]]: """Fold every selection's ``config`` into one enforced policy. Returns ``(job_limits, per_tool_call_budgets)``. Unknown keys are refused rather than ignored: a limit this worker does not understand must never be silently dropped, because dropping it is always the permissive outcome. """ job_limits: dict[str, Any] = { "max_calls": int(HARD_TOOL_LIMITS["max_calls"]), "timeout_seconds": float(HARD_TOOL_LIMITS["timeout_seconds"]), "max_input_bytes": int(HARD_TOOL_LIMITS["max_input_bytes"]), "max_output_bytes": int(HARD_TOOL_LIMITS["max_output_bytes"]), } per_tool: dict[str, int] = {} requested_calls: list[int] = [] for selection in selections: config = selection.config or {} if not isinstance(config, Mapping): raise _limit_error("tool config must be a mapping") unknown = sorted(set(map(str, config)) - set(HARD_TOOL_LIMITS)) if unknown: raise _limit_error( f"tool {selection.id!r} config contains unsupported limits: {', '.join(unknown)}" ) tool_calls = int(HARD_TOOL_LIMITS["max_calls"]) if "max_calls" in config: tool_calls = _clamped_int("max_calls", config["max_calls"]) per_tool[selection.id] = tool_calls requested_calls.append(tool_calls) for key in ("max_input_bytes", "max_output_bytes"): if key in config: job_limits[key] = min(job_limits[key], _clamped_int(key, config[key])) if "timeout_seconds" in config: job_limits["timeout_seconds"] = min( job_limits["timeout_seconds"], _clamped_timeout(config["timeout_seconds"]) ) if requested_calls: job_limits["max_calls"] = min( int(HARD_TOOL_LIMITS["max_calls"]), sum(requested_calls) ) return job_limits, per_tool def load_tool_broker(job: JobSpec) -> ToolBrokerAdapter: """Create the optional broker, or reject rather than widening access.""" if not isinstance(job, JobSpec): raise TypeError("job must be a distinct_protocol.JobSpec") try: module = importlib.import_module("distinct_tools") default_registry = module.default_registry broker_type = module.ToolBroker policy_type = module.JobPolicy ref_type = module.ToolRef except (ImportError, AttributeError) as exc: raise ToolBrokerUnavailable("distinct_tools is not installed or compatible") from exc allowed = {selection.id: selection for selection in job.allowed_tools} job_limits, per_tool_calls = effective_tool_limits(job.allowed_tools) # One snapshot of the environment for both reads below. The registry and # the operator policy are derived from the same variables, and if they were # read at two different moments an operator editing their approvals between # the two would get a broker that refuses every tool in its own registry. environ = dict(os.environ) try: registry = default_registry(environ=environ) specs = {spec.ref: spec for spec in registry.specs()} except Exception as exc: raise ToolBrokerUnavailable( f"tool registry could not be built: {type(exc).__name__}" ) from exc refs = {} for selection in job.allowed_tools: ref = ref_type(selection.id, selection.version) if ref not in specs: # Fail closed and say which tool: with an empty registry this is # the expected path for every selection, and "unavailable" without # a name is not a usable diagnostic. raise ToolBrokerUnavailable( f"tool {selection.id}@{selection.version} is not installed on this agent" ) refs[selection.id] = ref # What the operator approved, read here rather than assumed. Two things # were wrong on this path and both made a documented gate vacuous. # # `allowed_hosts` was the union of the hosts declared by the very specs it # is meant to constrain, so the broker's `missing_hosts` refusal could not # fire: every tool was allowed exactly the hosts it asked for. It is now # the hosts of the tools the operator separately cleared for egress, which # is what a second gate means. # # `operator_policy` was not passed at all, so the `OperatorRefusal` check # documented in `ToolBroker` as "checked first: a server cannot reach a # tool the operator withheld" never ran in a real worker. The property # held only because `default_registry()` filters by the same policy when # it builds, which is one indirection away from failing open and is not # where the code says the check is. operator = None egress_hosts: frozenset = frozenset() try: operator = module.operator_policy(environ=environ) cleared = getattr(operator, "egress_approved", frozenset()) egress_hosts = frozenset( host for ref in refs.values() if ref in cleared for host in specs[ref].required_hosts ) except Exception as exc: raise ToolBrokerUnavailable( f"the operator's tool policy could not be read: {type(exc).__name__}" ) from exc try: allowed_refs = frozenset(refs.values()) policy = policy_type( allowed_tools=allowed_refs, allowed_hosts=egress_hosts, max_calls=job_limits["max_calls"], per_tool_quotas={refs[tool_id]: budget for tool_id, budget in per_tool_calls.items()}, timeout_seconds=job_limits["timeout_seconds"], max_input_bytes=job_limits["max_input_bytes"], max_output_bytes=job_limits["max_output_bytes"], ) broker = broker_type( registry=registry, policy=policy, job_id=job.id, operator_policy=operator, ) except Exception as exc: raise ToolBrokerUnavailable( f"tool broker rejected the job allowlist: {type(exc).__name__}" ) from exc if not callable(getattr(broker, "invoke", None)): raise ToolBrokerUnavailable("tool broker has no invoke method") return ToolBrokerAdapter( broker=broker, allowed=allowed, limits=MappingProxyType(dict(job_limits)), ) def advertised_tool_refs() -> tuple[str, ...]: """Return registry refs, or an empty tuple when tools are unavailable.""" try: module = importlib.import_module("distinct_tools") registry = module.default_registry() except Exception: return () try: return tuple( sorted(f"{spec.tool_id}@{spec.version}" for spec in registry.specs()) ) except Exception: return () def _split_ref(value: str) -> tuple[str, str | None]: if not isinstance(value, str): raise ToolNotAllowed("tool id must be text") if "@" not in value: return value, None match = _TOOL_REF_RE.fullmatch(value) if not match: raise ToolNotAllowed("tool reference is malformed") return match.group(1), match.group(2)