From 1b886822deaf01a596493209a43a52197e974600 Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Sat, 1 Aug 2026 12:29:59 -0300 Subject: [PATCH] test(gateway): cover the run_generation guard and API-server disconnect reap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit dbbb10d39 shipped without direct test coverage for its own new logic — the same gap teknium's review flagged on the competing PR. Close it: - _reap_gateway_turn_processes: skips when is_still_current() is False, proceeds when True, fails open (reaps) if the check itself raises rather than silently disabling the leak fix. - _abandon_timed_out_gateway_turn: still marks the turn abandoned (interrupt fires) even when the reap itself is skipped. - api_server._reap_disconnected_agent_processes: reaps the baseline-diff for an owned turn, no-ops when the agent never recorded ownership markers. - APIServerAdapter._run_agent: markers are populated with the right task_id/baseline during the turn and cleared once it completes, closing the same race window fixed in gateway/run.py for this separate agent-lifecycle surface. --- .../test_abandoned_turn_process_cleanup.py | 100 ++++++++++++++++++ tests/gateway/test_api_server.py | 89 ++++++++++++++++ 2 files changed, 189 insertions(+) diff --git a/tests/gateway/test_abandoned_turn_process_cleanup.py b/tests/gateway/test_abandoned_turn_process_cleanup.py index 6d4a029e86758..67e62ffae8ed7 100644 --- a/tests/gateway/test_abandoned_turn_process_cleanup.py +++ b/tests/gateway/test_abandoned_turn_process_cleanup.py @@ -4,6 +4,7 @@ import threading from gateway.run import ( _abandon_timed_out_gateway_turn, + _reap_gateway_turn_processes, _watch_gateway_turn_inactivity, ) from tools.process_registry import process_registry @@ -112,3 +113,102 @@ def test_timeout_cleanup_is_idempotent(monkeypatch): assert not _abandon_timed_out_gateway_turn(**kwargs) assert len(calls) == 1 assert len(agent.interrupts) == 1 + + +# --------------------------------------------------------------------------- +# Cross-turn race guard (#76188 review): task_id is session-scoped, not +# turn-scoped, so a replacement turn on the same session could otherwise +# have its freshly-spawned process killed by a stale reaper. Gated on +# run_generation via an injected `is_still_current` check. +# --------------------------------------------------------------------------- + + +def test_reap_skips_when_a_newer_turn_has_claimed_the_session(monkeypatch): + calls = [] + monkeypatch.setattr( + process_registry, + "kill_started_since", + lambda *_a, **_k: calls.append(True) or 1, + ) + + killed = _reap_gateway_turn_processes( + "session-a", + frozenset({"proc_old"}), + source="gateway_turn_timeout", + is_still_current=lambda: False, + ) + + assert killed == 0 + assert calls == [] + + +def test_reap_proceeds_when_this_turn_is_still_current(monkeypatch): + calls = [] + monkeypatch.setattr( + process_registry, + "kill_started_since", + lambda task_id, baseline, *, source: calls.append( + (task_id, baseline, source) + ) + or 1, + ) + + killed = _reap_gateway_turn_processes( + "session-a", + frozenset({"proc_old"}), + source="gateway_turn_timeout", + is_still_current=lambda: True, + ) + + assert killed == 1 + assert calls == [("session-a", frozenset({"proc_old"}), "gateway_turn_timeout")] + + +def test_reap_fails_open_when_is_still_current_raises(monkeypatch): + """A bug in the generation-check closure must not silently disable the + underlying leak fix — it should log and fall through to reaping.""" + calls = [] + monkeypatch.setattr( + process_registry, + "kill_started_since", + lambda *_a, **_k: calls.append(True) or 1, + ) + + def _boom(): + raise RuntimeError("session state lookup failed") + + killed = _reap_gateway_turn_processes( + "session-a", + frozenset(), + source="gateway_turn_timeout", + is_still_current=_boom, + ) + + assert killed == 1 + assert calls == [True] + + +def test_timeout_abandon_propagates_is_still_current_to_the_reap(monkeypatch): + agent = _IdleAgent() + worker_done, timeout_fired, cleanup_lock = _state() + calls = [] + monkeypatch.setattr( + process_registry, + "kill_started_since", + lambda *_a, **_k: calls.append(True) or 1, + ) + + assert _abandon_timed_out_gateway_turn( + agent_holder=[agent], + task_id="session-a", + process_baseline=frozenset(), + worker_done=worker_done, + timeout_fired=timeout_fired, + cleanup_lock=cleanup_lock, + is_still_current=lambda: False, + ) + + # The turn was still marked abandoned (interrupt fired), but the actual + # reap was skipped because a newer turn already claimed the session. + assert agent.interrupts == ["Execution timed out (inactivity)"] + assert calls == [] diff --git a/tests/gateway/test_api_server.py b/tests/gateway/test_api_server.py index 86c3b03a77d06..ce921291e6e9b 100644 --- a/tests/gateway/test_api_server.py +++ b/tests/gateway/test_api_server.py @@ -392,6 +392,95 @@ class TestAgentExecution: task_id="session-123", ) + @pytest.mark.asyncio + async def test_run_agent_sets_and_clears_process_ownership_markers(self, adapter): + """#76188 review: this surface runs its own agent lifecycle outside + TurnRunner, so it needs its own baseline snapshot/clear — verify the + markers _reap_disconnected_agent_processes() reads are actually + populated during the turn and cleared once it finishes.""" + mock_agent = MagicMock() + mock_agent.session_prompt_tokens = 0 + mock_agent.session_completion_tokens = 0 + mock_agent.session_total_tokens = 0 + captured = {} + + def _capture_markers(**_kwargs): + captured["task_id"] = mock_agent._gateway_turn_process_task_id + captured["baseline"] = mock_agent._gateway_turn_process_baseline + return {"final_response": "ok"} + + mock_agent.run_conversation.side_effect = _capture_markers + + with patch.object(adapter, "_create_agent", return_value=mock_agent): + await adapter._run_agent( + user_message="hello", + conversation_history=[], + session_id="session-456", + requested_model="MiniMax-M3", + requested_provider="minimax", + model_options={"reasoning": {"enabled": False}, "fast": False}, + ) + + assert captured["task_id"] == "session-456" + assert isinstance(captured["baseline"], frozenset) + # Turn completed normally — markers must be cleared so a disconnect + # arriving after this point can't reap work this turn left running. + assert mock_agent._gateway_turn_process_task_id == "" + assert mock_agent._gateway_turn_process_baseline == frozenset() + + +class TestDisconnectedAgentReap: + """#76188 review: SSE disconnect handlers must reap only the background + processes the disconnected turn created, and must no-op when no turn + ownership was ever recorded on the agent.""" + + def test_reaps_baseline_diff_for_owned_turn(self, monkeypatch): + from gateway.platforms.api_server import _reap_disconnected_agent_processes + from tools.process_registry import process_registry + + calls = [] + monkeypatch.setattr( + process_registry, + "kill_started_since", + lambda task_id, baseline, *, source: calls.append( + (task_id, baseline, source) + ) + or 1, + ) + agent = types.SimpleNamespace( + _gateway_turn_process_task_id="session-abc", + _gateway_turn_process_baseline=frozenset({"proc-1"}), + ) + + _reap_disconnected_agent_processes(agent) + + deadline = time.time() + 1.0 + while not calls and time.time() < deadline: + time.sleep(0.01) + assert calls == [ + ("session-abc", frozenset({"proc-1"}), "api_server_sse_disconnect") + ] + + def test_noop_when_agent_has_no_ownership_markers(self, monkeypatch): + from gateway.platforms.api_server import _reap_disconnected_agent_processes + from tools.process_registry import process_registry + + calls = [] + monkeypatch.setattr( + process_registry, + "kill_started_since", + lambda *a, **k: calls.append(True), + ) + agent = types.SimpleNamespace( + _gateway_turn_process_task_id="", + _gateway_turn_process_baseline=None, + ) + + _reap_disconnected_agent_processes(agent) + + time.sleep(0.1) + assert calls == [] + class TestRunEventCallback: