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:
@@ -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:
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
# =========================================================================
|
||||
|
||||
Reference in New Issue
Block a user