diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 122412f..377264a 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -423,6 +423,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): _purge_expired_nodes_locked(server) def do_POST(self): + if self.path == "/v1/chat/completions": + self._handle_proxy_chat() + return if self.path == "/v1/nodes/register": self._handle_register() return @@ -463,6 +466,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self._handle_tracker_nodes(model) elif parsed.path == "/v1/raft/status": self._handle_raft_status() + elif parsed.path == "/v1/health": + self._send_json(200, {"status": "ok"}) else: self.send_response(404) self.end_headers() @@ -550,6 +555,119 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): ], }) + # ---------------------------------------------------------------- OpenAI proxy + + def _handle_proxy_chat(self) -> None: + """Proxy POST /v1/chat/completions to a tracker-mode (first-shard) node. + + Picks a live tracker-mode node for the requested model using round-robin, + then forwards the request verbatim and relays the response (including + streaming SSE chunks) back to the caller. + """ + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + + length = int(self.headers.get("Content-Length", 0)) + raw_body = self.rfile.read(length) if length else b"{}" + + try: + body = json.loads(raw_body) + except json.JSONDecodeError: + self._send_json(400, {"error": {"message": "invalid JSON", "type": "invalid_request_error", "code": "invalid_request"}}) + return + + model: str = body.get("model", "") + is_stream: bool = bool(body.get("stream", False)) + + # Find a live tracker-mode node for this model + with server.lock: + self._purge_expired_nodes() + candidates = [ + n for n in server.registry.values() + if n.tracker_mode and (n.model == model or n.hf_repo == model) + ] + + if not candidates: + # Fall back: any node serving shard_start=0 for this model + with server.lock: + candidates = [ + n for n in server.registry.values() + if n.shard_start == 0 and (n.model == model or n.hf_repo == model) + ] + + if not candidates: + self._send_json(503, {"error": { + "message": f"no nodes available for model {model!r}", + "type": "service_unavailable", + "code": "model_not_available", + }}) + return + + # Simple round-robin via list length modulo (stateless, good enough) + node = candidates[int(time.time() * 1000) % len(candidates)] + target_url = f"{node.endpoint}/v1/chat/completions" + + req = urllib.request.Request( + target_url, + data=raw_body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + # Copy Authorization header from client if present + auth = self.headers.get("Authorization") + if auth: + req.add_header("Authorization", auth) + + try: + upstream = urllib.request.urlopen(req, timeout=300.0) + except urllib.error.HTTPError as exc: + # Relay error status + body from node + err_body = exc.read() + self.send_response(exc.code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(err_body))) + self.end_headers() + try: + self.wfile.write(err_body) + except BrokenPipeError: + pass + return + except Exception as exc: + self._send_json(503, {"error": { + "message": f"upstream node unreachable: {exc}", + "type": "service_unavailable", + "code": "upstream_error", + }}) + return + + with upstream: + content_type = upstream.headers.get("Content-Type", "application/json") + if is_stream or "text/event-stream" in content_type: + # Relay SSE stream chunk-by-chunk + self.send_response(200) + self.send_header("Content-Type", "text/event-stream; charset=utf-8") + self.send_header("Cache-Control", "no-cache") + self.end_headers() + try: + while True: + line = upstream.readline() + if not line: + break + self.wfile.write(line) + self.wfile.flush() + except BrokenPipeError: + pass + else: + # Non-streaming: buffer and relay + resp_body = upstream.read() + self.send_response(200) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(resp_body))) + self.end_headers() + try: + self.wfile.write(resp_body) + except BrokenPipeError: + pass + def _handle_register(self): server: _TrackerHTTPServer = self.server # type: ignore[assignment] body = self._read_json_body()