fix(ws): TimeoutError wrap + zombie-subscription cleanup on gap-recovery failure

[TexasCoding]

Summary

Two HIGH-severity WS integrity fixes from the Round-3 audit. ChannelsManager._wait_for_response now wraps asyncio.wait_for's bare TimeoutError in a structured KalshiSubscriptionError so callers retain the channel / client_id / op recovery surface promised by #213. SubscriptionManager.broadcast_error now pops the dead Subscription from _subscriptions / _sid_to_client so a failed gap-recovery (or backpressure rollback) doesn't leave a zombie sub that the next reconnect's resubscribe_all resurrects on the server — preventing silent data loss into a closed queue and server-quota leak.

Issues closed

• Closes _wait_for_response leaks TimeoutError instead of KalshiSubscriptionError #314 — _wait_for_response leaks TimeoutError instead of KalshiSubscriptionError
• Closes Failed gap-recovery leaves zombie Subscription that resurrects on next reconnect #315 — failed gap-recovery leaves zombie Subscription that resurrects on next reconnect

Behavioral changes

• A subscribe/unsubscribe/update timeout now raises KalshiSubscriptionError (with channel/client_id/op populated) instead of the bare TimeoutError it used to leak from asyncio.wait_for.
• After broadcast_error is called for a subscription (gap-recovery failure, backpressure rollback), the subscription is removed from _subscriptions / _sid_to_client. The next reconnect's resubscribe_all no longer re-establishes it on the server, and unsubscribe(client_id) no longer needs to short-circuit on a stuck entry.

Tests

• tests/ws/test_client.py::TestIssue314WaitForResponseTimeout::test_issue_314_subscribe_raises_kalshi_subscription_error_on_timeout — drives _wait_for_response against a connection whose recv() hangs forever and asserts the structured KalshiSubscriptionError surfaces with channel/client_id/op.
• tests/ws/test_client.py::TestIssue315ZombieSubscriptionCleanup::test_issue_315_failed_resubscribe_removes_zombie_subscription — simulates the _handle_seq_gap failure path and asserts _subscriptions/_sid_to_client are clean and a subsequent resubscribe_all sends nothing.
• tests/ws/test_client.py::TestIssue315ZombieSubscriptionCleanup::test_issue_315_backpressure_path_also_removes_subscription — covers the KalshiBackpressureError rollback path through the same broadcast_error.

Source

Round-3 independent audit closure plan, wave W1.

[claude]

Code Review — PR #364

Overall: Two clean, targeted bug fixes with good test coverage. The root causes are correctly identified and the fixes are minimal. A few things worth discussing below.

---

Fix 1 — _wait_for_response TimeoutError wrapping (channels.py:213-221)

Correct diagnosis. When asyncio.wait_for fires, it cancels the inner coroutine and raises TimeoutError in the outer scope — the remaining <= 0 guard at the top of the next iteration never runs. The fix is exactly right.

from None choice. Using from None suppresses the original TimeoutError from the chain. This avoids surfacing asyncio internals to callers, which is the right instinct, but it does make it harder to distinguish a legitimate timeout from an unexpected TimeoutError raised by recv() itself for some other reason. Consider from e to preserve context while still surfacing the structured KalshiSubscriptionError. Minor — not a blocker.

# type: ignore[arg-type] on op=op. Consistent with the identical annotation on line 196 — fine.

---

Fix 2 — Zombie subscription cleanup in broadcast_error (channels.py:597-600)

Correct and safe. Both dict.pop(..., None) calls are idempotent, so double-calls (e.g., from concurrent error paths) are harmless. The server_sid is not None guard correctly handles the gap-recovery path where resubscribe_one has already cleared the sid before raising.

One edge case to verify: resubscribe_all iterates over old_subs = dict(self._subscriptions) (a copy), then on per-sub failure calls self._subscriptions.pop(client_id, None) directly — it does not call broadcast_error. So there's no double-pop between these two paths. Confirmed safe.

---

Tests

TestIssue314... — solid. NeverRecv.recv() hanging on asyncio.Event().wait() is a clean way to force wait_for to time out. Assertions cover channel, client_id, op, and the msg_id string — good surface area.

TestIssue315...test_issue_315_failed_resubscribe_removes_zombie_subscription — one concern:

sub.server_sid = None          # ← cleared by test setup
mgr._sid_to_client.pop(314, None)  # ← also cleared by test setup
await mgr.broadcast_error(7, ...)
assert 314 not in mgr._sid_to_client  # ← passes vacuously, not tested by broadcast_error

Because sub.server_sid is None when broadcast_error runs, the new _sid_to_client.pop line is never executed in this test. The assertion passes because the test already removed the entry, not because the code under test removed it. The _subscriptions cleanup is genuinely tested here.

The _sid_to_client cleanup path is covered by the backpressure test (where server_sid = 999 is still set), but it's worth noting the gap-recovery test doesn't fully exercise the _sid_to_client branch of the fix.

TestIssue315...test_issue_315_backpressure_path_also_removes_subscription — clean and does correctly exercise the _sid_to_client.pop path.

