Skip to content

feat(websocket): add heartbeat dead-threshold watchdog - #595

Merged
AlexanderZ-Band merged 11 commits into
mainfrom
core-sdk/add-shared-heartbeat-interval-dead-threshold-polic-INT-1323
Aug 30, 2026
Merged

AlexanderZ-Band merged 11 commits into
mainfrom
core-sdk/add-shared-heartbeat-interval-dead-threshold-polic-INT-1323

Conversation

@AlexanderZ-Band

@AlexanderZ-Band AlexanderZ-Band commented Aug 29, 2026

Copy link
Copy Markdown
Collaborator

Summary

Part of INT-1323: wires WebSocketClient up to band-sdk-core's shared SessionPolicy heartbeat/dead-threshold policy (band-ai/band-sdk-core#60, released as v2.0.0) and phoenix-channels-python-client's new on_heartbeat_ack/close_connection surface (band-ai/phoenix-channels-python-client#53).

Dead-connection detection lives in a new band.client.streaming.watchdog.HeartbeatWatchdog class, not inline in WebSocketClient:

  • HeartbeatWatchdog owns the deadline, the watchdog task, and the force-close decision. WebSocketClient holds one self._watchdog and delegates start/stop/reset_deadline to it.
  • heartbeat_interval_s (from SessionPolicy, default or an injected override) is passed to PHXChannelsClient; on_heartbeat_ack and on_reconnect are wired to reset the watchdog's deadline.
  • HeartbeatWatchdog.start(client) takes the specific PHXChannelsClient instance as an argument (not read dynamically later), so a stale watchdog from a superseded initial-connect attempt can never act on a later instance.
  • A force-close, once dead_threshold_s passes with no ack, flows into the client's own existing reconnect path — no new disconnect-handling branch.
  • Reconnects reset the deadline immediately (not just acks and the initial connect) — otherwise a reconnect can inherit a stale deadline from disconnected-state polling and get force-closed before its first heartbeat cycle completes.
  • The watchdog task is started only once, on a successful connect, and stopped cleanly in __aexit__.
  • A close_connection failure 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-client PR #53 has since released as a real PyPI version (0.2.4) containing on_heartbeat_ack/close_connection. dependencies pins phoenix-channels-python-client>=0.2.4 — a real, verified floor — and there is no [tool.uv.sources] git override. This PR is mergeable/releasable; the packaging job's "Install from wheel" step passes.

Addressed review feedback (amit-gazal-band)

  • test_aenter_restores_reconnect_after_successful_initial_connect called __aenter__ without a matching __aexit__, leaking a watchdog task against a test double with no real .connection — once dead_threshold_s elapsed, _force_close_if_stale would raise AttributeError inside 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 against HeartbeatWatchdog directly, no real socket needed: task cancels cleanly on stop(); a stale watchdog from a superseded attempt cannot close a newer connection; survives a close_connection failure 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 error clean).
  • 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.
  • Merged main to pick up an unrelated band-sdk-core/lockfile drift; resolved keeping this PR's band-sdk-core==2.0.0 and phoenix-channels-python-client>=0.2.4. Full CI green (20/20 checks) as of the merge commit, including packaging.

🤖 Generated with Claude Code

https://claude.ai/code/session_011EuhCwwdkV3gYphyNFkw9Y

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
@linear-code

linear-code Bot commented Aug 29, 2026

Copy link
Copy Markdown

INT-1323

AlexanderZ-Band and others added 6 commits August 29, 2026 17:17
… 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.
AlexanderZ-Band and others added 2 commits August 30, 2026 09:29
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
@AlexanderZ-Band
AlexanderZ-Band requested a review from a team August 30, 2026 07:09

@amit-gazal-band amit-gazal-band left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread src/band/client/streaming/client.py
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.
@AlexanderZ-Band

Copy link
Copy Markdown
Collaborator Author

@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 amit-gazal-band left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@AlexanderZ-Band
AlexanderZ-Band added this pull request to the merge queue Aug 30, 2026
Merged via the queue into main with commit 4182215 Aug 30, 2026
21 checks passed

@amit-gazal-band amit-gazal-band left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@AlexanderZ-Band
AlexanderZ-Band deleted the core-sdk/add-shared-heartbeat-interval-dead-threshold-polic-INT-1323 branch August 30, 2026 11:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants