diff --git a/agent/tool_executor.py b/agent/tool_executor.py index c7d25804ed..2bc34bcf2f 100644 --- a/agent/tool_executor.py +++ b/agent/tool_executor.py @@ -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: diff --git a/tests/run_agent/test_authorization_gate.py b/tests/run_agent/test_authorization_gate.py index f98c22344b..7b29c0a26e 100644 --- a/tests/run_agent/test_authorization_gate.py +++ b/tests/run_agent/test_authorization_gate.py @@ -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. diff --git a/tools/approval.py b/tools/approval.py index 723f27d843..f7338cb65d 100644 --- a/tools/approval.py +++ b/tools/approval.py @@ -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 # =========================================================================