resubscribe_all no-op assertion — the recorded == [] check and _subscriptions == {} assertion are a nice end-to-end confirmation that the zombie cleanup is complete.

---

Style / CLAUDE.md

• The multi-line # comment block added inside _wait_for_response (6 lines) and the paragraph added to the broadcast_error docstring both exceed the project's "one short line max" guideline for comments. The inline # #315: drop the dead subscription so resubscribe_all skips it. is fine; the block comment explaining the asyncio internals is more than needed at the call site — that level of detail belongs in the PR description (already there) or the referenced issue.

---

Summary

Area	Status
Bug fix correctness	✅ Both fixes are correct
Minimal footprint	✅ Only the necessary lines changed
Test coverage	✅ with one caveat: _sid_to_client branch in gap-recovery test passes vacuously
Comment style	⚠️ Multi-line comment blocks exceed project guidelines
from None suppression	⚠️ Minor — consider from e for better debuggability

No blockers. The _sid_to_client test observation is the most worth addressing before merge, as it could give false confidence that the gap-recovery path is fully covered.

[TexasCoding]

Round-2 review addressed in 09ea50a:

• Reworked test_issue_315_failed_resubscribe_removes_zombie_subscription so sub.server_sid is non-None when broadcast_error runs — now meaningfully exercises the _sid_to_client.pop branch instead of asserting vacuously.
• Trimmed multi-line comment block in _wait_for_response and broadcast_error docstring per CLAUDE.md (one short line).

[claude]

Code Review

Overview

This PR closes two HIGH-severity WS integrity issues: a bare TimeoutError escape from _wait_for_response (#314) and zombie subscriptions left behind by broadcast_error after gap-recovery or backpressure failures (#315). Both fixes are surgical and well-targeted.

---

Fix #314 — _wait_for_response TimeoutError wrap

kalshi/ws/channels.py lines 209–216

The placement is correct — except TimeoutError is ordered after except ConnectionClosed (a websockets subclass), avoiding any accidental shadowing. The structured fields (channel, client_id, op) are properly forwarded, consistent with the existing manual-deadline timeout path at lines 192–197.

Minor concern — from None vs from e:

except TimeoutError:
    raise KalshiSubscriptionError(...) from None   # ← suppresses chain entirely

from None clears the exception context, so a KalshiSubscriptionError traceback won't show the underlying TimeoutError. This is intentional (avoids asyncio internal noise) but it does strip context that could help distinguish a true network timeout from, say, an event loop stall. from e would keep the context without meaningfully polluting the message. Either is defensible — just flagging for awareness.

Pre-existing # type: ignore[arg-type]:

The new code at line 215 repeats the same suppress already on line 196. The root cause is that _wait_for_response accepts op: str | None while KalshiSubscriptionError.op is Literal["subscribe", "unsubscribe", "update_subscription"] | None. This PR correctly follows the existing pattern rather than widening the issue — but narrowing _wait_for_response's op parameter to the same Literal type would eliminate both ignores. Worth tracking as tech debt if not already.

---

Fix #315 — broadcast_error zombie cleanup

kalshi/ws/channels.py lines 586–589

The cleanup logic is correct:

await sub.queue.put_error(exc)
# #315: drop the dead subscription so resubscribe_all skips it.
self._subscriptions.pop(client_id, None)
if sub.server_sid is not None:
    self._sid_to_client.pop(sub.server_sid, None)

• put_error before cleanup is the right order — the queue is marked closed first, so any subsequent frame dispatch via _sid_to_client that arrives before the pop hits a closed queue (no-op), not a live one.
• The server_sid is not None guard is essential: a subscription that never received its ack has no entry in _sid_to_client, and the naked pop would panic without it. Good catch.
• Double-call safety is preserved: a second broadcast_error on the same client_id hits the early-return at line 577–578 since _subscriptions.get returns None.

---

Tests

The test coverage is solid and addresses the three relevant paths:

Test	Verdict
TestIssue314 — subscribe timeout surfaces KalshiSubscriptionError with correct fields	✅ Good — NeverRecv.recv() via asyncio.Event().wait() cleanly simulates a hung connection
TestIssue315 gap-recovery — both maps cleared, resubscribe_all is a no-op	✅ Good — verifies the _subscriptions == {} postcondition explicitly
TestIssue315 backpressure — same cleanup through the KalshiBackpressureError path	✅ Good — confirms broadcast_error is the single cleanup point regardless of error type

One minor nit — timeout=0.05 in test_issue_314:

await mgr._wait_for_response(
    msg_id=42,
    timeout=0.05,   # ← may be tight on a loaded CI runner
    ...
)

50 ms is fine for a mocked connection on a healthy machine but can produce occasional flakes on a heavily loaded CI host. 0.1 or even 0.2 would cost negligible wall time and provide a larger margin. Not a blocker, just a robustness suggestion.

---

Summary

Both fixes are correct, minimal, and well-tested. The only actionable suggestion is the _wait_for_response op parameter narrowing (eliminates two # type: ignore lines) and the optional timeout bump in the regression test. Neither is a blocker.

Verdict: Approve with optional nits.