diff --git a/gateway/run.py b/gateway/run.py index 5b6fd4cf44940..5bf2c0625a770 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -17201,6 +17201,59 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew return (evt_type, producer_id, started_at) return None + async def _classify_completion_target(self, parent_session_id: str) -> str: + """Classify an async-completion delivery target before adapter acceptance. + + Returns one of: + + - ``"deliver"`` — the spawning session is live, or ended by a + compression rotation with a verified live continuation. The inner + #55578 resolver (:meth:`_resolve_async_delegation_session`) still + owns the actual route retarget; this pre-flight only proves the + completion is deliverable so the durable ack stays honest. + - ``"terminal"`` — the spawning session is gone for good (unknown, or + ended at an explicit user boundary such as /new). Delivery can never + succeed; the durable row should be terminally dropped rather than + falsely acknowledged as delivered or replayed forever as pending. + - ``"retry"`` — transient uncertainty (session DB unavailable, lookup + error, or a compression rotation caught mid-flight before its + continuation exists). The claim should be released so a later + consumer can retry; the attempt cap bounds the churn. + """ + session_db = getattr(self, "_session_db", None) + if session_db is None: + return "retry" + try: + parent = await session_db.get_session(parent_session_id) + except Exception: + logger.debug( + "Async-completion pre-flight parent lookup failed for %s", + parent_session_id, exc_info=True, + ) + return "retry" + if parent is None: + return "terminal" + if not parent.get("ended_at"): + return "deliver" + if parent.get("end_reason") != "compression": + return "terminal" + try: + tip_session_id = await session_db.get_compression_tip(parent_session_id) + if not tip_session_id or tip_session_id == parent_session_id: + # Rotation caught mid-flight: parent is compression-ended but + # its continuation isn't visible yet. Retry, don't drop. + return "retry" + tip = await session_db.get_session(tip_session_id) + except Exception: + logger.debug( + "Async-completion pre-flight tip lookup failed for %s", + parent_session_id, exc_info=True, + ) + return "retry" + if tip is None or tip.get("ended_at"): + return "retry" + return "deliver" + async def _deliver_completion_notification( self, synth_text: str, evt: dict, ) -> Optional[bool]: @@ -17232,6 +17285,49 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew durable_delegation_id, exc, ) return False + parent_session_id = str(evt.get("parent_session_id") or "").strip() + if parent_session_id: + # Pre-flight (#65838-class): adapter acceptance is NOT proof of + # delivery — the inner #55578 resolver can still fail closed + # inside the message pipeline AFTER the adapter accepted, which + # would falsely acknowledge the durable row as delivered. + # Verify the target here, before acceptance, and give drops an + # honest durable disposition. + verdict = await self._classify_completion_target(parent_session_id) + if verdict == "terminal": + logger.warning( + "Async delegation %s targets permanently-gone session %s; " + "terminally dropping delivery (result remains in the " + "delegation records).", + durable_delegation_id or "", parent_session_id, + ) + if durable_claim_id: + try: + from tools.async_delegation import drop_completion_delivery + + drop_completion_delivery( + durable_delegation_id, durable_claim_id, + ) + except Exception: + logger.debug( + "Could not drop durable completion claim", + exc_info=True, + ) + return None + if verdict == "retry": + if durable_claim_id: + try: + from tools.async_delegation import release_completion_delivery + + release_completion_delivery( + durable_delegation_id, durable_claim_id, + ) + except Exception: + logger.debug( + "Could not release durable completion claim", + exc_info=True, + ) + return False if identity is not None: with self._completion_delivery_lock: if ( diff --git a/tests/gateway/test_completion_delivery.py b/tests/gateway/test_completion_delivery.py index 57c75f755a665..4face3841b271 100644 --- a/tests/gateway/test_completion_delivery.py +++ b/tests/gateway/test_completion_delivery.py @@ -186,6 +186,164 @@ def test_failed_async_injection_is_retried_and_only_success_is_acked( assert acknowledgements == ["deleg_duplicate"] +def _persist_pending_completion(event): + from tools import async_delegation + + async_delegation._persist_dispatch({ + "delegation_id": event["delegation_id"], + "session_key": event["session_key"], + "origin_ui_session_id": "", + "parent_session_id": event.get("parent_session_id"), + "dispatched_at": event["dispatched_at"], + }) + async_delegation._persist_completion(event, { + "status": "completed", + "summary": event["summary"], + }) + + +def test_compression_parent_delivery_targets_tip_and_is_acked( + monkeypatch, isolated_registry, +): + """A compression-rotated parent with a live tip is deliverable + acked.""" + from tools import async_delegation + + event = _async_event("deleg_compression") + event["parent_session_id"] = "sess_parent" + _persist_pending_completion(event) + + adapter = SimpleNamespace(handle_message=AsyncMock()) + runner = _runner(adapter) + runner._session_db = SimpleNamespace( + get_session=AsyncMock(side_effect=lambda session_id: { + "sess_parent": { + "id": "sess_parent", + "ended_at": "2026-07-16T12:00:00", + "end_reason": "compression", + }, + "sess_tip": {"id": "sess_tip", "ended_at": None, "end_reason": None}, + }.get(session_id)), + get_compression_tip=AsyncMock(return_value="sess_tip"), + ) + + assert asyncio.run( + runner._deliver_completion_notification("completion", event) + ) is True + + adapter.handle_message.assert_awaited_once() + durable = async_delegation.get_durable_delegation(event["delegation_id"]) + assert durable is not None + assert durable["delivery_state"] == "delivered" + + +def test_explicit_reset_drop_is_terminal_not_falsely_delivered( + monkeypatch, isolated_registry, +): + """An explicit /new boundary drop gets a terminal 'dropped' disposition. + + Not 'delivered' (the ack must stay honest — nothing was injected) and not + 'pending' (restart recovery would replay a completion that is fail-closed + dropped again on every boot). + """ + from tools import async_delegation + + event = _async_event("deleg_explicit_new") + event["parent_session_id"] = "sess_reset" + _persist_pending_completion(event) + + adapter = SimpleNamespace(handle_message=AsyncMock()) + runner = _runner(adapter) + runner._session_db = SimpleNamespace( + get_session=AsyncMock(return_value={ + "id": "sess_reset", + "ended_at": "2026-07-16T12:00:00", + "end_reason": "session_reset", + }), + get_compression_tip=AsyncMock(), + ) + + assert asyncio.run( + runner._deliver_completion_notification("completion", event) + ) is None + + adapter.handle_message.assert_not_awaited() + durable = async_delegation.get_durable_delegation(event["delegation_id"]) + assert durable is not None + assert durable["delivery_state"] == "dropped" + restored = queue.Queue() + assert async_delegation.restore_undelivered_completions(restored) == 0 + + +def test_midflight_compression_rotation_stays_pending_for_retry( + monkeypatch, isolated_registry, +): + """A rotation without a visible continuation yet is retryable, not dropped.""" + from tools import async_delegation + + event = _async_event("deleg_midflight") + event["parent_session_id"] = "sess_rotating" + _persist_pending_completion(event) + + adapter = SimpleNamespace(handle_message=AsyncMock()) + runner = _runner(adapter) + runner._session_db = SimpleNamespace( + get_session=AsyncMock(return_value={ + "id": "sess_rotating", + "ended_at": "2026-07-16T12:00:00", + "end_reason": "compression", + }), + get_compression_tip=AsyncMock(return_value=None), + ) + + assert asyncio.run( + runner._deliver_completion_notification("completion", event) + ) is False + + adapter.handle_message.assert_not_awaited() + durable = async_delegation.get_durable_delegation(event["delegation_id"]) + assert durable is not None + assert durable["delivery_state"] == "pending" + restored = queue.Queue() + assert async_delegation.restore_undelivered_completions(restored) == 1 + assert restored.get_nowait()["delegation_id"] == event["delegation_id"] + + +def test_retry_attempts_are_capped_to_a_terminal_drop( + monkeypatch, isolated_registry, +): + """Endless claim/release churn converges to a terminal 'dropped' state.""" + from tools import async_delegation + + event = _async_event("deleg_attempt_cap") + event["parent_session_id"] = "sess_rotating" + _persist_pending_completion(event) + + adapter = SimpleNamespace(handle_message=AsyncMock()) + runner = _runner(adapter) + runner._session_db = SimpleNamespace( + get_session=AsyncMock(return_value={ + "id": "sess_rotating", + "ended_at": "2026-07-16T12:00:00", + "end_reason": "compression", + }), + get_compression_tip=AsyncMock(return_value=None), + ) + + async def _churn(): + for _ in range(async_delegation._MAX_DELIVERY_ATTEMPTS + 2): + await runner._deliver_completion_notification("completion", event) + + asyncio.run(_churn()) + + adapter.handle_message.assert_not_awaited() + durable = async_delegation.get_durable_delegation(event["delegation_id"]) + assert durable is not None + assert durable["delivery_state"] == "dropped" + assert durable["delivery_attempts"] <= async_delegation._MAX_DELIVERY_ATTEMPTS + restored = queue.Queue() + assert async_delegation.restore_undelivered_completions(restored) == 0 + + def test_distinct_process_incarnations_are_not_deduplicated(): """Producer spawn time distinguishes a reused process session ID.""" adapter = SimpleNamespace(handle_message=AsyncMock()) diff --git a/tools/async_delegation.py b/tools/async_delegation.py index 6811d6867fdb4..5181149c4d3ac 100644 --- a/tools/async_delegation.py +++ b/tools/async_delegation.py @@ -77,6 +77,11 @@ _DEFAULT_MAX_ASYNC_CHILDREN = 3 _MAX_RETAINED_COMPLETED = 50 _DURABLE_RETENTION_SECONDS = 7 * 24 * 60 * 60 _MAX_DURABLE_PENDING = 1000 +# A pending completion whose delivery keeps failing is retried across claim +# cycles (and across restarts via restore_undelivered_completions). Cap the +# attempts so an unroutable row converges to a terminal 'dropped' state +# instead of replaying on every restart forever. +_MAX_DELIVERY_ATTEMPTS = 8 _DB_LOCK = threading.Lock() @@ -334,14 +339,59 @@ def claim_event_delivery(evt: Dict[str, Any], consumer: str) -> Optional[str]: def release_completion_delivery(delegation_id: str, claim_id: str) -> bool: - """Release a failed delivery claim so another consumer may retry.""" + """Release a failed delivery claim so another consumer may retry. + + Attempts are counted at claim time, so a row that keeps being claimed and + released has burned real delivery attempts. Once the budget is exhausted + the row converges to a terminal ``dropped`` state instead of returning to + ``pending`` — otherwise an undeliverable completion replays on every + gateway restart forever (restore_undelivered_completions only restores + pending rows). + """ + now = time.time() with _DB_LOCK, _connect() as conn: + capped = conn.execute( + """UPDATE async_delegations SET delivery_state='dropped', + delivery_claim=NULL, delivery_claimed_at=NULL, updated_at=? + WHERE delegation_id=? AND delivery_state='pending' + AND delivery_claim=? AND delivery_attempts>=?""", + (now, delegation_id, claim_id, _MAX_DELIVERY_ATTEMPTS), + ) + if capped.rowcount == 1: + logger.warning( + "Async delegation %s exhausted its %d delivery attempts; " + "marking terminally dropped (result remains queryable).", + delegation_id, _MAX_DELIVERY_ATTEMPTS, + ) + return True cur = conn.execute( """UPDATE async_delegations SET delivery_claim=NULL, delivery_claimed_at=NULL, updated_at=? WHERE delegation_id=? AND delivery_state='pending' AND delivery_claim=?""", - (time.time(), delegation_id, claim_id), + (now, delegation_id, claim_id), + ) + return cur.rowcount == 1 + + +def drop_completion_delivery(delegation_id: str, claim_id: str) -> bool: + """Terminally drop a claimed completion that can never be delivered. + + Used when the delivery target is permanently gone — the spawning session + ended at an explicit user boundary (/new, reset) rather than a compression + rotation. Marking the row ``dropped`` (not ``delivered``) keeps the ack + honest, and (not ``pending``) keeps restart recovery from replaying a + completion that will be fail-closed dropped again every time. + """ + now = time.time() + with _DB_LOCK, _connect() as conn: + cur = conn.execute( + """UPDATE async_delegations SET delivery_state='dropped', + updated_at=?, delivery_claim=NULL, + delivery_claimed_at=NULL + WHERE delegation_id=? AND delivery_state='pending' + AND delivery_claim=?""", + (now, delegation_id, claim_id), ) return cur.rowcount == 1