From 1be70d63548845eb8918c08ed698cda0674cf9a7 Mon Sep 17 00:00:00 2001 From: kshitij Date: Wed, 5 Aug 2026 13:00:22 +0530 Subject: [PATCH] fix: join heartbeat thread in finally + add error-path test Add activity_hb.join(timeout=2.0) after activity_hb_stop.set() in direct_api_call's finally block so the heartbeat thread is deterministically stopped before client teardown. Add test verifying no stray _touch_activity fires after direct_api_call raises an exception. Follow-up to PR #78548 by @xxxigm. --- agent/chat_completion_helpers.py | 1 + tests/cron/test_cron_direct_api_call_62151.py | 39 +++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index bc81774e4ae1d..313e7cd865be3 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -647,6 +647,7 @@ def direct_api_call(agent, api_kwargs: dict): return response finally: activity_hb_stop.set() + activity_hb.join(timeout=2.0) if getattr(agent, "_active_request_abort", None) is _abort_active_request: agent._active_request_abort = None with request_client_lock: diff --git a/tests/cron/test_cron_direct_api_call_62151.py b/tests/cron/test_cron_direct_api_call_62151.py index 56bbc4db1749a..222027e61e065 100644 --- a/tests/cron/test_cron_direct_api_call_62151.py +++ b/tests/cron/test_cron_direct_api_call_62151.py @@ -128,3 +128,42 @@ def test_direct_api_call_keeps_activity_alive_during_slow_wait(monkeypatch): call.args[0] == "waiting for non-streaming API response" for call in agent._touch_activity.call_args_list ) + + +def test_direct_api_call_heartbeat_stops_on_exception(monkeypatch): + """The activity heartbeat thread must be joined on error paths so no + stray _touch_activity fires after the call has failed. + """ + import threading + import time + + from agent import chat_completion_helpers as helpers + + monkeypatch.setattr(helpers, "_DIRECT_API_ACTIVITY_HEARTBEAT_SECONDS", 0.05) + + agent = _make_agent(platform="subagent") + fake_client = MagicMock() + fake_client.chat.completions.create.side_effect = RuntimeError("provider down") + agent._create_request_openai_client.return_value = fake_client + + raised = threading.Event() + + def _runner(): + try: + direct_api_call(agent, {"model": "m", "messages": []}) + except RuntimeError: + raised.set() + + worker = threading.Thread(target=_runner, daemon=True) + worker.start() + worker.join(timeout=3.0) + + assert raised.is_set(), "expected RuntimeError from direct_api_call" + # Give any stray heartbeat a chance to fire after the call returned. + time.sleep(0.2) + touches_before = agent._touch_activity.call_count + time.sleep(0.2) + touches_after = agent._touch_activity.call_count + assert touches_after == touches_before, ( + "heartbeat thread still firing after direct_api_call raised" + )