Fix exact gateway approval producer resolution

This commit is contained in:
Michael Snoyman authored and n committed 2026-09-22 06:29:27 +00:00
1 parent 571ef3eec2
commit 8c92d88225
4 files changed
+147 -50

No files matched your search

+53 -37
View File
@@ -10,15 +10,13 @@ Gateway approvals carry two pieces of state with one lifecycle:
* the **mirror** entry in ``_pending`` — this module's projection of that
producer for the UI.
Retirement resolves both together (``retire_gateway_pending_mirror``), and the
ownership direction matters: reconciliation rebuilds mirrors from live
producers, so a producer that survives its own resolution resurrects an
approval the user already dismissed. Retiring a specific ``approval_id``
therefore consumes that approval's producer, retiring a whole run consumes the
run's producers, and a same-run producer carrying no ``approval_id`` is retired
with it — it cannot be attributed to anything else, and leaving it queued is
what permits the resurrection. Retiring one approval must leave sibling
approvals and other runs queued.
A producer is consumed only by an explicit resolution operation, which assigns
its result and signals its waiter before reconciliation can observe it again.
``retire_gateway_pending_mirror`` cleans up the UI projection; it must not
silently delete a producer, because that would strand the agent thread parked
on the producer's event. For a provider-qualified response,
``resolve_gateway_pending_run`` matches the exact ``(run_id, approval_id)``
pair, settles that producer, and preserves same-run siblings and other runs.
"""
import queue
import threading
@@ -534,8 +532,6 @@ def retire_gateway_pending_mirror(
queue = _pending.get(session_key)
entries = queue if isinstance(queue, list) else [queue] if queue else []
normalized_run_id = str(run_id or "").strip()
gateway_queue = _gateway_queues.get(session_key) or []
retained_gateway_queue = gateway_queue
gateway_queue_changed = False
if approval_id:
match = _gateway_pending_mirror_locked(
@@ -560,13 +556,18 @@ def retire_gateway_pending_mirror(
and not str(entry.get("run_id") or "").strip()
]
if normalized_run_id:
gateway_queue = _gateway_queues.get(session_key) or []
retained_gateway_queue = []
for entry in gateway_queue:
data = getattr(entry, "data", None) or {}
if str(data.get("run_id") or "").strip() == normalized_run_id:
gateway_queue_changed = True
continue
retained_gateway_queue.append(entry)
gateway_queue_changed = len(retained_gateway_queue) != len(gateway_queue)
if retained_gateway_queue:
_gateway_queues[session_key] = retained_gateway_queue
else:
_gateway_queues.pop(session_key, None)
if not retired and not gateway_queue_changed:
head, total, changed = reconcile_gateway_pending_mirror_locked(session_key)
_approval_sse_notify_locked(session_key, head, total)
@@ -575,31 +576,6 @@ def retire_gateway_pending_mirror(
return changed
for match in retired:
entries.remove(match)
if normalized_run_id:
# A successful response may arrive with the exact approval_id. The
# mirror and its producer queue are one lifecycle: leaving the
# producer behind lets reconciliation recreate the just-retired
# mirror on the next poll (the completed-child resurrection bug).
# Remove only the matching approval for an active run; when the
# caller retires a run without an approval_id, retire the whole run.
def _same_retired_run_entry(entry) -> bool:
data = getattr(entry, "data", None) or {}
if str(data.get("run_id") or "").strip() != normalized_run_id:
return False
if not approval_id:
return True
queued_approval_id = str(data.get("approval_id") or "").strip()
return not queued_approval_id or queued_approval_id == approval_id
retained_gateway_queue = [
entry for entry in gateway_queue
if not _same_retired_run_entry(entry)
]
gateway_queue_changed = len(retained_gateway_queue) != len(gateway_queue)
if retained_gateway_queue:
_gateway_queues[session_key] = retained_gateway_queue
else:
_gateway_queues.pop(session_key, None)
if entries:
_pending[session_key] = entries
else:
@@ -853,6 +829,46 @@ def resolve_gateway_pending_local(
return 1, head, total
def resolve_gateway_pending_run(
session_key: str,
approval_id: str,
run_id: str,
choice: str,
reason: str | None = None,
) -> tuple[int, dict | None, int]:
"""Resolve one exact run-backed producer and wake its parked waiter."""
target = None
approval_id = str(approval_id or "").strip()
run_id = str(run_id or "").strip()
if not approval_id or not run_id:
return 0, None, 0
with _lock:
gateway_queue = _gateway_queues.get(session_key) or []
retained_gateway_queue = []
for entry in gateway_queue:
data = getattr(entry, "data", None) or {}
entry_approval_id = str(data.get("approval_id") or "").strip()
entry_run_id = str(data.get("run_id") or "").strip()
if target is None and entry_approval_id == approval_id and entry_run_id == run_id:
target = entry
else:
retained_gateway_queue.append(entry)
if retained_gateway_queue:
_gateway_queues[session_key] = retained_gateway_queue
else:
_gateway_queues.pop(session_key, None)
head, total, _changed = reconcile_gateway_pending_mirror_locked(session_key)
_approval_sse_notify_locked(session_key, head, total)
if target is None:
return 0, head, total
target.result = choice
if reason:
target.reason = reason
target.event.set()
publish_session_list_changed("attention_resolved")
return 1, head, total
def resolve_gateway_pending_local_no_run_mirror(
session_key: str, approval_id: str, choice: str, reason: str | None = None
) -> tuple[bool, int, dict | None, int]:
+5 -6
View File
@@ -10773,6 +10773,7 @@ from api.route_approvals import ( # noqa: F401 — re-exports for backward comp
retire_gateway_pending_mirror,
reconcile_gateway_pending_mirror_locked,
resolve_gateway_pending_local,
resolve_gateway_pending_run,
resolve_gateway_pending_local_all,
resolve_gateway_pending_local_no_run_mirror,
set_session_yolo_enabled,
@@ -26219,7 +26220,6 @@ def _resolve_approval_legacy(sid: str, approval_id: str, choice: str, run_id: st
found_target = False
gateway_keys = []
local_gateway_approval_id = ""
gateway_head_matches_target = False
with _lock:
reconcile_gateway_pending_mirror_locked(sid)
queue = _pending.get(sid)
@@ -26291,9 +26291,6 @@ def _resolve_approval_legacy(sid: str, approval_id: str, choice: str, run_id: st
gw_data = getattr(gw_entry, "data", None) or {}
gw_approval_id = str(gw_data.get("approval_id") or "").strip()
gw_run_id = str(gw_data.get("run_id") or "").strip()
gateway_head_matches_target = bool(
run_id and gw_approval_id == approval_id and gw_run_id == run_id
)
if gw_approval_id == approval_id and (not run_id or gw_run_id == run_id):
local_gateway_approval_id = approval_id
elif not run_id and found_target and pending:
@@ -26354,8 +26351,10 @@ def _resolve_approval_legacy(sid: str, approval_id: str, choice: str, run_id: st
local_gateway_resolved, _head, _total = resolve_gateway_pending_local(
sid, local_gateway_approval_id, choice
)
elif approval_id and found_target and run_id and gateway_head_matches_target:
gateway_resolved = resolve_gateway_approval(sid, choice, resolve_all=False) or 0
elif approval_id and found_target and run_id:
gateway_resolved, _head, _total = resolve_gateway_pending_run(
sid, approval_id, run_id, choice
)
elif not approval_id:
gateway_resolved = resolve_gateway_approval(sid, choice, resolve_all=False) or 0
# Keep the historical no-id response path truthy for old clients/tests while
+79
View File
@@ -601,6 +601,85 @@ class TestApprovalHTTPEndpoints:
r._gateway_queues.pop(sid, None)
_STREAM_RUN_IDS.pop(stream_id, None)
def test_identity_v1_non_head_run_approval_settles_exact_producer(self, monkeypatch):
"""A relayed non-head run approval must wake only its own parked waiter."""
from api import route_approvals as ra
from api import routes as r
sid = f"identity-v1-non-head-{uuid.uuid4().hex[:8]}"
run_id = f"run-identity-v1-{uuid.uuid4().hex[:8]}"
head = {
"approval_id": "head-approval",
"run_id": run_id,
"command": "head command",
"_gateway_agent_identity_v1": True,
}
target = {
"approval_id": "target-approval",
"run_id": run_id,
"command": "target command",
"_gateway_agent_identity_v1": True,
}
head_entry = _ApprovalEntry(dict(head))
target_entry = _ApprovalEntry(dict(target))
captured = {}
relays = []
def fake_j(_handler, data, status=200, extra_headers=None):
captured.update(payload=data, status=status)
return data
def fake_respond(_self, got_run_id, got_approval_id, choice):
relays.append((got_run_id, got_approval_id, choice))
return {"resolved": 1}
monkeypatch.setattr(r, "j", fake_j)
monkeypatch.setattr("api.runner_client.HttpRunnerClient.respond_approval", fake_respond)
monkeypatch.setattr(
"api.config.gateway_supports_approval_identity_v1", lambda *_args: True
)
with _lock:
r._pending.pop(sid, None)
r._gateway_queues[sid] = [head_entry, target_entry]
try:
ra.submit_gateway_pending_mirror(sid, dict(target_entry.data))
mirror = ra.gateway_pending_mirror(
sid, approval_id=target["approval_id"], run_id=run_id
)
assert mirror is not None
r._handle_approval_respond(
object(),
{
"session_id": sid,
"choice": "deny",
"approval_id": target["approval_id"],
"run_id": run_id,
"mirror_token": mirror[ra._GATEWAY_MIRROR_TOKEN],
},
)
assert captured == {
"payload": {"ok": True, "choice": "deny", "relayed": True},
"status": 200,
}
assert relays == [(run_id, target["approval_id"], "deny")]
assert target_entry.event.is_set()
assert target_entry.result == "deny"
assert not head_entry.event.is_set()
with _lock:
assert r._gateway_queues[sid] == [head_entry]
for _ in range(2):
with _lock:
ra.reconcile_gateway_pending_mirror_locked(sid)
assert ra.gateway_pending_mirror(
sid, approval_id=target["approval_id"], run_id=run_id
) is None
finally:
with _lock:
r._pending.pop(sid, None)
r._gateway_queues.pop(sid, None)
def test_gateway_mirror_without_run_id_with_one_producer_resolves_exactly(self, monkeypatch):
"""A no-run mirror retires only after its exact local producer resolves."""
from api import routes as r
+10 -7
View File
@@ -875,13 +875,13 @@ def test_legacy_approval_without_run_id_retires_locally():
_STREAM_RUN_IDS.pop(stream_id, None)
def test_deny_retires_run_producer_so_reconciliation_cannot_resurrect_mirror():
"""Deny must consume the producer as well as its WebUI mirror.
def test_deny_settles_run_producer_so_reconciliation_cannot_resurrect_mirror():
"""Deny must settle the producer as well as its WebUI mirror.
Selection is the behaviour under test, so the queue carries more than the
one entry being retired: a sibling approval for the same run, and a producer
belonging to a different run. Retiring one approval must consume exactly
that producer and leave both others queued.
one entry being resolved: a sibling approval for the same run, and a producer
belonging to a different run. Resolving one approval must consume exactly
that producer, wake its waiter, and leave both others queued.
"""
from types import SimpleNamespace
from api import route_approvals as ra
@@ -925,12 +925,13 @@ def test_deny_retires_run_producer_so_reconciliation_cannot_resurrect_mirror():
mirror = ra.gateway_pending_mirror(sid, approval_id=approval_id, run_id=run_id)
assert mirror is not None
assert ra.retire_gateway_pending_mirror(
resolved, _head, _total = ra.resolve_gateway_pending_run(
sid,
approval_id=approval_id,
run_id=run_id,
mirror_token=mirror[ra._GATEWAY_MIRROR_TOKEN],
choice="deny",
)
assert resolved == 1
assert queued_before == [target, same_run_sibling, other_run], (
"the queue should have held all three producers before retirement"
@@ -938,6 +939,8 @@ def test_deny_retires_run_producer_so_reconciliation_cannot_resurrect_mirror():
with ra._lock:
remaining = list(ra._gateway_queues.get(sid) or [])
assert target not in remaining, "the denied approval's producer must be consumed"
assert target.event.is_set(), "the denied approval's waiter must be woken"
assert target.result == "deny"
assert same_run_sibling in remaining, (
"a sibling approval on the same run must not be consumed by retiring another approval"
)