API Reference¶

class BudgetContext:
    """Convenient wrapper around DeadlineBudget with per-call timeout caps.

    Encapsulates deadline budget and per-call timeout configuration for cleaner orchestration code.

    Example:
        # Define caps for each specific call in the operation
        call_caps = {
            "identity_create_user": 3.0,
            "verification_verify_code": 2.0,
            "credential_set_password": 3.0,
        }
        ctx = BudgetContext.create(total_seconds=10.0, call_caps=call_caps)

        # Get timeout for specific call (with configured cap)
        timeout = ctx.timeout_for_call("identity_create_user")
        await identity_client.create_user(timeout=timeout)

        # Get timeout for unconfigured call (uses remaining budget)
        timeout = ctx.timeout_for_call("identity_get_user")
        await identity_client.get_user(timeout=timeout)
    """

    def __init__(self, budget: DeadlineBudget, call_caps: dict[str, float]) -> None:
        """Initialize context with a budget and per-call caps.

        Args:
            budget: The underlying DeadlineBudget instance.
            call_caps: Mapping of call names to their timeout caps in seconds.
                If a call is not in this dict, it will use remaining budget without cap.
                Copied, so the context is not tied to the lifetime of the mapping given.
        """
        self._budget: DeadlineBudget = budget
        self._call_caps: dict[str, float] = dict(call_caps)

    @classmethod
    def create(
        cls,
        total_seconds: float,
        call_caps: dict[str, float] | None = None,
        *,
        min_timeout: float = 0.1,
        safety_margin: float = 0.0,
    ) -> BudgetContext:
        """Create a new budget context with the given parameters.

        Args:
            total_seconds: Total budget for this request in seconds.
            call_caps: Mapping of call names to their timeout caps (optional).
            min_timeout: Minimum timeout to return (default 0.1s).
            safety_margin: Safety margin reserved at the end (default 0.0s).

        Returns:
            A new BudgetContext instance.
        """
        budget = DeadlineBudget(
            total_seconds=total_seconds,
            min_timeout=min_timeout,
            safety_margin=safety_margin,
        )
        return cls(budget=budget, call_caps=call_caps or {})

    def timeout_for_call(self, call_name: str, reserve_for_next: float = 0.0) -> float:
        """Compute timeout for the next downstream call.

        If call_name has a configured cap, applies it. Otherwise uses remaining budget.

        Args:
            call_name: Name of the call (e.g., "identity_create_user", "verification_verify_code").
            reserve_for_next: Reserve this many seconds for subsequent steps.

        Returns:
            Computed timeout in seconds, from DeadlineBudget.timeout_for with the configured cap.
            May exceed the remaining budget when min_timeout does.

        Raises:
            DeadlineExceededError: If remaining budget is already exhausted.
        """
        cap = self._call_caps.get(call_name)
        return self._budget.timeout_for(cap=cap, reserve_for_next=reserve_for_next)

    def check_expired(self) -> None:
        """Raise DeadlineExceededError if budget is exhausted.

        Raises:
            DeadlineExceededError: If the deadline has passed.
        """
        self._budget.check_expired()

    def remaining(self) -> float:
        """Return remaining time budget in seconds.

        Returns negative value if deadline already exceeded.
        """
        return self._budget.remaining()

    def elapsed(self) -> float:
        """Return elapsed time since budget start in seconds."""
        return self._budget.elapsed()

    def expired(self) -> bool:
        """Check if budget is exhausted."""
        return self._budget.expired()

    @property
    def budget(self) -> DeadlineBudget:
        """Access to underlying budget for advanced use cases."""
        return self._budget

    @property
    def call_caps(self) -> dict[str, float]:
        """Access to this context's own copy of the configured call caps."""
        return self._call_caps

