railtracks.prebuilt.middleware.pre_verifier

  1from __future__ import annotations
  2
  3import asyncio
  4import functools
  5from typing import Any, Awaitable, Callable, ParamSpec, TypeVar, overload
  6
  7from typing_extensions import Never
  8
  9from railtracks.events.middleware import (
 10    MiddlewareVerifierPreFailureEvent,
 11    MiddlewareVerifierPreInvocationEvent,
 12    MiddlewareVerifierPreResponseEvent,
 13)
 14from railtracks.events.send import emit
 15from railtracks.middleware.core import Middleware, _MiddlewareSignature, wrap_node
 16from railtracks.middleware.verdict import (
 17    Verdict,
 18    VerifierDecision,
 19    VerifierRejectedError,
 20)
 21from railtracks.utils.logging.create import get_rt_logger
 22from railtracks.utils.unpack import unpack_async_sync
 23
 24logger = get_rt_logger(__name__)
 25
 26_P = ParamSpec("_P")
 27_R = TypeVar("_R")
 28
 29_ApproveFn = Callable[_P, Verdict | Awaitable[Verdict]]
 30
 31
 32# A pre-verifier constrains the call arguments but leaves the node output free.
 33@overload
 34def pre_verifier(
 35    approve_fn: _ApproveFn[_P],
 36    /,
 37    *,
 38    timeout: float | None = None,
 39    name: str | None = None,
 40) -> Middleware[_P, Any, _MiddlewareSignature[_P, Never]]: ...
 41@overload
 42def pre_verifier(
 43    *, timeout: float | None = None, name: str | None = None
 44) -> Callable[
 45    [_ApproveFn[_P]], Middleware[_P, Any, _MiddlewareSignature[_P, Never]]
 46]: ...
 47
 48
 49def pre_verifier(
 50    approve_fn: _ApproveFn[_P] | None = None,
 51    /,
 52    *,
 53    timeout: float | None = None,
 54    name: str | None = None,
 55) -> (
 56    Middleware[_P, Any, _MiddlewareSignature[_P, Never]]
 57    | Callable[[_ApproveFn[_P]], Middleware[_P, Any, _MiddlewareSignature[_P, Never]]]
 58):
 59    """Build a node-verification middleware around ``approve_fn`` that gates a call
 60    BEFORE it runs.
 61
 62    ``approve_fn`` is called with the exact ``*args, **kwargs`` the wrapped
 63    node was called with — sync or async, both supported — and must return a
 64    `Verdict`. On decline, `VerifierRejectedError` is raised and the node's own
 65    body never runs. On accept, the call is forwarded onward, using the
 66    verdict's ``args``/``kwargs`` if it supplied overrides, otherwise the
 67    original ones unchanged.
 68
 69    If ``timeout`` is set and ``approve_fn`` doesn't respond in time, the call
 70    is treated as declined with reason ``"timeout"``.
 71
 72    See also :func:`~railtracks.prebuilt.middleware.post_verifier.post_verifier`,
 73    which gates a call's output AFTER it has already run. For the full
 74    picture (composing with other middleware, custom approval backends,
 75    guided walkthroughs), see the Verifiers docs.
 76    """
 77
 78    if approve_fn is None:
 79        return lambda fn: pre_verifier(fn, timeout=timeout, name=name)
 80
 81    return wrap_node(_wrapper(approve_fn, timeout), name=name)
 82
 83
 84def _wrapper(approve_fn: _ApproveFn[_P], timeout: float | None):
 85    @functools.wraps(approve_fn)
 86    async def wrapped(
 87        call: Callable[_P, Awaitable[_R]], *args: _P.args, **kwargs: _P.kwargs
 88    ) -> _R:
 89        await emit(MiddlewareVerifierPreInvocationEvent(args=args, kwargs=kwargs))
 90
 91        is_timeout = False
 92        try:
 93            review = unpack_async_sync(approve_fn(*args, **kwargs))
 94            if timeout is None:
 95                verdict = await review
 96            else:
 97                verdict = await asyncio.wait_for(review, timeout=timeout)
 98        except asyncio.TimeoutError:
 99            verdict = Verdict(accepted=False, comment="timeout")
100            is_timeout = True
101        except Exception as e:
102            await emit(MiddlewareVerifierPreFailureEvent.from_exception(e))
103            raise
104
105        new_args = verdict.args if verdict.args is not None else args
106        new_kwargs = verdict.kwargs if verdict.kwargs is not None else kwargs
107        overridden = verdict.args is not None or verdict.kwargs is not None
108
109        await emit(
110            MiddlewareVerifierPreResponseEvent(
111                decision=VerifierDecision.from_verdict(
112                    verdict, overridden=overridden, timeout=is_timeout
113                ),
114                args=new_args,
115                kwargs=new_kwargs,
116            )
117        )
118
119        if not verdict.accepted:
120            raise VerifierRejectedError(verdict.comment or "rejected")
121
122        if verdict.comment:
123            logger.info("pre_verifier accepted with comment: %s", verdict.comment)
124
125        return await call(*new_args, **new_kwargs)
126
127    return wrapped