connect: heartbeat can race the init handshake on reconnect, causing 1008 connect_worker_hello_invalid_msg

[kuba415]

Summary

On a reconnect, _heartbeat_sender can send WORKER_HEARTBEAT on the new WebSocket before the
init handshake has sent WORKER_CONNECT. The gateway closes the connection with
1008 connect_worker_hello_invalid_msg.

The SDK's own comment predicts the symptom exactly:

# We can't send heartbeats until the handshake is complete. Doing so
# isn't fatal, but will result in "connect_worker_hello_invalid_msg" logs
await self._state.init_handshake_complete.wait_for(True)

Verified against 0.5.18.

Why the guard doesn't hold

State.close_ws() clears the socket but leaves init_handshake_complete at its previous value
(_internal/models.py):

def close_ws(self) -> None:
    if self.allow_reconnect():
        self.conn_state.value = ConnectionState.RECONNECTING
    self.conn_init.value = None
    self.ws.value = None
    # init_handshake_complete is not reset

The flag is only reset once the handshake handler receives GATEWAY_HELLO on the new connection
(_internal/init_handshake_handler.py):

if msg.kind != connect_pb2.GatewayMessageType.GATEWAY_HELLO:
    ...
# Reset because we were told to redo the initial handshake
self._state.init_handshake_complete.value = False

So there is a window on every reconnect:

1. close_ws() — ws = None, init_handshake_complete still True
2. the new socket is established — state.ws.value is set
3. _heartbeat_sender — wait_for(True) passes on the stale flag, wait_for_not_none()
   returns the new socket, and it sends WORKER_HEARTBEAT
4. only later does GATEWAY_HELLO arrive, reset the flag, and trigger WORKER_CONNECT

Between (2) and (4) the heartbeat can reach the gateway first, which is what produces the close.

Observed

A worker that had been connected for ~13 hours (a gateway drain, not a deploy):

18:22:19.544  POST /v0/connect/start 200
18:22:19.629  Connection closed abnormally      <- 85ms later, i.e. pre-handshake
18:22:20.654  POST /v0/connect/start 200        <- recovered in 1.03s
19:03:53.606  POST /v0/connect/start 200        <- later reconnect, clean

It is self-healing — the reconnect loop recovers in about a second — so the practical impact is
noise rather than downtime. But it is noisy enough to be the loudest thing in our error tracking,
which is its own cost: it trains people to ignore that connection's errors.

A genuinely invalid hello would produce the same string, so the two are hard to tell apart from the
message alone.

Suggested fix

Reset the flag where the connection it describes is torn down, so it can't outlive it:

def close_ws(self) -> None:
    if self.allow_reconnect():
        self.conn_state.value = ConnectionState.RECONNECTING
    self.conn_init.value = None
    self.ws.value = None
    self.init_handshake_complete.value = False

Happy to open a PR if that looks right to you.