diff --git a/agent/tool_executor.py b/agent/tool_executor.py index 6b545f9e6f..efcf9be09d 100644 --- a/agent/tool_executor.py +++ b/agent/tool_executor.py @@ -102,6 +102,29 @@ _DEFAULT_CONCURRENT_TOOL_TIMEOUT_S = 420.0 # to cover slow-but-legitimate authorization (e.g. an approval round-trip), # short enough that one wedged dispatch cannot starve the batch forever. _START_ORDER_GATE_TIMEOUT_S = 120.0 +# Fallback bound a concurrent worker will wait for the authorization gate's +# serialization lock before running its prompt unserialized. The effective +# bound is derived from ``approvals.timeout`` plus a margin (see +# _authorization_gate_lock_timeout): a legitimate holder is at worst a human +# answering an approval prompt, which self-terminates at approvals.timeout — +# so a holder that overstays it is wedged and must not starve the batch. +_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). + """ + try: + from tools.approval import _get_approval_timeout + + return float(_get_approval_timeout()) + 60.0 + except Exception: + return _AUTHORIZATION_GATE_LOCK_TIMEOUT_S class _BatchAbandoned(BaseException): @@ -356,43 +379,72 @@ class _ManagedToolResult: class _ConcurrentToolAuthorizationGate: - """Serialize policy prompts and exclude their queue from batch deadlines.""" + """Serialize policy prompts and exclude human approval waits from batch deadlines. - def __init__(self) -> None: + Serialization keeps concurrent approval prompts from interleaving on the + user's screen. The acquire is BOUNDED: a worker wedged inside the gate (a + hanging ``pre_tool_call`` plugin, or an approval round-trip to a client + that went away) must not park every other worker forever. On expiry the + worker runs its prompt unserialized — worst case is interleaved prompts, + strictly better than permanent starvation (same tradeoff as the + start-order gate, #79705). + + Deadline exclusion is measured at the SOURCE of the human wait + (``tools.approval.human_wait_seconds``: the CLI prompt and the gateway + approval poll loop mark their own blocking windows), NOT as residency in + this gate. Gate residency is arbitrary code — using it as the exclusion + signal let a wedged plugin grow the exclusion 1:1 with wall clock, keeping + the batch deadline's ``remaining`` constant so it never fired and the turn + hung forever (#79719). A wedged plugin now contributes nothing to the + exclusion and the batch times out normally, while a genuine approval wait + (which can legitimately exceed any fixed bound) is still excluded in full. + """ + + def __init__(self, *, lock_timeout: float | None = None) -> None: self._serialization_lock = threading.Lock() - self._state_lock = threading.Lock() - self._pending = 0 - self._window_started: float | None = None - self._excluded_seconds = 0.0 + 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 + + # 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: + pass + self._baseline_wait_seconds = self._human_wait_seconds() + + def _human_wait_seconds(self) -> float: + try: + from tools.approval import human_wait_seconds + + return human_wait_seconds(self._session_key) + except Exception: + return 0.0 def run(self, callback): - now = time.monotonic() - with self._state_lock: - if self._pending == 0: - self._window_started = now - self._pending += 1 + acquired = self._serialization_lock.acquire(timeout=self._lock_timeout) + if not acquired: + logger.warning( + "authorization gate lock not acquired after %.1fs " + "(holder wedged in a pre_tool_call plugin or approval " + "round-trip?); running prompt unserialized", + self._lock_timeout, + ) + return callback() try: - with self._serialization_lock: - return callback() + return callback() finally: - now = time.monotonic() - with self._state_lock: - self._pending -= 1 - if self._pending == 0: - if self._window_started is not None: - self._excluded_seconds += max( - 0.0, now - self._window_started - ) - self._window_started = None + self._serialization_lock.release() def excluded_seconds(self) -> float: - """Return completed plus currently active authorization wait time.""" - now = time.monotonic() - with self._state_lock: - excluded = self._excluded_seconds - if self._window_started is not None: - excluded += max(0.0, now - self._window_started) - return excluded + """Return human-approval wait seconds accrued since the batch started.""" + return max(0.0, self._human_wait_seconds() - self._baseline_wait_seconds) def _managed_values( diff --git a/tests/run_agent/test_authorization_gate.py b/tests/run_agent/test_authorization_gate.py new file mode 100644 index 0000000000..cffef52bc4 --- /dev/null +++ b/tests/run_agent/test_authorization_gate.py @@ -0,0 +1,308 @@ +"""Tests for the concurrent authorization gate and human-wait accounting (#79719). + +Before the fix, a worker wedged *inside* the authorization gate — a hanging +``pre_tool_call`` plugin, or an approval round-trip to a client that went +away — had two coupled failure modes: + +1. The serialization lock was an unbounded blocking acquire: every other + worker needing authorization blocked behind the wedged holder forever. +2. ``excluded_seconds()`` measured residency in ``gate.run()`` (arbitrary + code), so an open window grew 1:1 with wall clock while the batch-deadline + loop added it to the deadline on every poll — ``remaining`` was constant + and the deadline NEVER fired. Algebraically: + ``remaining = (deadline + (now - window_started)) - now = deadline - window_started``. + +The fix moves deadline exclusion to the source of the human wait +(``tools.approval.human_wait_window`` around the CLI prompt and the gateway +approval poll loop) and bounds the serialization lock acquire. A wedged +plugin now contributes nothing to the exclusion, so the batch times out +normally; a genuine approval wait is still excluded in full. +""" + +import threading +import time + +import pytest + +from agent.tool_executor import _ConcurrentToolAuthorizationGate +from tools import approval as approval_mod + + +@pytest.fixture(autouse=True) +def _clean_human_wait_state(): + with approval_mod._human_wait_lock: + approval_mod._human_wait_states.clear() + yield + with approval_mod._human_wait_lock: + approval_mod._human_wait_states.clear() + + +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 + + +class TestHumanWaitTracker: + def test_no_wait_reports_zero(self): + assert approval_mod.human_wait_seconds(SESSION) == 0.0 + + def test_open_window_counts(self): + opened = threading.Event() + release = threading.Event() + + def _wait(): + with approval_mod.human_wait_window(SESSION): + opened.set() + release.wait(timeout=5) + + t = threading.Thread(target=_wait, daemon=True) + t.start() + assert opened.wait(timeout=5) + time.sleep(0.05) + assert approval_mod.human_wait_seconds(SESSION) > 0.0 + release.set() + t.join(timeout=5) + # Window closed: total is frozen (completed_seconds), not still growing. + first = approval_mod.human_wait_seconds(SESSION) + time.sleep(0.05) + assert approval_mod.human_wait_seconds(SESSION) == pytest.approx(first) + + def test_overlapping_windows_coalesce(self): + """Two concurrent windows on one session must not double-count wall clock.""" + release = threading.Event() + started = threading.Barrier(3) + + def _wait(): + with approval_mod.human_wait_window(SESSION): + started.wait(timeout=5) + release.wait(timeout=5) + + threads = [threading.Thread(target=_wait, daemon=True) for _ in range(2)] + start = time.monotonic() + for t in threads: + t.start() + started.wait(timeout=5) + time.sleep(0.1) + release.set() + for t in threads: + t.join(timeout=5) + elapsed = time.monotonic() - start + # Coalesced: recorded ≤ wall clock (a double count would be ~2×). + assert approval_mod.human_wait_seconds(SESSION) <= elapsed + 0.05 + + def test_sessions_are_isolated(self): + with approval_mod.human_wait_window("other-session"): + time.sleep(0.05) + assert approval_mod.human_wait_seconds(SESSION) == 0.0 + assert approval_mod.human_wait_seconds("other-session") > 0.0 + + def test_open_window_clamped_to_approval_timeout(self, monkeypatch): + """A window that overstays approvals.timeout is itself wedged and must + stop extending the exclusion (belt-and-braces for #79719).""" + monkeypatch.setattr(approval_mod, "_get_approval_timeout", lambda: 300) + with approval_mod.human_wait_window(SESSION): + state = approval_mod._human_wait_states[SESSION] + # Simulate a window that has been open for a full day. + state.window_started = time.monotonic() - 86_400.0 + assert approval_mod.human_wait_seconds(SESSION) <= 300.0 + 60.0 + + def test_eviction_keeps_pending_sessions(self): + with approval_mod.human_wait_window(SESSION): + 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. + assert SESSION in approval_mod._human_wait_states + assert approval_mod._human_wait_states[SESSION].pending == 1 + + +class TestAuthorizationGate: + def test_serializes_callbacks(self): + gate = _make_gate() + state_lock = threading.Lock() + active = 0 + max_active = 0 + + def _callback(): + nonlocal active, max_active + with state_lock: + active += 1 + max_active = max(max_active, active) + try: + time.sleep(0.03) + finally: + with state_lock: + active -= 1 + + threads = [ + threading.Thread(target=lambda: gate.run(_callback), daemon=True) + for _ in range(4) + ] + for t in threads: + t.start() + for t in threads: + t.join(timeout=5) + assert max_active == 1 + + def test_lock_timeout_degrades_to_unserialized(self): + """A wedged lock holder must not park later callers forever.""" + gate = _make_gate(lock_timeout=0.1) + holder_in = threading.Event() + release = threading.Event() + + def _wedged(): + holder_in.set() + release.wait(timeout=10) + + holder = threading.Thread(target=lambda: gate.run(_wedged), daemon=True) + holder.start() + assert holder_in.wait(timeout=5) + + done = threading.Event() + result = {} + + def _second(): + result["value"] = gate.run(lambda: "ran-unserialized") + done.set() + + t = threading.Thread(target=_second, daemon=True) + start = time.monotonic() + t.start() + assert done.wait(timeout=5), "second caller starved behind wedged holder" + assert result["value"] == "ran-unserialized" + assert time.monotonic() - start < 2.0 + release.set() + holder.join(timeout=5) + + def test_wedged_callback_contributes_nothing_to_exclusion(self): + """THE #79719 regression: gate residency is not deadline exclusion.""" + gate = _make_gate() + wedged_in = threading.Event() + release = threading.Event() + + def _wedged(): + wedged_in.set() + release.wait(timeout=10) + + t = threading.Thread(target=lambda: gate.run(_wedged), daemon=True) + t.start() + assert wedged_in.wait(timeout=5) + time.sleep(0.15) + # No human prompt is pending — the wedge is invisible to the deadline. + assert gate.excluded_seconds() == 0.0 + release.set() + t.join(timeout=5) + + def test_deadline_arithmetic_converges_with_wedged_worker(self): + """The issue's repro: remaining must DECREASE while a worker is wedged. + + Pre-fix, ``remaining = deadline - window_started`` was constant for + the life of the wedge (24h simulated in the issue). Now the exclusion + stays 0 for a wedge, so remaining tracks wall clock down to zero. + """ + gate = _make_gate() + wedged_in = threading.Event() + release = threading.Event() + + def _wedged(): + wedged_in.set() + release.wait(timeout=10) + + t = threading.Thread(target=lambda: gate.run(_wedged), daemon=True) + t.start() + assert wedged_in.wait(timeout=5) + + timeout_s = 0.3 + deadline = time.monotonic() + timeout_s + first = deadline + gate.excluded_seconds() - time.monotonic() + time.sleep(0.15) + second = deadline + gate.excluded_seconds() - time.monotonic() + assert second < first, "remaining is constant — deadline never fires (#79719)" + time.sleep(0.25) + assert deadline + gate.excluded_seconds() - time.monotonic() <= 0, ( + "deadline never became due despite the wedge" + ) + release.set() + t.join(timeout=5) + + def test_human_wait_is_excluded(self): + """A genuine approval wait during the batch extends the deadline.""" + gate = _make_gate() + with approval_mod.human_wait_window(SESSION): + time.sleep(0.1) + assert gate.excluded_seconds() >= 0.09 + + def test_baseline_ignores_waits_before_batch(self): + """Approval waits from BEFORE this batch must not extend its deadline.""" + with approval_mod.human_wait_window(SESSION): + time.sleep(0.1) + gate = _make_gate() + assert gate.excluded_seconds() == 0.0 + + def test_other_sessions_wait_not_excluded(self): + gate = _make_gate() + with approval_mod.human_wait_window("unrelated-session"): + time.sleep(0.05) + assert gate.excluded_seconds() == 0.0 + + +class TestApprovalPathsRecordHumanWait: + def test_await_gateway_decision_records_wait(self, monkeypatch): + """The gateway approval poll loop must mark itself as human wait.""" + monkeypatch.setattr(approval_mod, "_get_approval_timeout", lambda: 300) + approval_data = { + "command": "rm -rf /tmp/x", + "description": "test", + "pattern_key": "k", + "pattern_keys": ["k"], + } + notified = threading.Event() + result_holder = {} + + def _worker(): + result_holder["result"] = approval_mod._await_gateway_decision( + SESSION, lambda _data: notified.set(), approval_data + ) + + t = threading.Thread(target=_worker, daemon=True) + t.start() + assert notified.wait(timeout=5) + time.sleep(0.1) + 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() + t.join(timeout=5) + assert not t.is_alive() + # Window closed once the wait resolved. + assert approval_mod._human_wait_states[SESSION].pending == 0 + + def test_prompt_dangerous_approval_records_wait(self, monkeypatch): + """The CLI prompt path must mark itself as human wait.""" + observed = {} + + def _callback(_command, _description, **_kwargs): + observed["during"] = approval_mod.human_wait_seconds() + return "deny" + + choice = approval_mod.prompt_dangerous_approval( + "rm -rf /tmp/x", "test", approval_callback=_callback + ) + assert choice == "deny" + # The window was open while the callback (the human prompt) ran. + state = approval_mod._human_wait_states.get( + approval_mod.get_current_session_key() + ) + assert state is not None + assert state.pending == 0 diff --git a/tools/approval.py b/tools/approval.py index f5e7f0beb1..41842044ea 100644 --- a/tools/approval.py +++ b/tools/approval.py @@ -8,6 +8,7 @@ This module is the single source of truth for the dangerous command system: - Permanent allowlist persistence (config.yaml) """ +import contextlib import contextvars import fnmatch import functools @@ -2208,6 +2209,120 @@ _session_approved: dict[str, set] = {} _session_yolo: set[str] = set() _permanent_approved: set = set() + +# ========================================================================= +# Human-wait accounting (per session) +# ========================================================================= +# Tracks the wall-clock time the agent spends verifiably blocked on a HUMAN +# prompt (CLI approval prompt, gateway approval round-trip). The concurrent +# 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 +# approval client from growing the exclusion 1:1 with wall clock and defeating +# the deadline entirely (#79719). +# +# Keyed by session so one gateway session's pending approval cannot extend a +# different session's batch deadline. State is process-global like the rest +# of this module's approval state; entries are bounded by _HUMAN_WAIT_MAX_SESSIONS. + + +class _HumanWaitState: + __slots__ = ("pending", "window_started", "completed_seconds") + + def __init__(self) -> None: + self.pending = 0 + self.window_started: float | None = None + self.completed_seconds = 0.0 + + +_human_wait_lock = threading.Lock() +_human_wait_states: dict[str, _HumanWaitState] = {} +_HUMAN_WAIT_MAX_SESSIONS = 256 + + +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. + """ + 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 _human_wait_states[key].pending == 0: + del _human_wait_states[key] + break + state = _HumanWaitState() + _human_wait_states[session_key] = state + return state + + +@contextlib.contextmanager +def human_wait_window(session_key: str | None = None): + """Mark the enclosed block as time spent blocked on a human prompt. + + Wrap ONLY code that is genuinely parked waiting for a user's answer (the + CLI approval prompt, the gateway approval poll loop). The concurrent tool + batch deadline excludes this time; wrapping anything else re-creates the + #79719 hang where arbitrary wedged code pushes the deadline out forever. + + Overlapping windows for the same session coalesce (pending counter), so + two serialized approval prompts don't double-count the same wall clock. + """ + key = session_key if session_key is not None else get_current_session_key() + now = time.monotonic() + with _human_wait_lock: + state = _human_wait_state(key) + if state.pending == 0: + state.window_started = now + state.pending += 1 + try: + yield + finally: + now = time.monotonic() + 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.window_started = None + + +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. + + 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). + """ + 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 = float(_get_approval_timeout()) + 60.0 + 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) + return total + # ========================================================================= # Consecutive-denial circuit breaker for smart approvals # ========================================================================= @@ -2585,6 +2700,26 @@ def prompt_dangerous_approval(command: str, description: str, if timeout_seconds is None: timeout_seconds = _get_approval_timeout() + # Everything below is a human prompt: either the registered CLI callback + # (prompt_toolkit panel, bounded by the approval deadline) or the input() + # fallback (bounded by thread.join(timeout_seconds)). Record it as + # human-wait time so the concurrent batch deadline excludes it (#79719). + with human_wait_window(): + return _prompt_dangerous_approval_inner( + command, + description, + timeout_seconds, + allow_permanent, + approval_callback, + smart_denied=smart_denied, + ) + + +def _prompt_dangerous_approval_inner(command: str, description: str, + timeout_seconds: int, + allow_permanent: bool = True, + approval_callback=None, + *, smart_denied: bool = False) -> str: # Redact secrets before any user-visible rendering. The original # `command` is still what executes after approval; only the displayed # copy is scrubbed. Reuses the same redaction module used for memory @@ -3511,32 +3646,37 @@ def _await_gateway_decision(session_key: str, notify_cb, approval_data: dict, _deadline = _now + max(timeout, 0) _activity_state = {"last_touch": _now, "start": _now} resolved = False - while True: - # Respect interrupt signals (e.g. /stop, /new, or an inactivity - # timeout from the gateway) so a pending approval doesn't keep the - # session wedged on threading.Event.wait() until the 5-minute approval - # timeout. The wait runs on the agent's execution thread, which is the - # exact thread AIAgent.interrupt() flags — so is_interrupted() here - # sees the signal. Resolve as "deny" so the agent loop receives a - # normal denial and unwinds cleanly (#8697). - if is_interrupted(): - logger.info( - "Approval wait interrupted by user signal — " - "returning deny for session %s", - session_key, - ) - entry.result = "deny" - entry.event.set() - resolved = True - break - _remaining = _deadline - time.monotonic() - if _remaining <= 0: - break - if entry.event.wait(timeout=min(1.0, _remaining)): - resolved = True - break - if touch_activity_if_due is not None: - touch_activity_if_due(_activity_state, "waiting for user approval") + # The poll loop below is verifiably blocked on a human answer (the user + # tapping approve/deny on the gateway surface), bounded by the approval + # timeout. Record it as human-wait time so the concurrent batch deadline + # excludes it (#79719). + with human_wait_window(session_key): + while True: + # Respect interrupt signals (e.g. /stop, /new, or an inactivity + # timeout from the gateway) so a pending approval doesn't keep the + # session wedged on threading.Event.wait() until the 5-minute approval + # timeout. The wait runs on the agent's execution thread, which is the + # exact thread AIAgent.interrupt() flags — so is_interrupted() here + # sees the signal. Resolve as "deny" so the agent loop receives a + # normal denial and unwinds cleanly (#8697). + if is_interrupted(): + logger.info( + "Approval wait interrupted by user signal — " + "returning deny for session %s", + session_key, + ) + entry.result = "deny" + entry.event.set() + resolved = True + break + _remaining = _deadline - time.monotonic() + if _remaining <= 0: + break + if entry.event.wait(timeout=min(1.0, _remaining)): + resolved = True + break + if touch_activity_if_due is not None: + touch_activity_if_due(_activity_state, "waiting for user approval") _drop_entry()