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.
This commit is contained in:
kshitij 2026-08-05 13:00:22 +05:30 committed by kshitij
parent 62800ddadb
commit 1be70d6354
2 changed files with 40 additions and 0 deletions

View File

@ -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:

View File

@ -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"
)