phase 1.5: auto-reconnect on publisher + connect (production gap fix)

both publish() and connect() runners died permanently on WebSocket drop —
hub restart, network blip, transient error all left the publisher dead.
hand-test in last batch surfaced this; now fixed.

implementation:
- ZhubConnection gains _stop_event field (ZhubPublication already had one)
- publish().runner: outer reconnect loop with exponential backoff
  (1s → 2x → 60s max). _serve_one_session opens WS, registers (with
  pub.api_key from prior session if set so key pinning re-attaches the
  same name), serves until disconnect. AuthError = terminal (no retry on
  register_failed); any other exception = reconnect with backoff.
- connect().runner: same shape — _serve_one_session + outer loop.
- on registered envelope, only update pub.api_key if the new value is
  non-empty (so a transient empty re-register response can't clear it).

tests:
- tests/test_reconnect.py: spin up uvicorn hub thread, register publisher,
  set should_exit on the server (proper shutdown), restart hub on same
  port + same db_path, send a chat through a new client, assert response
  arrives — only succeeds if publisher auto-re-registered.

result: 33/33 pytest passing. publisher now survives:
  - hub restart (re-registers via key pinning)
  - transient network blip (reconnects with backoff)
  - clean reconnect if hub was down momentarily

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: Zawwarsami16

tests/test_reconnect.py

@@ -0,0 +1,136 @@
+"""Auto-reconnect after hub WS drops or hub restarts.
+
+Today (pre-1.5) the publish/connect runners die permanently on disconnect.
+This test asserts they re-register automatically when the hub comes back.
+"""
+
+import asyncio
+import socket
+import threading
+import time
+
+import pytest
+
+try:
+    import fastapi  # noqa
+    import uvicorn  # noqa
+    SERVER_AVAILABLE = True
+except ImportError:
+    SERVER_AVAILABLE = False
+
+if SERVER_AVAILABLE:
+    from zhub.server import create_app
+from zhub import publish
+
+
+def _free_port() -> int:
+    with socket.socket() as s:
+        s.bind(("", 0))
+        return s.getsockname()[1]
+
+
+class _HubThread:
+    """Owns a uvicorn server in a daemon thread. Allows clean stop()."""
+
+    def __init__(self, port: int, db_path: str):
+        self.port = port
+        self.db_path = db_path
+        self.server: uvicorn.Server | None = None
+        self.thread: threading.Thread | None = None
+
+    def start(self):
+        config = uvicorn.Config(
+            create_app(db_path=self.db_path), host="127.0.0.1", port=self.port,
+            log_level="warning",
+        )
+        self.server = uvicorn.Server(config)
+
+        def _run():
+            asyncio.run(self.server.serve())
+
+        self.thread = threading.Thread(target=_run, daemon=True)
+        self.thread.start()
+        # wait for port
+        for _ in range(50):
+            try:
+                with socket.create_connection(("127.0.0.1", self.port), timeout=0.1):
+                    return
+            except OSError:
+                time.sleep(0.1)
+
+    def stop(self):
+        if self.server:
+            self.server.should_exit = True
+        if self.thread:
+            self.thread.join(timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_publisher_auto_reregisters_after_hub_restart():
+    """Publisher's WS dies when hub stops. After hub comes back on the same
+    port + db, publisher reconnects + re-registers via key pinning."""
+    if not SERVER_AVAILABLE:
+        pytest.skip("fastapi/uvicorn not installed")
+
+    import tempfile, os
+    db_fd, db_path = tempfile.mkstemp(suffix=".db")
+    os.close(db_fd)
+    try:
+        port = _free_port()
+        hub = _HubThread(port, db_path)
+        hub.start()
+
+        # First registration — fresh hub, no prior key.
+        pub = publish(
+            name="reconn",
+            description="reconnect test",
+            chat_handler=lambda m, o: "alive",
+            hub_url=f"ws://127.0.0.1:{port}",
+        )
+        for _ in range(50):
+            if pub.api_key:
+                break
+            await asyncio.sleep(0.05)
+        first_key = pub.api_key
+        assert first_key
+
+        # Stop the hub — publisher's WS will drop.
+        hub.stop()
+        await asyncio.sleep(0.3)
+
+        # Restart the hub on the same port + same db.
+        hub2 = _HubThread(port, db_path)
+        hub2.start()
+
+        # Wait long enough for the reconnect backoff to fire (initial 1s + jitter).
+        for _ in range(50):
+            if pub.api_key == first_key:
+                # api_key still set — but now actually serving via hub2?
+                # try a chat to confirm reconnection is alive
+                pass
+            await asyncio.sleep(0.2)
+
+        # Verify by sending a chat through the new hub — only succeeds if
+        # publisher has reconnected and re-registered.
+        from zhub import connect
+        client = connect(
+            ai_name="reconn", api_key=first_key,
+            hub_url=f"ws://127.0.0.1:{port}",
+            capabilities={},
+        )
+        await asyncio.sleep(0.5)
+
+        try:
+            resp = await asyncio.wait_for(
+                client.chat(messages=[{"role": "user", "content": "ping"}]),
+                timeout=10.0,
+            )
+            assert resp.get("text") == "alive", \
+                f"reconnected publisher didn't serve chat: {resp!r}"
+        finally:
+            hub2.stop()
+    finally:
+        try:
+            os.unlink(db_path)
+        except OSError:
+            pass

zhub/client.py

@@ -149,18 +149,22 @@ def publish(
     pub._pending: dict[str, asyncio.Future] = {}  # type: ignore[attr-defined]
     pub._ws = None  # type: ignore[attr-defined]
 
-    async def runner() -> None:
-        url = _to_ws_url(hub_url, "/ws/publish")
-        log.info("publisher connecting to %s", url)
+    async def _serve_one_session(url: str) -> None:
+        """Open one WebSocket session, register, serve until disconnect.
+        On disconnect, returns normally (the outer loop reconnects)."""
         async with websockets.connect(url, max_size=10_000_000) as ws:
             pub._ws = ws  # type: ignore[attr-defined]
             manifest_dict = manifest.to_dict()
             if private_key:
                 from .signing import sign_manifest as _sign
                 manifest_dict = _sign(manifest_dict, private_key)
             register_env = register_publisher(manifest_dict, name)
-            if api_key:
-                register_env.payload["api_key"] = api_key
+            # Use the api_key from the prior registration if we already have
+            # one (re-register on the same name via key pinning); fall back
+            # to the operator-supplied api_key on first connect.
+            register_key = pub.api_key or api_key
+            if register_key:
+                register_env.payload["api_key"] = register_key
             await ws.send(register_env.to_json())
 
             async for raw in ws:
@@ -169,7 +173,11 @@ async def runner() -> None:
                 if env.type == "registered":
                     pub.name = env.payload.get("name", name)
                     pub.base_url = env.payload.get("base_url", "")
-                    pub.api_key = env.payload.get("api_key", "")
+                    # Hub returns the same key on re-registration; only set
+                    # if non-empty so a transient empty value doesn't clear it.
+                    new_key = env.payload.get("api_key", "")
+                    if new_key:
+                        pub.api_key = new_key
                     log.info("publisher registered as %s with key %s",
                              pub.name, pub.api_key[:10] + "…")
 
@@ -196,6 +204,32 @@ async def runner() -> None:
 
                 elif env.type == "error":
                     log.warning("hub error: %s", env.payload)
+                    # Auth/registration failures are terminal — don't loop.
+                    if env.payload.get("code") == "register_failed":
+                        raise AuthError(env.payload.get("message", "register failed"))
+
+    async def runner() -> None:
+        url = _to_ws_url(hub_url, "/ws/publish")
+        backoff = 1.0
+        while not pub._stop_event.is_set():
+            try:
+                log.info("publisher connecting to %s", url)
+                await _serve_one_session(url)
+                # Clean disconnect — reset backoff before retrying
+                backoff = 1.0
+            except AuthError:
+                # Terminal — stop the loop
+                log.error("publisher registration failed — stopping reconnect loop")
+                return
+            except Exception as e:
+                log.warning("publisher session ended: %s — reconnecting in %.1fs", e, backoff)
+            if pub._stop_event.is_set():
+                return
+            try:
+                await asyncio.wait_for(pub._stop_event.wait(), timeout=backoff)
+            except asyncio.TimeoutError:
+                pass
+            backoff = min(backoff * 2.0, 60.0)
 
     pub._task = asyncio.ensure_future(runner())
     return pub
@@ -278,6 +312,7 @@ class ZhubConnection:
     _ws: Any = None
     _pending: dict[str, asyncio.Future] = field(default_factory=dict)
     _streams: dict[str, asyncio.Queue] = field(default_factory=dict)
+    _stop_event: asyncio.Event = field(default_factory=asyncio.Event)
 
     async def chat(self, messages: list[dict[str, Any]],
                    model: str = "default", temperature: float = 0.4,
@@ -349,9 +384,7 @@ def connect(
         capabilities=handlers,
     )
 
-    async def runner() -> None:
-        url = _to_ws_url(hub_url, "/ws/connect")
-        log.info("client connecting to %s for AI %s", url, ai_name)
+    async def _serve_one_session(url: str) -> None:
         async with websockets.connect(url, max_size=10_000_000) as ws:
             conn._ws = ws
             await ws.send(register_connection(ai_name, api_key, cm.to_dict()).to_json())
@@ -384,6 +417,27 @@ async def runner() -> None:
                     if env.payload.get("code") == "register_failed":
                         raise AuthError(env.payload.get("message", "register failed"))
 
+    async def runner() -> None:
+        url = _to_ws_url(hub_url, "/ws/connect")
+        backoff = 1.0
+        while not conn._stop_event.is_set():
+            try:
+                log.info("client connecting to %s for AI %s", url, ai_name)
+                await _serve_one_session(url)
+                backoff = 1.0
+            except AuthError:
+                log.error("client registration failed — stopping reconnect loop")
+                return
+            except Exception as e:
+                log.warning("client session ended: %s — reconnecting in %.1fs", e, backoff)
+            if conn._stop_event.is_set():
+                return
+            try:
+                await asyncio.wait_for(conn._stop_event.wait(), timeout=backoff)
+            except asyncio.TimeoutError:
+                pass
+            backoff = min(backoff * 2.0, 60.0)
+
     conn._task = asyncio.ensure_future(runner())
     return conn