feat: add tracker registration routing

packages/gateway/meshnet_gateway/server.py

@@ -4,13 +4,16 @@ import http.server
 import json
 import threading
 import time
+import urllib.error
+import urllib.parse
 import urllib.request
 
 
 class _GatewayHTTPServer(http.server.HTTPServer):
-    def __init__(self, addr, handler, inference_route: list[str]):
+    def __init__(self, addr, handler, inference_route: list[str] | None, tracker_url: str | None):
         super().__init__(addr, handler)
         self.inference_route = inference_route
+        self.tracker_url = tracker_url
 
 
 class _GatewayHandler(http.server.BaseHTTPRequestHandler):
@@ -24,11 +27,41 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler):
             self.send_response(404)
             self.end_headers()
 
+    def _send_json_error(self, status: int, message: str) -> None:
+        resp = json.dumps({"error": message}).encode()
+        self.send_response(status)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(resp)))
+        self.end_headers()
+        self.wfile.write(resp)
+
     def _handle_chat_completions(self):
+        server: _GatewayHTTPServer = self.server  # type: ignore[assignment]
         length = int(self.headers.get("Content-Length", 0))
         body = json.loads(self.rfile.read(length))
 
-        route: list[str] = self.server.inference_route  # type: ignore[attr-defined]
+        if server.tracker_url is not None:
+            model = body.get("model", "stub-model")
+            tracker_path = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(model)}"
+            try:
+                tracker_resp = _get_json(tracker_path)
+                route = tracker_resp["route"]
+                if (
+                    not isinstance(route, list)
+                    or not route
+                    or not all(isinstance(node_url, str) and node_url for node_url in route)
+                ):
+                    raise ValueError("invalid tracker route")
+            except urllib.error.HTTPError as exc:
+                error_body = _safe_error_body(exc)
+                self._send_json_error(503, error_body.get("error", "tracker unavailable"))
+                return
+            except (urllib.error.URLError, TimeoutError, json.JSONDecodeError, KeyError, TypeError, ValueError):
+                self._send_json_error(503, "tracker unavailable")
+                return
+        else:
+            route = server.inference_route  # type: ignore[assignment]
+
         messages = body.get("messages", [])
 
         # Send messages to the first node in the inference route.
@@ -65,6 +98,14 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler):
         self.wfile.write(response_body)
 
 
+def _safe_error_body(exc: urllib.error.HTTPError) -> dict:
+    try:
+        body = json.loads(exc.read())
+    except json.JSONDecodeError:
+        return {}
+    return body if isinstance(body, dict) else {}
+
+
 def _post_json(url: str, payload: dict, timeout: float = 5.0) -> dict:
     data = json.dumps(payload).encode()
     req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}, method="POST")
@@ -72,11 +113,18 @@ def _post_json(url: str, payload: dict, timeout: float = 5.0) -> dict:
         return json.loads(r.read())
 
 
+def _get_json(url: str, timeout: float = 5.0) -> dict:
+    with urllib.request.urlopen(url, timeout=timeout) as r:
+        return json.loads(r.read())
+
+
 class GatewayServer:
     """HTTP gateway that routes /v1/chat/completions through an ordered inference route.
 
-    Supply either ``node_url`` (single-node, US-001 compat) or ``inference_route``
-    (ordered list of node base URLs for multi-shard routing).
+    Routing modes (mutually exclusive):
+    - ``tracker_url``: gateway queries the tracker per request using the model name.
+    - ``inference_route``: static ordered list of node base URLs.
+    - ``node_url``: single-node shortcut (US-001 compat), equivalent to inference_route=[node_url].
     """
 
     def __init__(
@@ -84,18 +132,23 @@ class GatewayServer:
         node_url: str | None = None,
         *,
         inference_route: list[str] | None = None,
+        tracker_url: str | None = None,
         host: str = "127.0.0.1",
         port: int = 0,
     ):
-        if node_url is not None and inference_route is not None:
-            raise ValueError("Provide node_url or inference_route, not both")
+        routing_modes = sum(x is not None for x in (node_url, inference_route, tracker_url))
+        if routing_modes > 1:
+            raise ValueError("Provide exactly one of: node_url, inference_route, or tracker_url")
+        if routing_modes == 0:
+            raise ValueError("Provide one of: node_url, inference_route, or tracker_url")
+
         if node_url is not None:
             inference_route = [node_url]
-        elif inference_route is None:
-            raise ValueError("Provide either node_url or inference_route")
-        if not inference_route:
+        if inference_route is not None and not inference_route:
             raise ValueError("inference_route must contain at least one node URL")
-        self._inference_route = inference_route
+
+        self._inference_route: list[str] | None = inference_route
+        self._tracker_url: str | None = tracker_url.rstrip("/") if tracker_url is not None else None
         self._host = host
         self._requested_port = port
         self._server: _GatewayHTTPServer | None = None
@@ -110,6 +163,7 @@ class GatewayServer:
             (self._host, self._requested_port),
             _GatewayHandler,
             self._inference_route,
+            self._tracker_url,
         )
         self.port = self._server.server_address[1]
         self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)