iii-sdk 0.21.8.dev1__tar.gz → 0.21.8.dev3__tar.gz

{iii_sdk-0.21.8.dev1 → iii_sdk-0.21.8.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: iii-sdk
-Version: 0.21.8.dev1
+Version: 0.21.8.dev3
 Summary: III SDK for Python
 Project-URL: Homepage, https://github.com/iii-hq/iii
 Project-URL: Repository, https://github.com/iii-hq/iii
@@ -14,7 +14,7 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Python: >=3.10
-Requires-Dist: iii-helpers==0.21.8.dev1
+Requires-Dist: iii-helpers==0.21.8.dev3
 Requires-Dist: opentelemetry-api>=1.25
 Requires-Dist: pydantic>=2.0
 Requires-Dist: websockets>=12.0

{iii_sdk-0.21.8.dev1 → iii_sdk-0.21.8.dev3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "iii-sdk"
-version = "0.21.8.dev1"
+version = "0.21.8.dev3"
 description = "III SDK for Python"
 authors = [{ name = "III" }]
 license = { text = "Apache-2.0" }
@@ -23,7 +23,7 @@ dependencies = [
     "websockets>=12.0",
     "pydantic>=2.0",
     "opentelemetry-api>=1.25",
-    "iii-helpers==0.21.8.dev1",
+    "iii-helpers==0.21.8.dev3",
 ]
 
 [project.urls]

{iii_sdk-0.21.8.dev1 → iii_sdk-0.21.8.dev3}/src/iii/__init__.py

@@ -5,6 +5,8 @@ from iii_helpers.queue import EnqueueResult
 from .errors import InvocationError
 from .iii import TriggerAction, register_worker
 from .iii_constants import (
+    ConnectionStateCallback,
+    IIIConnectionState,
     InitOptions,
     TelemetryOptions,
 )
@@ -24,6 +26,8 @@ __all__ = [
     # Errors
     "InvocationError",
     # Core
+    "ConnectionStateCallback",
+    "IIIConnectionState",
     "InitOptions",
     "register_worker",
     "TelemetryOptions",

{iii_sdk-0.21.8.dev1 → iii_sdk-0.21.8.dev3}/src/iii/iii.py

@@ -32,6 +32,7 @@ from .format_utils import extract_request_format, extract_response_format
 from .iii_constants import (
     DEFAULT_RECONNECTION_CONFIG,
     MAX_QUEUE_SIZE,
+    ConnectionStateCallback,
     FunctionRef,
     IIIConnectionState,
     InitOptions,
@@ -202,7 +203,11 @@ class III:
         )
         self._reconnect_attempt = 0
         self._connection_state: IIIConnectionState = "disconnected"
+        # Must exist before the auto-connect below: the loop thread calls
+        # _set_connection_state as soon as connect_async starts.
+        self._connection_listeners: list[ConnectionStateCallback] = []
         self._worker_id: str | None = None
+        self._reattach_token: str | None = None
 
         # Background event loop thread
         self._loop = asyncio.new_event_loop()
@@ -232,11 +237,16 @@ class III:
             return
         if self._connection_state == "failed":
             raise ConnectionError(f"Connection to {self._address} failed")
-        self._connected_event.wait(timeout=30)
+        connected = self._connected_event.wait(timeout=30)
         if cast(IIIConnectionState, self._connection_state) == "failed":
             raise ConnectionError(
                 f"Connection to {self._address} failed after max retries"
             )
+        if not connected:
+            log.warning(
+                f"Timed out after 30s waiting for connection to {self._address}; "
+                "continuing to retry in background"
+            )
 
     def shutdown(self) -> None:
         """Gracefully shut down the client, releasing all resources.
@@ -337,8 +347,18 @@ class III:
             )
             log.info(f"Connected to {self._address}")
             await self._on_connected()
-        except (ConnectionError, OSError, TimeoutError, asyncio.TimeoutError) as e:
-            log.warning(f"Connection failed: {e}")
+        except Exception as e:
+            # Catch everything: a connect-time failure must never kill the
+            # reconnect loop. Enumerating exception types here has failed
+            # twice: ConnectionError/OSError missed websockets'
+            # InvalidHandshake (MOT-3931), and that in turn missed
+            # InvalidMessage, raised when the engine-boot listener race
+            # surfaces through iii-network's accept-then-dial as an EOF
+            # between TCP connect and the WS 101 (MOT-3966): the exception
+            # escaped, the connect task died silently, and the worker stayed
+            # a zombie with zero retries. CancelledError is BaseException,
+            # so shutdown cancellation still propagates.
+            log.warning(f"Connection failed: {type(e).__name__}: {e}")
             if self._running:
                 self._schedule_reconnect()
 
@@ -378,6 +398,18 @@ class III:
     async def _on_connected(self) -> None:
         self._reconnect_attempt = 0
         self._set_connection_state("connected")
+        # Reconnect: present the previous engine-assigned identity BEFORE the
+        # registration replay so the engine retires the old connection and the
+        # replay lands on a clean slate instead of racing its cleanup. The
+        # token proves we ARE that worker (ids alone are publicly listable).
+        if self._worker_id:
+            reattach: dict[str, Any] = {
+                "type": MessageType.REATTACH.value,
+                "previous_worker_id": self._worker_id,
+            }
+            if self._reattach_token:
+                reattach["reattach_token"] = self._reattach_token
+            await self._send(reattach)
         # Re-register all (snapshot to avoid mutation from caller thread)
         for trigger_type_data in list(self._trigger_types.values()):
             await self._send(trigger_type_data.message)
@@ -425,7 +457,12 @@ class III:
     async def _send(self, msg: Any) -> None:
         data = self._to_dict(msg)
         if self._ws and self._ws.state.name == "OPEN":
-            log.debug(f"Send: {json.dumps(data)[:200]}")
+            # Redact the reattach secret from debug logs; it authorizes
+            # evicting this worker's previous connection.
+            log_safe = {
+                k: ("***" if k == "reattach_token" else v) for k, v in data.items()
+            }
+            log.debug(f"Send: {json.dumps(log_safe)[:200]}")
             await self._ws.send(json.dumps(data))
         else:
             if len(self._queue) >= MAX_QUEUE_SIZE:
@@ -486,6 +523,7 @@ class III:
         elif msg_type == MessageType.WORKER_REGISTERED.value:
             worker_id = data.get("worker_id", "")
             self._worker_id = worker_id
+            self._reattach_token = data.get("reattach_token")
             log.debug(f"Worker registered with ID: {worker_id}")
 
     def _handle_result(self, invocation_id: str, result: Any, error: Any) -> None:
@@ -807,6 +845,50 @@ class III:
                 self._connected_event.set()
             else:
                 self._connected_event.clear()
+            for handler in list(self._connection_listeners):
+                try:
+                    handler(state)
+                except Exception:
+                    log.exception("Connection state listener raised")
+
+    def add_connection_state_listener(
+        self, handler: ConnectionStateCallback
+    ) -> Callable[[], None]:
+        """Subscribe to connection-state transitions.
+
+        The handler is fired immediately with the current state (on the
+        caller's thread), then once per transition. Transitions fire on the
+        SDK's background event-loop thread: keep handlers fast and do not
+        call sync SDK methods from them (they would raise ``RuntimeError``).
+        Treat calls as state notifications, not edges. A state may rarely
+        be observed twice around subscription. Registering the same handler
+        twice fires it twice. Returns an idempotent unsubscribe function
+        that removes only its own registration.
+
+        Examples:
+            >>> unsubscribe = worker.add_connection_state_listener(
+            ...     lambda state: print(f"engine link: {state}")
+            ... )
+        """
+        self._connection_listeners.append(handler)
+        try:
+            handler(self._connection_state)
+        except Exception:
+            log.exception("Connection state listener raised on initial fire")
+
+        unsubscribed = False
+
+        def unsubscribe() -> None:
+            nonlocal unsubscribed
+            if unsubscribed:
+                return
+            unsubscribed = True
+            try:
+                self._connection_listeners.remove(handler)
+            except ValueError:
+                pass
+
+        return unsubscribe
 
     def get_connection_state(self) -> IIIConnectionState:
         """Return the current WebSocket connection state.
@@ -1394,7 +1476,7 @@ class TriggerAction:
 
     Examples:
         >>> from iii import TriggerAction
-        >>> # The queue must be declared in the iii-queue worker's queue_configs.
+        >>> # The queue must be declared in the queue worker's queue_configs.
         >>> worker.trigger({'function_id': 'process', 'payload': {}, 'action': TriggerAction.Enqueue(queue='jobs')})
         >>> worker.trigger({'function_id': 'notify', 'payload': {}, 'action': TriggerAction.Void()})
     """
@@ -1417,14 +1499,19 @@ class TriggerAction:
 def register_worker(address: str, options: InitOptions | None = None) -> III:
     """Register the worker with a iii instance, returns a connected worker client.
 
-    Blocks until the WebSocket connection is established and ready.
+    Blocks up to 30 seconds for the WebSocket connection to be established.
+    If the engine is not reachable in time, a warning is logged and the
+    client is returned anyway; it keeps retrying in the background and
+    flushes registrations once connected. Use
+    ``add_connection_state_listener`` to observe the actual transition.
 
     Args:
         address: WebSocket URL of the III engine (e.g. ``ws://localhost:49134``).
         options: Optional configuration for worker name, timeouts, reconnection, and OTel.
 
     Returns:
-        A connected III client instance ready to use.
+        An III client instance (connected, or still connecting after the
+        30s wait timed out).
 
     Raises:
         ConnectionError: If the connection fails or exceeds max retries.

{iii_sdk-0.21.8.dev1 → iii_sdk-0.21.8.dev3}/src/iii/iii_types.py

@@ -21,6 +21,7 @@ class MessageType(str, Enum):
     UNREGISTER_TRIGGER_TYPE = "unregistertriggertype"
     TRIGGER_REGISTRATION_RESULT = "triggerregistrationresult"
     WORKER_REGISTERED = "workerregistered"
+    REATTACH = "reattach"
 
 
 class RegisterTriggerTypeMessage(BaseModel):

{iii_sdk-0.21.8.dev1 → iii_sdk-0.21.8.dev3}/src/iii/types.py

@@ -21,6 +21,7 @@ from typing import (
 from iii_helpers.http import HttpInvocationConfig
 from pydantic import BaseModel, ConfigDict
 
+from .iii_constants import ConnectionStateCallback
 from .iii_types import (
     RegisterFunctionMessage,
     RegisterTriggerInput,
@@ -111,6 +112,8 @@ class IIIClient(Protocol):
 
     def unregister_trigger_type(self, trigger_type: RegisterTriggerTypeInput | dict[str, Any]) -> None: ...
 
+    def add_connection_state_listener(self, handler: ConnectionStateCallback) -> Callable[[], None]: ...
+
     def shutdown(self) -> None: ...
 
 

iii_sdk-0.21.8.dev3/tests/test_connection_state_listener.py

@@ -0,0 +1,149 @@
+"""Unit tests for III.add_connection_state_listener (no live engine needed)."""
+
+import time
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+from iii_helpers.observability import ReconnectionConfig
+from websockets.exceptions import InvalidMessage
+
+import iii.iii as iii_module
+from iii import InitOptions
+from iii.iii import III
+
+
+class FakeWebSocket:
+    def __init__(self) -> None:
+        self.state = SimpleNamespace(name="OPEN")
+
+    async def send(self, payload: str) -> None:
+        pass
+
+    async def close(self) -> None:
+        self.state = SimpleNamespace(name="CLOSED")
+
+    def __aiter__(self) -> "FakeWebSocket":
+        return self
+
+    async def __anext__(self) -> Any:
+        raise StopAsyncIteration
+
+
+def _patch_transport(monkeypatch: pytest.MonkeyPatch, connect) -> None:
+    monkeypatch.setattr(iii_module.websockets, "connect", connect)
+    monkeypatch.setattr("iii_helpers.observability.telemetry.init_otel", lambda **kwargs: None)
+    monkeypatch.setattr("iii_helpers.observability.telemetry.attach_event_loop", lambda loop: None)
+    monkeypatch.setattr(iii_module.III, "_register_worker_metadata", lambda self: None)
+
+
+def _connected_client(monkeypatch: pytest.MonkeyPatch) -> III:
+    async def fake_connect(_addr: str, **kwargs: object) -> FakeWebSocket:
+        return FakeWebSocket()
+
+    _patch_transport(monkeypatch, fake_connect)
+
+    client = III("ws://fake")
+    client._wait_until_connected()
+    time.sleep(0.05)
+    assert client.get_connection_state() == "connected"
+    return client
+
+
+def test_immediate_fire_transitions_dedup_and_shutdown(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    client = _connected_client(monkeypatch)
+    events: list[str] = []
+
+    client.add_connection_state_listener(events.append)
+    assert events == ["connected"]  # fired immediately with current state
+
+    client._set_connection_state("reconnecting")
+    client._set_connection_state("connected")
+    client._set_connection_state("connected")  # same state -> no fire
+    assert events == ["connected", "reconnecting", "connected"]
+
+    client.shutdown()
+    assert events == ["connected", "reconnecting", "connected", "disconnected"]
+
+
+def test_raising_listener_isolated_and_unsubscribe_idempotent(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    client = _connected_client(monkeypatch)
+    events: list[str] = []
+
+    def bad(state: str) -> None:
+        raise RuntimeError("boom")
+
+    client.add_connection_state_listener(bad)  # immediate fire raises, is swallowed
+    unsubscribe = client.add_connection_state_listener(events.append)
+    assert events == ["connected"]
+
+    client._set_connection_state("reconnecting")
+    assert events == ["connected", "reconnecting"]  # bad didn't block delivery
+
+    unsubscribe()
+    unsubscribe()  # idempotent
+    client._set_connection_state("connected")
+    assert events == ["connected", "reconnecting"]
+
+    client.shutdown()  # fires "disconnected" into bad -> swallowed, no crash
+    assert events == ["connected", "reconnecting"]
+
+
+def test_duplicate_subscription_unsubscribe_removes_only_its_own(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    client = _connected_client(monkeypatch)
+    events: list[str] = []
+
+    unsub1 = client.add_connection_state_listener(events.append)
+    unsub2 = client.add_connection_state_listener(events.append)
+    assert events == ["connected", "connected"]
+
+    unsub1()
+    unsub1()  # second call must not remove the other registration
+    events.clear()
+    client._set_connection_state("reconnecting")
+    assert events == ["reconnecting"]
+
+    unsub2()
+    client._set_connection_state("connected")
+    assert events == ["reconnecting"]
+
+    client.shutdown()
+
+
+def test_handshake_failure_keeps_retrying(monkeypatch: pytest.MonkeyPatch) -> None:
+    attempts = 0
+
+    async def flaky_connect(_addr: str, **kwargs: object) -> FakeWebSocket:
+        nonlocal attempts
+        attempts += 1
+        if attempts <= 2:
+            # Peer accepted TCP but the WS upgrade died (not an OSError):
+            # must not kill the connect/reconnect task.
+            raise InvalidMessage("did not receive a valid HTTP response")
+        return FakeWebSocket()
+
+    _patch_transport(monkeypatch, flaky_connect)
+
+    client = III(
+        "ws://fake",
+        InitOptions(reconnection_config=ReconnectionConfig(initial_delay_ms=5, max_delay_ms=10)),
+    )
+    states: list[str] = []
+    client.add_connection_state_listener(states.append)
+    client._wait_until_connected()
+    assert client.get_connection_state() == "connected"
+    assert attempts == 3
+    # _set_connection_state sets the connected event before dispatching
+    # listeners, so give the loop thread a bounded window to deliver.
+    deadline = time.time() + 2
+    while "connected" not in states and time.time() < deadline:
+        time.sleep(0.01)
+    assert "connected" in states
+
+    client.shutdown()

iii_sdk-0.21.8.dev3/tests/test_reconnect_sends_reattach.py

@@ -0,0 +1,77 @@
+"""iii-hq/iii#1975 regression: on reconnect the SDK must present the previous
+engine-assigned worker id via a ``reattach`` frame BEFORE replaying its
+registration batch, so the engine retires the old connection and the replay
+lands on a clean slate instead of racing its cleanup. On the first connect,
+with no prior id, no reattach is sent.
+
+Driven directly against ``_on_connected`` (no socket needed): the reattach is
+the first thing sent, gated on a stored ``_worker_id``.
+"""
+
+import asyncio
+
+from iii.iii import III
+from iii.iii_constants import InitOptions
+from iii.iii_types import MessageType
+
+
+def test_reconnect_sends_reattach_first(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "iii_helpers.observability.telemetry.init_otel", lambda **kwargs: None
+    )
+    monkeypatch.setattr(
+        "iii_helpers.observability.telemetry.attach_event_loop", lambda loop: None
+    )
+
+    # The constructor auto-connects on a background thread; it never reaches a
+    # real engine here (unused.test fails to resolve), so it never calls
+    # _on_connected on its own — we drive that directly below. shutdown() in
+    # the finally stops the (non-daemon) loop thread.
+    client = III("ws://unused.test", InitOptions())
+
+    sent: list = []
+
+    async def fake_send(msg) -> None:
+        sent.append(msg)
+
+    monkeypatch.setattr(client, "_send", fake_send)
+    monkeypatch.setattr(client, "_register_worker_metadata", lambda: None)
+    client._ws = None  # _receive_loop returns immediately
+
+    # One binding to replay, so ordering (reattach before replay) is observable.
+    client._triggers["t1"] = {"type": "registertrigger", "id": "t1"}
+
+    async def drive() -> None:
+        # First connect: no stored id yet -> no reattach.
+        await client._on_connected()
+        assert not any(
+            isinstance(m, dict) and m.get("type") == MessageType.REATTACH.value
+            for m in sent
+        ), f"no reattach must be sent on first connect: {sent}"
+
+        # Engine assigns an identity + secret; the socket then drops and
+        # reconnects.
+        client._worker_id = "w-old"
+        client._reattach_token = "tok-old"
+        sent.clear()
+        await client._on_connected()
+
+    loop = asyncio.new_event_loop()
+    try:
+        loop.run_until_complete(drive())
+    finally:
+        loop.close()
+        client.shutdown()
+
+    assert sent, "reconnect must send frames"
+    first = sent[0]
+    assert isinstance(first, dict) and first.get("type") == MessageType.REATTACH.value, (
+        f"reattach must be the first frame on reconnect, got: {sent}"
+    )
+    assert first.get("previous_worker_id") == "w-old"
+    assert first.get("reattach_token") == "tok-old"
+    # And it must precede the replayed registration.
+    assert any(
+        (m.get("type") if isinstance(m, dict) else getattr(m, "type", None)) == "registertrigger"
+        for m in sent[1:]
+    ), f"registration replay must follow the reattach: {sent}"

iii_sdk-0.21.8.dev3/tests/test_reconnect_survives_handshake_eof.py

@@ -0,0 +1,106 @@
+"""MOT-3966 regression: the reconnect loop must survive handshake-level failures.
+
+First-boot conditions: the engine autostarts bundle workers before its worker
+listener is accepting, and the VM's userspace network stack (iii-network)
+accepts the guest's TCP connection before dialing the host — so "nothing
+listening yet" surfaces to the client as an EOF between TCP connect and the
+WS 101. The websockets library raises that as ``InvalidMessage`` (an
+``InvalidHandshake``, NOT a ``ConnectionError``/``OSError``). The SDK's old
+narrow ``except`` let it escape ``_do_connect``, silently killing the connect
+task with zero retries: the worker stayed a zombie for the VM's lifetime.
+"""
+
+import base64
+import hashlib
+import socket
+import threading
+
+from iii.iii import III
+from iii.iii_constants import InitOptions
+from iii_helpers.observability import ReconnectionConfig
+
+WS_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
+
+
+def _minimal_ws_server(listener: socket.socket, eof_first_n: int) -> None:
+    """Accept loop: the first ``eof_first_n`` connections are closed without
+    sending a byte (EOF during the client's handshake — the engine-boot race
+    as seen through iii-network's accept-then-dial). Later connections get a
+    real RFC 6455 101 reply and are held open."""
+    held = []
+    accepted = 0
+    while True:
+        try:
+            conn, _ = listener.accept()
+        except OSError:
+            return  # listener closed: test over
+        accepted += 1
+        if accepted <= eof_first_n:
+            conn.close()
+            continue
+
+        data = b""
+        try:
+            while b"\r\n\r\n" not in data:
+                chunk = conn.recv(4096)
+                if not chunk:
+                    break
+                data += chunk
+            key = ""
+            for line in data.decode("latin1").split("\r\n"):
+                if line.lower().startswith("sec-websocket-key:"):
+                    key = line.split(":", 1)[1].strip()
+            accept = base64.b64encode(
+                hashlib.sha1((key + WS_GUID).encode()).digest()
+            ).decode()
+            conn.sendall(
+                b"HTTP/1.1 101 Switching Protocols\r\n"
+                b"Upgrade: websocket\r\n"
+                b"Connection: Upgrade\r\n"
+                b"Sec-WebSocket-Accept: " + accept.encode() + b"\r\n\r\n"
+            )
+        except OSError:
+            conn.close()
+            continue
+        held.append(conn)  # keep open; the client streams registrations
+
+
+def test_reconnect_survives_handshake_eof(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "iii_helpers.observability.telemetry.init_otel", lambda **kwargs: None
+    )
+    monkeypatch.setattr(
+        "iii_helpers.observability.telemetry.attach_event_loop", lambda loop: None
+    )
+
+    listener = socket.socket()
+    listener.bind(("127.0.0.1", 0))
+    listener.listen(8)
+    port = listener.getsockname()[1]
+    server = threading.Thread(
+        target=_minimal_ws_server, args=(listener, 1), daemon=True
+    )
+    server.start()
+
+    client = III(
+        f"ws://127.0.0.1:{port}",
+        InitOptions(
+            reconnection_config=ReconnectionConfig(
+                initial_delay_ms=100,
+                max_delay_ms=200,
+                jitter_factor=0.0,
+            ),
+        ),
+    )
+    try:
+        # Pre-fix: the first connection's EOF raised InvalidMessage out of
+        # _do_connect, the connect task died silently, no retry ever ran,
+        # and this wait timed out with the client stuck "connecting".
+        client._connected_event.wait(timeout=10)
+        assert client._connection_state == "connected", (
+            f"client never recovered from the handshake EOF "
+            f"(state={client._connection_state!r}; reconnect loop died)"
+        )
+    finally:
+        client.shutdown()
+        listener.close()