class DeadlineBudget:
    """Tracks request deadline budget using monotonic time.

    Example:
        budget = DeadlineBudget(total_seconds=9.5)
        timeout1 = budget.timeout_for(cap=5.0)
        await some_call(timeout=timeout1)
        timeout2 = budget.timeout_for(cap=5.0)
        await another_call(timeout=timeout2)
        if budget.expired():
            raise DeadlineExceededError(...)
    """

    def __init__(
        self,
        total_seconds: float,
        *,
        min_timeout: float = 0.1,
        safety_margin: float = 0.0,
    ) -> None:
        """Initialize deadline budget.

        Args:
            total_seconds: Total budget for this request in seconds.
            min_timeout: Minimum timeout to return from timeout_for (default 0.1s).
            safety_margin: Safety margin reserved at the end (subtracted from total).
        """
        if total_seconds <= 0:
            raise ValueError("total_seconds must be positive")
        if min_timeout < 0:
            raise ValueError("min_timeout must be non-negative")
        if safety_margin < 0:
            raise ValueError("safety_margin must be non-negative")
        if safety_margin >= total_seconds:
            raise ValueError("safety_margin must be less than total_seconds")

        self._total_seconds = total_seconds - safety_margin
        self._min_timeout = min_timeout
        self._started_at = time.monotonic()

    def remaining(self) -> float:
        """Return remaining time budget in seconds.

        Returns negative value if deadline already exceeded.
        """
        elapsed = time.monotonic() - self._started_at
        return self._total_seconds - elapsed

    def elapsed(self) -> float:
        """Return elapsed time since budget start in seconds."""
        return time.monotonic() - self._started_at

    def expired(self) -> bool:
        """Check if budget is exhausted."""
        return self.remaining() <= 0

    def timeout_for(
        self,
        cap: float | None = None,
        min_timeout: float | None = None,
        reserve_for_next: float = 0.0,
    ) -> float:
        """Compute timeout for the next downstream call.

        The remaining budget, less reserve_for_next, raised to min_timeout and then limited by cap.
        Both bounds outrank what is left: min_timeout is returned even when the budget holds less
        than that, and cap is applied last, so a cap below min_timeout wins over it.

        Args:
            cap: Maximum allowed timeout for this call (service-level cap).
            min_timeout: Minimum timeout override (default: use budget min_timeout).
            reserve_for_next: Reserve this many seconds for subsequent steps.

        Returns:
            Computed timeout in seconds. May exceed the remaining budget when min_timeout does.

        Raises:
            DeadlineExceededError: If remaining budget is already exhausted.
        """
        remaining = self.remaining()
        if remaining <= 0:
            raise DeadlineExceededError(
                budget_seconds=self._total_seconds,
                elapsed_seconds=self.elapsed(),
            )

        effective_min = min_timeout if min_timeout is not None else self._min_timeout
        available = max(remaining - reserve_for_next, effective_min)

        if cap is not None:
            return min(available, cap)
        return available

    def check_expired(self) -> None:
        """Raise DeadlineExceededError if budget is exhausted."""
        if self.expired():
            raise DeadlineExceededError(
                budget_seconds=self._total_seconds,
                elapsed_seconds=self.elapsed(),
            )

    @property
    def total_seconds(self) -> float:
        """Return total budget in seconds."""
        return self._total_seconds

class DeadlineExceededError(Exception):
    """Raised when request deadline budget is exhausted.

    Initialises ``Exception`` directly rather than via ``super()``, so it is safe to use as the first base
    of a class that also inherits from an exception with an incompatible ``__init__`` signature.
    Round-trips through ``pickle`` and ``copy`` with both attributes intact.

    Attributes:
        budget_seconds: Usable budget in seconds (total minus safety margin).
        elapsed_seconds: Time elapsed since the budget started, in seconds.
    """

    def __init__(self, budget_seconds: float, elapsed_seconds: float) -> None:
        self.budget_seconds = budget_seconds
        self.elapsed_seconds = elapsed_seconds
        # Not super(): under multiple inheritance super() follows the instance MRO and would hand this
        # message to a sibling base's __init__ instead of Exception's.
        Exception.__init__(
            self,
            f"Request deadline exceeded: {elapsed_seconds:.2f}s elapsed, budget was {budget_seconds:.2f}s",
        )

    def __reduce__(self) -> tuple[type[DeadlineExceededError], tuple[float, float], dict[str, Any]]:
        # BaseException pickles self.args, which holds the message rather than the constructor arguments.
        return type(self), (self.budget_seconds, self.elapsed_seconds), self.__dict__