fix: node auto-re-registers with tracker after 404 heartbeat (tracker restart)

When the tracker restarts, nodes' registrations are lost. The heartbeat loop
was catching the 404 and printing a warning but never re-registering, leaving
the node permanently invisible until manually restarted.

Fix: on HTTP 404 heartbeat response, the loop re-posts the original
registration payload to /v1/nodes/register and updates the node_id and
heartbeat URL for subsequent beats. This also handles tracker expiry races.

The register_payload is now passed into _start_heartbeat so the thread has
everything needed for re-registration without reaching back into run_startup.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

packages/node/meshnet_node/startup.py

@@ -34,15 +34,32 @@ def _get_json(url: str, timeout: float = 10.0) -> dict:
         return json.loads(r.read())
 
 
-def _start_heartbeat(tracker_url: str, node_id: str, interval: float = 20.0) -> threading.Thread:
-    """Daemon thread that sends periodic heartbeats to the tracker."""
+def _start_heartbeat(
+    tracker_url: str,
+    node_id: str,
+    register_payload: dict,
+    interval: float = 20.0,
+) -> threading.Thread:
+    """Daemon thread: sends heartbeats; re-registers automatically on 404 (tracker restart)."""
     def _loop() -> None:
+        nonlocal node_id
         hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
         while True:
             time.sleep(interval)
             try:
                 _post_json(hb_url, {})
-                print(f"  [node] heartbeat sent → tracker (node {node_id[:8]})", flush=True)
+            except urllib.error.HTTPError as exc:
+                if exc.code == 404:
+                    print("  [node] tracker lost registration — re-registering...", flush=True)
+                    try:
+                        resp = _post_json(f"{tracker_url}/v1/nodes/register", register_payload)
+                        node_id = resp.get("node_id", node_id)
+                        hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
+                        print(f"  [node] re-registered — node ID: {node_id}", flush=True)
+                    except Exception as re_exc:
+                        print(f"  [node] WARNING: re-registration failed: {re_exc}", flush=True)
+                else:
+                    print(f"  [node] WARNING: heartbeat failed: {exc}", flush=True)
             except Exception as exc:
                 print(f"  [node] WARNING: heartbeat failed: {exc}", flush=True)
 
