feat(sdk-rust): reconnect a dropped WebSocket stream - #270
Conversation
The Rust SDK's WebSocket client had no reconnect logic, so a stream that dropped for any reason — server restart, idle NAT rebind, transient network loss — ended permanently: `recv()` returned `None` and the caller silently stopped receiving. The TypeScript SDK has handled this since it shipped, via `reconnect` / `reconnectInterval` / `maxReconnectAttempts`. Adds `ReconnectPolicy` (defaults matching the TS options: enabled, 3s, 10 attempts) and re-dials inside `WebSocketConnection::recv`, so the recovery is transparent to callers. `TinyPlaceWebSocket::with_reconnect` overrides it; `ReconnectPolicy::disabled()` restores the previous end-on-drop behaviour. Each attempt re-runs the full dial rather than replaying the original request. This is required, not incidental: `sign_websocket_upgrade` binds a timestamp and nonce, so a replayed upgrade presents a stale credential and is rejected with a 401. Two deliberate deviations from the TS implementation: - `connect()` stays a single attempt, so a misconfigured or unauthorised client still fails fast instead of retrying a 401 ten times. The policy governs only an already-established stream that later drops. - The attempt counter resets on a delivered message rather than on open. TS resets in `onopen`, which lets a server that accepts and immediately drops be retried forever. Covered by five tests against a local server: reconnect after a drop, the disabled opt-out, giving up after max_attempts, default parity with the TS SDK, and an assertion that the second upgrade carries a freshly signed credential rather than a replayed one.
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe Rust WebSocket SDK adds configurable automatic reconnection. Dropped streams are re-dialed with bounded retries, pauses, timeouts, and freshly signed upgrade credentials. Tests cover recovery, keepalive handling, disabled and exhausted reconnects, stalled handshakes, defaults, and signature renewal. ChangesWebSocket reconnect support
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant WebSocketConnection
participant TinyPlaceWebSocket
participant WebSocketServer
Client->>WebSocketConnection: recv()
WebSocketServer-->>WebSocketConnection: Close or transport error
WebSocketConnection->>TinyPlaceWebSocket: wait and dial()
TinyPlaceWebSocket->>WebSocketServer: fresh signed upgrade
WebSocketServer-->>TinyPlaceWebSocket: new connection
WebSocketConnection-->>Client: next data frame
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0d64817f17
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| } | ||
| return None; | ||
| } | ||
| Some(Ok(_)) => continue, |
There was a problem hiding this comment.
When a reconnected stream remains idle, its keepalive ping/pong frames reach this arm but never reset attempts; only JSON text or binary frames do. The server explicitly uses keepalive pings for live streams (sdk/typescript/src/websocket.ts:135-140), so ten unrelated outages separated by otherwise healthy but quiet connections—such as an inbox with no new items—eventually exhaust max_attempts and permanently end recv(). Reset the budget after evidence that the replacement connection is healthy, such as a keepalive/stability threshold, while still bounding sockets that immediately drop.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch — fixed in e20de4f.
This was a real bug and the most important of the four findings. My doc comment claimed the budget bounded consecutive failures, but since only Text/Binary reset it, that was only true on a stream carrying data. On a quiet one — an inbox with no new items, which is a completely normal state — the server keepalives were the only traffic, so unrelated outages days apart accumulated until max_attempts was spent and recv() ended permanently.
Fix is the minimal one you suggested: any frame counts as evidence of health, so control frames reset the budget too (mark_healthy). A socket that produces nothing at all before dropping still counts against the budget, so the anti-flap property I wanted is preserved — a server that accepts and immediately drops stays bounded.
Regression test: ws_keepalive_traffic_resets_the_retry_budget, with max_attempts: 1 and a server that pings-then-closes twice before delivering data. Reaching the third connection is only possible if the pings reset the budget. Verified it fails without the fix.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
sdk/rust/src/websocket.rs (1)
272-289: 🧹 Nitpick | 🔵 TrivialFixed-interval retries with no jitter can synchronize reconnect storms.
All clients disconnected by the same server-side event will redial at the exact same fixed cadence (default 3s), which can produce thundering-herd load spikes against the backend right as it's recovering. This mirrors the TS SDK by design, so it's advisory rather than a blocker — worth keeping in mind if this SDK is used by many concurrent clients against a shared backend.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sdk/rust/src/websocket.rs` around lines 272 - 289, Update reconnect to add randomized jitter to each retry delay, preventing clients from redialing on the same fixed cadence; preserve the existing policy.enabled, max_attempts, interval pause-before-attempt behavior, and successful dial handling in reconnect.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@sdk/rust/src/websocket.rs`:
- Around line 272-289: Update reconnect to enforce a timeout for every
self.dialer.dial() attempt, using tokio::time::timeout with the existing
reconnect policy duration or an appropriate per-attempt timeout. Handle timeout
and dial errors as failed attempts so the while loop still respects
policy.max_attempts and returns false when no connection is established.
---
Nitpick comments:
In `@sdk/rust/src/websocket.rs`:
- Around line 272-289: Update reconnect to add randomized jitter to each retry
delay, preventing clients from redialing on the same fixed cadence; preserve the
existing policy.enabled, max_attempts, interval pause-before-attempt behavior,
and successful dial handling in reconnect.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 374b4d8f-4bae-4f5a-9086-743e0a79c27e
📒 Files selected for processing (3)
sdk/rust/src/lib.rssdk/rust/src/websocket.rssdk/rust/tests/websocket.rs
|
| Filename | Overview |
|---|---|
| sdk/rust/src/websocket.rs | Core reconnect implementation: adds ReconnectPolicy, dial(), reconnect(), and mark_healthy() to WebSocketConnection. Logic is sound and well-commented; the only open concern (already in review threads) is that Message::Close is matched without inspecting the frame code, so permanent-error codes still trigger the full retry sequence. |
| sdk/rust/src/lib.rs | One-line addition to re-export ReconnectPolicy from the crate root alongside the existing WebSocket types. |
| sdk/rust/tests/websocket.rs | Seven new integration tests covering the happy path, opt-out, budget exhaustion, keepalive-based budget reset, stalled-handshake timeout, and fresh-credential signing. Tests use short timeouts to keep CI fast and use server.abort() correctly for the stall test. |
Sequence Diagram
sequenceDiagram
participant Caller
participant recv as WebSocketConnection::recv
participant reconnect as reconnect()
participant dialer as TinyPlaceWebSocket::dial()
participant Server
Caller->>recv: recv()
recv->>Server: inner.next()
Server-->>recv: Message::Text / Binary
recv->>recv: "mark_healthy() [attempts=0]"
recv-->>Caller: Some(Ok(json))
Caller->>recv: recv()
recv->>Server: inner.next()
Server-->>recv: Message::Close(_) or None
recv->>reconnect: reconnect()
loop "while attempts < max_attempts"
reconnect->>reconnect: "attempts += 1, sleep(interval)"
reconnect->>dialer: dial() [fresh sign]
dialer->>dialer: sign_websocket_upgrade() [new timestamp+nonce]
dialer->>Server: connect_async(signed_request)
alt dial succeeds
Server-->>dialer: 101 Switching Protocols
dialer-->>reconnect: Ok(stream)
reconnect->>reconnect: "inner = stream, last_dial_error = None"
reconnect-->>recv: true
recv->>Server: inner.next() [new stream]
Server-->>recv: Message::Text
recv->>recv: "mark_healthy() [attempts=0]"
recv-->>Caller: Some(Ok(json))
else dial fails or times out
dialer-->>reconnect: Err / Timeout
reconnect->>reconnect: "last_dial_error = Some(error)"
end
end
reconnect-->>recv: false [budget exhausted]
recv-->>Caller: None (Close arm) or Some(Err(last_dial_error)) (Err arm)
Reviews (2): Last reviewed commit: "fix(sdk-rust): bound the reconnect dial ..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| Some(Ok(Message::Close(_))) | None => { | ||
| if self.reconnect().await { | ||
| continue; | ||
| } | ||
| return None; |
There was a problem hiding this comment.
Message::Close(_) is matched without inspecting the close frame's code, so a server-initiated close carrying a permanent-error code (e.g. 4001 Unauthorized, 4003 Forbidden, 1008 Policy Violation) triggers up to max_attempts reconnect attempts at full interval spacing — 30 seconds of retries by default for a token that is already expired. Only transient codes like 1001 (Going Away) or an unexpected 1006 (Abnormal Closure) are good candidates for reconnect; permanent ones will all fail for the same reason and just delay the caller in noticing.
A minimal guard would match Some(Ok(Message::Close(Some(frame)))) and skip reconnect when frame.code is in the application-error range (4000–4999) or other well-known permanent codes before falling through to the current reconnect path.
There was a problem hiding this comment.
Thanks — accurate about the code, but I am going to decline this one, and I want to be explicit about why so a maintainer can overrule me.
The specific example does not hold for this SDK. Every reconnect re-runs the full dial and re-signs the credential (sign_websocket_upgrade binds a fresh timestamp and nonce), so an expired token is precisely the case where retrying is correct — the new signature may well succeed. That was the main argument for the guard.
More importantly, I do not know what close codes this backend actually sends. The auth rejection path documented in e2e_docker.rs is a 401 on the upgrade, which never produces a close frame at all — it fails in dial() and burns an attempt. I could not find a case where the backend emits an app-range close code. Implementing an allowlist against guessed codes risks the worse failure: not reconnecting when we should, which is silent.
The cost of being wrong in the current direction is bounded and visible — at most max_attempts × interval of futile retries. The cost of a wrong allowlist is a stream that quietly stops.
Happy to add this the moment the backend documents its close-code semantics, or if someone who knows them wants to specify the set.
Three fixes from review feedback on the reconnect path. Keepalive traffic now resets the retry budget. Only text and binary frames did, so a quiet stream — an inbox with no new items, where the server's pings are the only traffic — never reset it. Ten outages days apart would exhaust `max_attempts` cumulatively and end `recv()` for good, which is the opposite of what the budget is for. Any frame is proof the replacement socket works; a socket that produces nothing at all before dropping still counts against the budget, so a server that accepts and immediately drops stays bounded. Each re-dial is now bounded by `ReconnectPolicy::connect_timeout` (default 10s). `connect_async` has no timeout of its own, so a black-holed connect or a handshake the server never answers hung `recv` indefinitely and left `max_attempts` bounding nothing. This is reachable only through the reconnect loop, where the caller cannot wrap it themselves; `connect()` keeps its current semantics. Exhausting the budget now reports why the *last* re-dial failed rather than the older error that started the run, which is what describes the current state of the world. Both behavioural fixes have regression tests, each verified to fail without its fix (the stalled-handshake one by hanging until the test's own deadline).
1 participant
Problem
The Rust SDK's WebSocket client has no reconnect logic. A stream that drops for any reason — server restart, idle NAT rebind, transient network loss — ends permanently:
recv()returnsNoneand the caller silently stops receiving, with nothing distinguishing "the server closed this stream" from "the network blipped".The TypeScript SDK has handled this since it shipped (
websocket.ts), viareconnect/reconnectInterval/maxReconnectAttempts. This closes that parity gap.It matters most for long-lived consumers — an agent sitting on
/inbox/streamor/a2a/{id}/streamfor hours is exactly the case where a single blip currently ends the subscription for good, and the agent has no way to notice.Change
ReconnectPolicy { enabled, interval, max_attempts }, defaults matching the TS options: enabled, 3s, 10 attempts.WebSocketConnection::recvre-dials on close or transport error, so recovery is transparent to callers.TinyPlaceWebSocket::with_reconnect(policy)to override;ReconnectPolicy::disabled()restores the previous end-on-drop behaviour.Each attempt re-runs the full dial rather than replaying the stored request. This is required rather than incidental:
sign_websocket_upgradebinds a timestamp and nonce, so a replayed upgrade presents a stale credential and gets a 401. There's a test pinning this.Two deliberate deviations from the TS implementation
Flagging these explicitly since they're judgement calls, and I'm happy to change either:
connect()stays a single attempt. In TS an initial failure both rejects and schedules a retry. Here a refused or unauthorised upgrade returns the error immediately, so a misconfigured client fails fast instead of retrying a 401 ten times over 30s. The policy governs only an already-established stream that later drops.reconnectCountinonopen, which means a server that accepts and immediately drops is retried forever. Resetting on a delivered frame keepsmax_attemptsa real bound.Behaviour change
Reconnect defaults to on, matching TS. Callers that today treat
recv() -> Noneas terminal still getNone, just after the policy is exhausted rather than immediately. Anyone wanting the old semantics can opt out withReconnectPolicy::disabled().Worth noting a reconnect re-subscribes: messages published while the socket was down are not replayed, and a stream that opens with a snapshot frame delivers a fresh one. Documented on
recv.If you'd rather ship this opt-in (default
enabled: false) to avoid any semantic change on a patch release, that's a one-line change — say the word.Tests
Five new tests in
sdk/rust/tests/websocket.rs, all offline against a localtokio-tungsteniteserver:ws_recv_reconnects_after_the_server_drops_the_streamws_recv_reports_the_drop_when_reconnecting_is_disabledws_recv_gives_up_after_max_attemptsrecvdoesn't hangws_reconnect_defaults_match_the_typescript_sdkTinyPlaceWebSocketOptionsws_reconnect_signs_a_fresh_upgradev1:credential, not a replayValidation
Run from
sdk/rust, matching the CI job:cargo fmt --check✅cargo clippy --all-targets -- -D warnings✅cargo clippy --all-targets --features cli -- -D warnings✅cargo test✅ (websocket suite 16 passed, full suite green)cargo test --features cli✅cargo test --doc✅Not run: the
#[ignore]de2e_dockersuite, which needs the compose stack or staging.Summary by CodeRabbit
New Features
ReconnectPolicy(enabled/interval/max attempts/connect timeout) and can be applied withwith_reconnect.Bug Fixes
Tests