feat(websocket): add heartbeat dead-threshold watchdog - #595
AlexanderZ-Band merged 11 commits into
Conversation
WebSocketClient now drives phoenix-channels-python-client's heartbeat interval from band-sdk-core's SessionPolicy and runs a watchdog task that force-closes the connection if no heartbeat ack arrives within dead_threshold_s, so a silently dead socket gets reconnected instead of sitting idle. The watchdog is scoped 1:1 to the PHXChannelsClient instance it watches, so a stale watchdog from a superseded initial-connect attempt can never act on a later instance. Temporarily depends on phoenix-channels-python-client's PR #53 branch via a documented [tool.uv.sources] override, since on_heartbeat_ack and close_connection aren't released yet. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011EuhCwwdkV3gYphyNFkw9Y
… spam close_connection failures previously killed the watchdog task silently for the rest of that PHXChannelsClient instance's life, since nothing recreates it across internal reconnects. Also guard against warning/closing when already disconnected (backoff periods produced misleading "forcing reconnect" log spam for a connection that wasn't there to close). Drops a now-dead deadline-changed check made redundant by resetting the deadline before the connection check. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011EuhCwwdkV3gYphyNFkw9Y
- Drop the eager watchdog task creation in __aenter__'s retry loop; only create it after a connect attempt actually succeeds, so a flaky initial connect no longer creates and immediately cancels a task per attempt. The except branch's cancel call is now dead (no watchdog exists yet on a failed attempt) and is removed. - Drop the _on_heartbeat_ack wrapper -- _reset_watchdog_deadline already matches PHXChannelsClient's on_heartbeat_ack signature, so it's passed directly. - Split _watchdog_loop into _sleep_until_watchdog_deadline (timing) and _force_close_if_stale (the close decision), so each has one concern and one short docstring instead of one function documenting three facts. - Use contextlib.suppress(asyncio.CancelledError) in _cancel_watchdog, matching this repo's own established idiom (and the vendored phoenix client's own task-cleanup code) instead of a manual try/except. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011EuhCwwdkV3gYphyNFkw9Y
A reconnect could inherit a deadline set by the watchdog's own disconnected-state polling (which pushes it forward every dead_threshold_s regardless of when the new socket appears), expiring before the fresh connection's first heartbeat cycle completes and killing a healthy connection. Wraps on_reconnect to reset the deadline immediately, mirroring the same reset already done on initial connect and on ack. Also corrects the phoenix-channels-python-client version-floor comment: the PR #53 branch itself is still versioned 0.2.3 (unbumped), so >=0.3.0 was a guessed number, not a verified one -- now documented as a deliberate placeholder, not a real version, with an explicit note against lowering it to silence CI (that would let a broken package install cleanly instead of failing loudly). Widens two real-socket watchdog tests' timing margins (50ms/150ms -> 150ms/400ms) to reduce CI scheduler-flake risk. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011EuhCwwdkV3gYphyNFkw9Y
Dead-connection detection (deadline tracking, force-close decision, task lifecycle) was mixed into WebSocketClient alongside topic subscriptions, event routing, and reconnect wiring. Pulls it into band.client.streaming.watchdog.HeartbeatWatchdog: WebSocketClient now holds one `self._watchdog` and delegates start/stop/reset_deadline to it instead of owning 3 attributes and 5 methods directly. Watchdog-internal behavior (instance-binding, close-failure recovery, disconnected-state no-op, reconnect deadline reset) moves to a new tests/websocket/test_watchdog.py testing HeartbeatWatchdog directly -- faster and more isolated than reaching into WebSocketClient's private state. test_client.py keeps only the tests that exercise the real wiring (heartbeat_interval_s/on_heartbeat_ack/on_reconnect passed to PHXChannelsClient, __aexit__ stopping the watchdog) against a real in-process Phoenix peer. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011EuhCwwdkV3gYphyNFkw9Y
The temporary git-source override pointed at the personal fork (AlexanderZ-Band/phoenix-channels-python-client), whose branch was deleted once PR #53 was reopened as a same-repo PR (#54) with write access. Point at band-ai/phoenix-channels-python-client instead, and bump the pinned commit to the branch's current HEAD (2b1fe6d), which includes fixes landed after the fork was cut over: close_connection's reason-length guard and the forced-close reconnect-classification fix. Still temporary until #54 merges and releases; the floor and this table get resolved then, same as before.
Removes the [tool.uv.sources] override early rather than waiting for PR band-ai/phoenix-channels-python-client#54 to merge and release. The >=0.3.0 floor now resolves against PyPI as a real consumer would see it, which currently fails (only <=0.2.3 is published) -- expected, and self-resolving the moment #54 releases 0.3.0, with no further edit needed here. uv.lock is intentionally left unregenerated (uv lock cannot resolve without the override); it will be regenerated once #54 releases.
PR #54 released as 0.2.4, not the >=0.3.0 placeholder floor (release-please bumps patch, not minor, pre-1.0). Point the dependency floor at the real PyPI release and drop the resolver's git pin now that it's unnecessary. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TW3zAmveYg4H7YfkWemQJ9
test_aenter_retries_unclassified_initial_connection_errors patches asyncio.sleep globally to capture backoff delays. The watchdog started on a successful connect calls the same patched sleep in a tight deadline-check loop with no real yield point, and the test never stopped it -- on the session-scoped asyncio loop this froze whichever later test's task next got scheduled, failing Windows CI only (its ProactorEventLoop apparently starves worse under the spin). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TW3zAmveYg4H7YfkWemQJ9
amit-gazal-band
left a comment
There was a problem hiding this comment.
Two test-isolation issues from the new HeartbeatWatchdog, neither blocking merge but worth a fast follow-up.
tests/websocket/test_client.py:417 (pre-existing line, not touched by this diff, so GitHub won't let me anchor an inline comment there): test_aenter_retries_unclassified_initial_connection_errors does monkeypatch.setattr("band.client.streaming.client.asyncio.sleep", fake_sleep), which patches the real asyncio module globally, not just this module's usage. The new HeartbeatWatchdog also calls asyncio.sleep internally (in _sleep_until_deadline), so it gets patched too — and since fake_sleep returns instantly, the watchdog's sleep loop busy-spins forever once it starts. Confirmed by running the test directly: it now takes ~32s instead of near-instant, and logs an unhandled Task exception was never retrieved / pytest-timeout error rooted in HeartbeatWatchdog._loop.
test_aenter_restores_reconnect_after_successful_initial_connect calls __aenter__ without a matching __aexit__, leaking a watchdog task against a test double with no real `.connection`. Once SessionPolicy.default()'s dead_threshold_s elapses, _force_close_if_stale raises AttributeError inside that fire-and-forget task, on the session-scoped event loop, polluting whatever test is running when it fires -- same class of leak already fixed for the retry-backoff test.
|
@amit-gazal-band addressed the review comment in de8673a — pushed and thread resolved. Full unit suite green (5175 passed). Ready for another look. |
…eartbeat-interval-dead-threshold-polic-INT-1323 # Conflicts: # pyproject.toml # uv.lock
amit-gazal-band
left a comment
There was a problem hiding this comment.
Verified Alex's fixes cover both review comments — the watchdog leak in client.py:431 and the monkeypatched asyncio.sleep spin in the retry-backoff test are both resolved (each affected test now stops the watchdog via __aexit__). Re-ran tests/websocket/test_client.py + tests/websocket/test_watchdog.py: 54 passed in 3.8s, no warnings or leaked tasks.
amit-gazal-band
left a comment
There was a problem hiding this comment.
Reviewed and confirmed: the leaked-watchdog-task fix in de8673a resolves the review comment — verified locally (tests/websocket/test_client.py, 49/49 passed, no leaked-task warnings). Note this PR is still marked not-mergeable pending the phoenix-channels-python-client dependency release per the PR description.
Summary
Part of INT-1323: wires
WebSocketClientup toband-sdk-core's sharedSessionPolicyheartbeat/dead-threshold policy (band-ai/band-sdk-core#60, released as v2.0.0) andphoenix-channels-python-client's newon_heartbeat_ack/close_connectionsurface (band-ai/phoenix-channels-python-client#53).Dead-connection detection lives in a new
band.client.streaming.watchdog.HeartbeatWatchdogclass, not inline inWebSocketClient:HeartbeatWatchdogowns the deadline, the watchdog task, and the force-close decision.WebSocketClientholds oneself._watchdogand delegatesstart/stop/reset_deadlineto it.heartbeat_interval_s(fromSessionPolicy, default or an injected override) is passed toPHXChannelsClient;on_heartbeat_ackandon_reconnectare wired to reset the watchdog's deadline.HeartbeatWatchdog.start(client)takes the specificPHXChannelsClientinstance as an argument (not read dynamically later), so a stale watchdog from a superseded initial-connect attempt can never act on a later instance.dead_threshold_spasses with no ack, flows into the client's own existing reconnect path — no new disconnect-handling branch.__aexit__.close_connectionfailure is caught and logged rather than killing the watchdog task, and the watchdog skips warning/closing when there's no live connection to avoid misleading log spam during backoff.Dependency floor: resolved
phoenix-channels-python-clientPR #53 has since released as a real PyPI version (0.2.4) containingon_heartbeat_ack/close_connection.dependenciespinsphoenix-channels-python-client>=0.2.4— a real, verified floor — and there is no[tool.uv.sources]git override. This PR is mergeable/releasable; thepackagingjob's "Install from wheel" step passes.Addressed review feedback (amit-gazal-band)
test_aenter_restores_reconnect_after_successful_initial_connectcalled__aenter__without a matching__aexit__, leaking a watchdog task against a test double with no real.connection— oncedead_threshold_selapsed,_force_close_if_stalewould raiseAttributeErrorinside that fire-and-forget task on the session-scoped event loop, polluting whichever test ran when it fired. Gave the double a no-op__aexit__and added the matching cleanup call, same shape as the sibling retry-backoff test's own earlier fix. Checked every other bare__aenter__()call in the file: the rest all raise before reaching the watchdog-start line, so this was the only actual leak.Review thread resolved.
Test plan
tests/websocket/test_watchdog.py— 5 tests againstHeartbeatWatchdogdirectly, no real socket needed: task cancels cleanly onstop(); a stale watchdog from a superseded attempt cannot close a newer connection; survives aclose_connectionfailure and keeps enforcing the threshold; neither warns nor closes while already disconnected; a reset deadline lets a fresh socket survive a reconnect that lands late in a disconnected-state polling window.tests/websocket/test_client.py— real-wire tests against an in-process Phoenix-protocol peer, covering the actual wiring: ack keeps the connection alive across several heartbeat cycles; a withheld ack forces close + reconnect at ~dead_threshold_s;__aexit__stops the watchdog; the successful-initial-connect test no longer leaks a watchdog task (-W errorclean).uv run pytest tests/ --ignore=tests/integration/ --ignore=tests/e2e/ --all-packages— 5175 passed, 122 skipped.uv run pyrefly check(project-wide, matching CI) — 0 errors.uv run ruff check/ruff format --check— clean.mainto pick up an unrelatedband-sdk-core/lockfile drift; resolved keeping this PR'sband-sdk-core==2.0.0andphoenix-channels-python-client>=0.2.4. Full CI green (20/20 checks) as of the merge commit, includingpackaging.🤖 Generated with Claude Code
https://claude.ai/code/session_011EuhCwwdkV3gYphyNFkw9Y