@@ -154,26 +171,24 @@ def run_startup(
         endpoint = f"http://{public_host}:{actual_port}"
         # Register with tracker so other nodes can auto-join this model.
         total_layers = getattr(node.backend, "total_layers", None)
+        reg_payload = {
+            "endpoint": endpoint,
+            "model": model_id.split("/")[-1],
+            "hf_repo": model_id,
+            "num_layers": total_layers,
+            "shard_start": shard_start,
+            "shard_end": shard_end,
+            "hardware_profile": hw,
+            "wallet_address": address,
+            "quantization": quantization,
+            "score": 1.0,
+            "tracker_mode": (shard_start == 0),
+        }
         try:
-            reg_resp = _post_json(
-                f"{tracker_url}/v1/nodes/register",
-                {
-                    "endpoint": endpoint,
-                    "model": model_id.split("/")[-1],
-                    "hf_repo": model_id,
-                    "num_layers": total_layers,
-                    "shard_start": shard_start,
-                    "shard_end": shard_end,
-                    "hardware_profile": hw,
-                    "wallet_address": address,
-                    "quantization": quantization,
-                    "score": 1.0,
-                    "tracker_mode": (shard_start == 0),
-                },
-            )
+            reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", reg_payload)
             node_id = reg_resp.get("node_id", "?")
             print(f"  Registered with tracker — node ID: {node_id}", flush=True)
-            _start_heartbeat(tracker_url, node_id)
+            _start_heartbeat(tracker_url, node_id, reg_payload)
         except Exception as exc:
             print(f"  Warning: tracker registration failed: {exc}", flush=True)
 
@@ -227,26 +242,24 @@ def run_startup(
         actual_port = node.start()
         public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
         endpoint = f"http://{public_host}:{actual_port}"
+        auto_reg_payload = {
+            "endpoint": endpoint,
+            "model": assigned_hf_repo.split("/")[-1],
+            "hf_repo": assigned_hf_repo,
+            "num_layers": assigned_num_layers,
+            "shard_start": assigned_shard_start,
+            "shard_end": assigned_shard_end,
+            "hardware_profile": hw,
+            "wallet_address": address,
+            "quantization": quantization,
+            "score": 1.0,
+            "tracker_mode": (assigned_shard_start == 0),
+        }
         try:
-            reg_resp = _post_json(
-                f"{tracker_url}/v1/nodes/register",
-                {
-                    "endpoint": endpoint,
-                    "model": assigned_hf_repo.split("/")[-1],
-                    "hf_repo": assigned_hf_repo,
-                    "num_layers": assigned_num_layers,
-                    "shard_start": assigned_shard_start,
-                    "shard_end": assigned_shard_end,
-                    "hardware_profile": hw,
-                    "wallet_address": address,
-                    "quantization": quantization,
-                    "score": 1.0,
-                    "tracker_mode": (assigned_shard_start == 0),
-                },
-            )
+            reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", auto_reg_payload)
             node_id = reg_resp.get("node_id", "?")
             print(f"  Registered with tracker — node ID: {node_id}", flush=True)
-            _start_heartbeat(tracker_url, node_id)
+            _start_heartbeat(tracker_url, node_id, auto_reg_payload)
         except Exception as exc:
             print(f"  Warning: tracker registration failed: {exc}", flush=True)
         shard_count = assigned_shard_end - assigned_shard_start + 1

packages/node/meshnet_node/torch_server.py

@@ -238,33 +238,75 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
                 self._send_json(500, {"error": f"generation failed: {exc}"})
             return
 
-        # Distributed path: encode prompt at the head, forward activations along the route.
-        prompt = " ".join(
-            str(m.get("content", ""))
-            for m in messages
-            if isinstance(m, dict) and m.get("role") == "user"
-        )
-        try:
-            payload = server.backend.encode_prompt(prompt)
-        except Exception as exc:
-            self._send_json(500, {"error": f"encode_prompt failed: {exc}"})
-            return
+        # Distributed path: autoregressive generation across shards.
+        # We do N single-step forward passes (no cross-node KV cache), which is slow
+        # but correct. Each step: head encodes current sequence → forwards through route
+        # → tail returns the next token string → append → repeat.
         remaining_route = self._get_remaining_route(model_name)
-        result_text = self._run_downstream_pipeline(payload, remaining_route)
+        if not remaining_route:
+            self._send_openai_response(
+                "error: no downstream route — check tracker connectivity",
+                model_name, False, messages,
+            )
+            return
+
+        backend = server.backend
+        # Format with chat template so the model knows it's in assistant mode.
+        try:
+            if hasattr(backend.tokenizer, "apply_chat_template"):
+                prompt_text: str = backend.tokenizer.apply_chat_template(
+                    messages, add_generation_prompt=True, tokenize=False,
+                )
+            else:
+                raise AttributeError("no apply_chat_template")
+        except Exception:
+            prompt_text = " ".join(
+                str(m.get("content", ""))
+                for m in messages
+                if isinstance(m, dict) and m.get("role") == "user"
+            )
+
+        eos_token: str = getattr(backend.tokenizer, "eos_token", "") or ""
+        generated: list[str] = []
+        current_text = prompt_text
+
+        for _ in range(max_tokens):
+            try:
+                payload = backend.encode_prompt(current_text)
+            except Exception as exc:
+                print(f"  [node] distributed encode error: {exc}", flush=True)
+                break
+            token_str = self._run_downstream_pipeline(payload, remaining_route)
+            if not token_str:
+                break
+            # Stop on error responses or EOS.
+            if token_str.startswith(("pipeline error", "decode error", "no downstream", "error:")):
+                break
+            if eos_token and token_str == eos_token:
+                break
+            generated.append(token_str)
+            current_text = current_text + token_str
+
+        result_text = "".join(generated)
         self._send_openai_response(result_text, model_name, stream, messages)
 
     def _get_remaining_route(self, model: str) -> list[str]:
         server: _TorchHTTPServer = self.server  # type: ignore[assignment]
         if server.tracker_url is None:
             return []
+        # Use the backend's actual hf_repo, not the client-provided model name (which may be
+        # a lowercased or abbreviated alias that doesn't match what the tracker registered).
+        route_model = getattr(server.backend, "model_id", None) or model
         try:
-            url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(model)}"
+            url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}"
             with urllib.request.urlopen(url, timeout=5.0) as r:
                 route_resp = json.loads(r.read())
             route = route_resp.get("route", [])
-            # Skip the first node in the route (self) since we're already the head
-            return list(route[1:])
-        except Exception:
+            # Skip our own endpoint from the route (match by port so host aliases don't matter).
+            own_port = server.server_address[1]
+            return [ep for ep in route if not ep.rstrip("/").endswith(f":{own_port}")]
+        except Exception as exc:
+            print(f"  [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
             return []
 
     def _run_downstream_pipeline(self, payload: object, route: list[str]) -> str:

packages/tracker/meshnet_tracker/server.py

@@ -697,10 +697,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             _rebalance_model_locked(server, entry.model or "stub-model")
             directives = list(entry.pending_directives)
             entry.pending_directives.clear()
-        print(
-            f"[tracker] heartbeat: {node_id[:8]}  {entry.endpoint}",
-            flush=True,
-        )
+        # print(
+        #     f"[tracker] heartbeat: {node_id[:8]}  {entry.endpoint}",
+        #     flush=True,
+        # )
         if directives:
             self._send_json(200, {"directives": directives})
         else: