Spaces:

User1342
/

distinct

Sleeping

Offload what the card can hold, survive a dropped connection, and stop sending prompts the model cannot read

95200a7 about 2 months ago

12.1 kB

"""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)

	"""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)