Python SDK does not reconnect after a normal WebSocket closure

Description

The Python SDK may remain stuck in the connected state and never reconnect when the Engine WebSocket closes normally, for example with close code 1000 or 1001.

This is reproducible when the worker connects to the Engine through an Nginx WebSocket reverse proxy and the Engine is restarted. The equivalent Node.js worker reconnects correctly through the same proxy and endpoint.

A direct ws:// connection may appear to work because restarting the Engine can result in an abnormal TCP/WebSocket closure. In that case, websockets raises ConnectionClosedError, which reaches the SDK's exception handler. Through a reverse proxy, the downstream connection may instead be observed as a normal WebSocket closure, exposing the issue.

Environment

• Python SDK: reproduced with iii-sdk==0.21.4
• Python WebSocket library: websockets==16.0
• Reverse proxy: Nginx
• Comparison: Node.js iii-sdk==0.21.4 reconnects correctly through the same endpoint

I also inspected the published Python wheels for the following versions, and the relevant receive-loop behavior remains unchanged:

• 0.21.5
• 0.21.6
• 0.21.8
• 0.22.0
• 0.22.1
• 0.23.0
• 0.23.1rc6

Nginx configuration

A dedicated WebSocket location is configured approximately as follows:

location /iii-api/ws-server/ {
    proxy_pass http://engine:49134/;
    proxy_http_version 1.1;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection "upgrade";
}

The initial wss:// connection and worker registration succeed. The problem only occurs after the Engine is restarted.

Steps to reproduce

1. Start an III Engine behind an Nginx WebSocket reverse proxy.

2. Start a Python worker using the proxied endpoint:

   python复制

   from iii import register_worker
   
   worker = register_worker("wss://example.com/iii-api/ws-server/")

3. Confirm that the worker connects and registers its functions.

4. Restart the III Engine while leaving Nginx and the Python worker running.

5. Wait for the Engine to become available again.

6. Observe that the Python worker does not reconnect or re-register its functions.

7. Repeat with a Node.js worker using the same proxied endpoint; it reconnects and re-registers correctly.

Expected behavior

After any WebSocket termination, including a normal closure, the Python SDK should:

1. Clear the current WebSocket reference.
2. Change the connection state to disconnected or reconnecting.
3. Schedule reconnection according to ReconnectionConfig.
4. Re-register the worker, functions, triggers, and trigger types after reconnecting.

Actual behavior

The receive loop can finish normally without entering the except websockets.ConnectionClosed block. As a result:

• self._ws remains non-None.
• The connection state may remain connected.
• _schedule_reconnect() is not called.
• The reconnect loop would not proceed anyway because it uses while self._running and not self._ws.
• The process remains alive but the worker is no longer connected or registered with the restarted Engine.

Root cause

The Python SDK currently handles cleanup and reconnection only inside the exception handler:

async def _receive_loop(self) -> None:
    if not self._ws:
        return
    try:
        async for msg in self._ws:
            await self._handle_message(msg)
    except websockets.ConnectionClosed:
        log.debug("Connection closed")
        self._ws = None
        self._set_connection_state("disconnected")
        if self._running:
            self._schedule_reconnect()

However, in websockets==16.0, Connection.__aiter__() handles a normal closure internally:

async def __aiter__(self):
    try:
        while True:
            yield await self.recv()
    except ConnectionClosedOK:
        return

Therefore, close codes such as 1000 and 1001 cause async for to end normally. No ConnectionClosed exception reaches _receive_loop().

The Node.js SDK does not have this gap because its WebSocket close event always invokes onSocketClose(), which clears the socket and schedules reconnection regardless of the close code.

Suggested fix

Move disconnect cleanup and reconnection scheduling into a finally block so that both normal and abnormal receive-loop termination are handled:

async def _receive_loop(self) -> None:
    if not self._ws:
        return

    try:
        async for msg in self._ws:
            await self._handle_message(msg)
    except websockets.ConnectionClosed:
        log.debug("Connection closed")
    finally:
        self._ws = None

        # Preserve the fatal registration state introduced in newer versions.
        if self._fatal_error is None:
            self._set_connection_state("disconnected")

        if self._running:
            self._schedule_reconnect()

For versions without _fatal_error, the corresponding state update can remain unconditional.

It may also be useful to add regression tests for:

1. Normal closure with code 1000.
2. Going-away closure with code 1001.
3. Abnormal network termination.
4. Shutdown cancellation, ensuring that _running == False prevents reconnection.
5. Successful re-registration of functions after reconnecting.

Additional note

Upgrading from Python SDK 0.21.4 to the current 0.23.0 does not appear to resolve this issue because the normal receive-loop completion path is still not handled.