From 10fb01e725bf251136891b283b0991d04aa374a4 Mon Sep 17 00:00:00 2001 From: kshitij Date: Thu, 6 Aug 2026 16:47:20 +0530 Subject: [PATCH] fix: harden human-wait tracker from review findings Review-driven follow-up to the #79719 fix: - Clamp the CLOSE-side accrual too: a wedged window that eventually closed used to inject its full unclamped overstay into completed_seconds, retroactively extending a running batch's deadline by hours. Both clamps now share one ceiling helper (_human_wait_ceiling = approvals.timeout + HUMAN_WAIT_MARGIN_S), and the gate's lock-timeout uses the same margin constant so the bounds cannot drift apart. - Evict idle sessions until the table is under the cap (was: at most one per insert, so churn could outgrow _HUMAN_WAIT_MAX_SESSIONS). Entries with an open window are still never evicted. - Log (debug) instead of silently swallowing a failed session-key snapshot in the gate constructor. Tests: close-side clamp regression + table-cap assertion added; suite at 17 passed. --- agent/tool_executor.py | 13 +++++--- tests/run_agent/test_authorization_gate.py | 17 +++++++++- tools/approval.py | 37 ++++++++++++++++++---- 3 files changed, 56 insertions(+), 11 deletions(-) diff --git a/agent/tool_executor.py b/agent/tool_executor.py index efcf9be09dfbf..c7d25804ed200 100644 --- a/agent/tool_executor.py +++ b/agent/tool_executor.py @@ -117,12 +117,13 @@ def _authorization_gate_lock_timeout() -> float: 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). + 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 _get_approval_timeout + from tools.approval import HUMAN_WAIT_MARGIN_S, _get_approval_timeout - return float(_get_approval_timeout()) + 60.0 + return float(_get_approval_timeout()) + HUMAN_WAIT_MARGIN_S except Exception: return _AUTHORIZATION_GATE_LOCK_TIMEOUT_S @@ -416,7 +417,11 @@ class _ConcurrentToolAuthorizationGate: # context may differ from the workers'. self._session_key = get_current_session_key() except Exception: - pass + 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 cffef52bc4fc0..f98c22344b5ab 100644 --- a/tests/run_agent/test_authorization_gate.py +++ b/tests/run_agent/test_authorization_gate.py @@ -118,9 +118,24 @@ class TestHumanWaitTracker: for i in range(approval_mod._HUMAN_WAIT_MAX_SESSIONS + 8): with approval_mod.human_wait_window(f"burst-{i}"): pass - # The active session survived the eviction pressure. + # The active session survived the eviction pressure and the table + # stayed at (or under) its cap. assert SESSION in approval_mod._human_wait_states assert approval_mod._human_wait_states[SESSION].pending == 1 + assert ( + len(approval_mod._human_wait_states) + <= approval_mod._HUMAN_WAIT_MAX_SESSIONS + ) + + def test_late_close_of_wedged_window_is_clamped(self, monkeypatch): + """A wedged window that eventually CLOSES must not retroactively inject + its full overstay into completed_seconds (close-side clamp).""" + monkeypatch.setattr(approval_mod, "_get_approval_timeout", lambda: 300) + with approval_mod.human_wait_window(SESSION): + state = approval_mod._human_wait_states[SESSION] + # Simulate the window having been open for a full day before close. + state.window_started = time.monotonic() - 86_400.0 + assert approval_mod.human_wait_seconds(SESSION) <= 300.0 + 60.0 class TestAuthorizationGate: diff --git a/tools/approval.py b/tools/approval.py index 41842044ea602..723f27d84307d 100644 --- a/tools/approval.py +++ b/tools/approval.py @@ -2239,22 +2239,42 @@ class _HumanWaitState: _human_wait_lock = threading.Lock() _human_wait_states: dict[str, _HumanWaitState] = {} _HUMAN_WAIT_MAX_SESSIONS = 256 +# Margin added on top of approvals.timeout when clamping a window's +# contribution (read-side AND close-side) and when bounding the authorization +# gate's serialization-lock acquire in agent/tool_executor.py. One constant so +# the clamps can't drift apart. +HUMAN_WAIT_MARGIN_S = 60.0 + + +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. + """ + return float(_get_approval_timeout()) + HUMAN_WAIT_MARGIN_S def _human_wait_state(session_key: str) -> _HumanWaitState: """Return (creating if needed) the wait state for *session_key*. Caller must hold ``_human_wait_lock``. Evicts idle entries (no pending - waiter) oldest-first when the table is full so an army of short-lived - session keys cannot grow it without bound. + waiter) insertion-order-first until the table is under the cap so an army + of short-lived session keys cannot grow it without bound. Entries with an + open window are never evicted (that would corrupt live accounting), so + the cap is best-effort under 256+ concurrently-pending sessions. """ state = _human_wait_states.get(session_key) if state is None: if len(_human_wait_states) >= _HUMAN_WAIT_MAX_SESSIONS: for key in list(_human_wait_states): + if len(_human_wait_states) < _HUMAN_WAIT_MAX_SESSIONS: + break if _human_wait_states[key].pending == 0: del _human_wait_states[key] - break state = _HumanWaitState() _human_wait_states[session_key] = state return state @@ -2283,14 +2303,19 @@ 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() 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 += max( - 0.0, now - state.window_started + state.completed_seconds += min( + max(0.0, now - state.window_started), ceiling ) state.window_started = None @@ -2312,7 +2337,7 @@ def human_wait_seconds(session_key: str | None = None) -> float: now = time.monotonic() # Resolve the clamp outside the lock: it reads the config cache, which # must never nest under _human_wait_lock. - ceiling = float(_get_approval_timeout()) + 60.0 + ceiling = _human_wait_ceiling() with _human_wait_lock: state = _human_wait_states.get(key) if state is None: