refactor: fold /simplify-code findings

- Single source for the approval-derived bound: public human_wait_ceiling()
  in tools/approval.py; the gate's lock-timeout helper delegates to it
  instead of re-deriving timeout + margin (was duplicated in two modules
  and reached for a private _get_approval_timeout).
- Shared _clamped_window_seconds() for the close-time accrual and the
  open-window read, so the two clamps are identical by construction.
- Gate __init__ grows session_key kwarg; tests construct via the real
  constructor instead of mutating privates post-hoc.
- Gateway test resolves its pending approval via resolve_gateway_approval()
  (the production /deny path) instead of hand-rolling queue-entry internals.
- Docstring accuracy: human_wait_seconds monotonicity caveat under cap
  eviction; s/pre_tool_block/pre_tool_call/ hook name.
This commit is contained in:
kshitij
2026-08-06 16:59:50 +05:30
committed by kshitij
parent 10fb01e725
commit ea0d54db1d
3 changed files with 65 additions and 51 deletions

View File

@@ -114,16 +114,18 @@ _AUTHORIZATION_GATE_LOCK_TIMEOUT_S = 360.0
def _authorization_gate_lock_timeout() -> float:
"""Bound for the authorization serialization lock: approval timeout + margin.
Long enough that serialization is never broken while a legitimate approval
prompt is still answerable; short enough that a wedged holder (hanging
``pre_tool_call`` plugin, dead approval client) cannot park other workers
forever (#79719). Resolved once per gate (per batch), so a mid-process
``approvals.timeout`` change applies from the next batch.
Delegates to ``tools.approval.human_wait_ceiling`` — the same bound that
clamps a human-wait window's deadline contribution — so the two can't
drift. Long enough that serialization is never broken while a legitimate
approval prompt is still answerable; short enough that a wedged holder
(hanging ``pre_tool_call`` plugin, dead approval client) cannot park other
workers forever (#79719). Resolved once per gate (per batch), so a
mid-process ``approvals.timeout`` change applies from the next batch.
"""
try:
from tools.approval import HUMAN_WAIT_MARGIN_S, _get_approval_timeout
from tools.approval import human_wait_ceiling
return float(_get_approval_timeout()) + HUMAN_WAIT_MARGIN_S
return human_wait_ceiling()
except Exception:
return _AUTHORIZATION_GATE_LOCK_TIMEOUT_S
@@ -401,27 +403,33 @@ class _ConcurrentToolAuthorizationGate:
(which can legitimately exceed any fixed bound) is still excluded in full.
"""
def __init__(self, *, lock_timeout: float | None = None) -> None:
def __init__(
self,
*,
lock_timeout: float | None = None,
session_key: str | None = None,
) -> None:
self._serialization_lock = threading.Lock()
self._lock_timeout = (
_authorization_gate_lock_timeout()
if lock_timeout is None
else lock_timeout
)
self._session_key: str | None = None
try:
from tools.approval import get_current_session_key
self._session_key = session_key
if self._session_key is None:
try:
from tools.approval import get_current_session_key
# Snapshot the batch's session identity on the SUBMITTING thread:
# excluded_seconds() is polled from the batch wait loop, whose
# context may differ from the workers'.
self._session_key = get_current_session_key()
except Exception:
logger.debug(
"authorization gate could not snapshot the session key; "
"human-wait exclusion will re-resolve it at poll time",
exc_info=True,
)
# Snapshot the batch's session identity on the SUBMITTING
# thread: excluded_seconds() is polled from the batch wait
# loop, whose context may differ from the workers'.
self._session_key = get_current_session_key()
except Exception:
logger.debug(
"authorization gate could not snapshot the session key; "
"human-wait exclusion will re-resolve it at poll time",
exc_info=True,
)
self._baseline_wait_seconds = self._human_wait_seconds()
def _human_wait_seconds(self) -> float:

View File

@@ -41,12 +41,9 @@ SESSION = "test-session-79719"
def _make_gate(**kwargs) -> _ConcurrentToolAuthorizationGate:
gate = _ConcurrentToolAuthorizationGate(**kwargs)
# Pin the session key so contextvar/env noise from other tests can't
# change which wait state the gate reads.
gate._session_key = SESSION
gate._baseline_wait_seconds = gate._human_wait_seconds()
return gate
return _ConcurrentToolAuthorizationGate(session_key=SESSION, **kwargs)
class TestHumanWaitTracker:
@@ -293,11 +290,8 @@ class TestApprovalPathsRecordHumanWait:
try:
assert approval_mod.human_wait_seconds(SESSION) > 0.0
finally:
# Resolve the pending entry so the worker unwinds.
with approval_mod._lock:
for entry in approval_mod._gateway_queues.get(SESSION, []):
entry.result = "deny"
entry.event.set()
# Resolve the pending entry via the real production path.
approval_mod.resolve_gateway_approval(SESSION, "deny", resolve_all=True)
t.join(timeout=5)
assert not t.is_alive()
# Window closed once the wait resolved.

View File

@@ -2218,7 +2218,7 @@ _permanent_approved: set = set()
# tool batch deadline in agent/tool_executor.py excludes this time so a slow
# human answer never times a batch out — but ONLY this time. Measuring human
# waits at the source (rather than residency in the authorization gate, which
# is arbitrary code) is what keeps a wedged pre_tool_block plugin or a dead
# is arbitrary code) is what keeps a wedged pre_tool_call plugin or a dead
# approval client from growing the exclusion 1:1 with wall clock and defeating
# the deadline entirely (#79719).
#
@@ -2246,18 +2246,30 @@ _HUMAN_WAIT_MAX_SESSIONS = 256
HUMAN_WAIT_MARGIN_S = 60.0
def _human_wait_ceiling() -> float:
def human_wait_ceiling() -> float:
"""Max seconds a single window may contribute: approvals.timeout + margin.
Every legitimate human wait self-terminates at ``approvals.timeout`` (the
CLI prompt join and the gateway poll loop both enforce it), so a window
that overstays this ceiling is itself wedged and must not keep extending
a batch deadline. Never call while holding ``_human_wait_lock`` — it
reads the config cache.
a batch deadline. Also used by agent/tool_executor.py as the bound on the
authorization gate's serialization-lock acquire, so the two bounds cannot
drift. Never call while holding ``_human_wait_lock`` — it reads the
config cache.
"""
return float(_get_approval_timeout()) + HUMAN_WAIT_MARGIN_S
def _clamped_window_seconds(started: float, now: float, ceiling: float) -> float:
"""Seconds an open window contributes: elapsed, floored at 0, capped.
Shared by the close-time accrual in :func:`human_wait_window` and the
open-window read in :func:`human_wait_seconds` so the two clamps stay
identical by construction.
"""
return min(max(0.0, now - started), ceiling)
def _human_wait_state(session_key: str) -> _HumanWaitState:
"""Return (creating if needed) the wait state for *session_key*.
@@ -2303,19 +2315,18 @@ def human_wait_window(session_key: str | None = None):
yield
finally:
now = time.monotonic()
# Same ceiling as the open-window read in human_wait_seconds(): every
# legitimate wait self-terminates at approvals.timeout, so a window
# that overstayed it was wedged — record at most the ceiling instead
# of retroactively injecting the whole overstay into the exclusion.
ceiling = _human_wait_ceiling()
# Clamp the accrual too: a window that overstayed the ceiling was
# wedged — record at most the ceiling instead of retroactively
# injecting the whole overstay into the exclusion.
ceiling = human_wait_ceiling()
with _human_wait_lock:
state = _human_wait_states.get(key)
if state is not None:
state.pending -= 1
if state.pending == 0:
if state.window_started is not None:
state.completed_seconds += min(
max(0.0, now - state.window_started), ceiling
state.completed_seconds += _clamped_window_seconds(
state.window_started, now, ceiling
)
state.window_started = None
@@ -2324,28 +2335,29 @@ def human_wait_seconds(session_key: str | None = None) -> float:
"""Return total human-wait seconds recorded for the session.
Completed windows plus the currently open one (if any). Monotonically
non-decreasing for the life of the process, so deadline consumers snapshot
a baseline at batch start and use the delta.
non-decreasing for the life of the process — except when an idle session's
entry is evicted under cap pressure, which can only shrink a consumer's
baseline delta to zero (the safe direction: the deadline fires sooner).
Deadline consumers snapshot a baseline at batch start and use the delta.
The open window's contribution is clamped to the approval timeout plus a
small margin: every legitimate human wait self-terminates at
``approvals.timeout`` (both the CLI prompt join and the gateway poll loop
enforce it), so a window that overstays that bound is itself wedged and
must not keep extending a batch deadline (belt-and-braces for #79719).
Each window's contribution is clamped to :func:`human_wait_ceiling`:
every legitimate human wait self-terminates at ``approvals.timeout``
(both the CLI prompt join and the gateway poll loop enforce it), so a
window that overstays that bound is itself wedged and must not keep
extending a batch deadline (belt-and-braces for #79719).
"""
key = session_key if session_key is not None else get_current_session_key()
now = time.monotonic()
# Resolve the clamp outside the lock: it reads the config cache, which
# must never nest under _human_wait_lock.
ceiling = _human_wait_ceiling()
ceiling = human_wait_ceiling()
with _human_wait_lock:
state = _human_wait_states.get(key)
if state is None:
return 0.0
total = state.completed_seconds
if state.window_started is not None:
open_seconds = max(0.0, now - state.window_started)
total += min(open_seconds, ceiling)
total += _clamped_window_seconds(state.window_started, now, ceiling)
return total
# =========================================================================