neuron-tai/packages/gateway/meshnet_gateway/server.py

"""Gateway HTTP server — accepts OpenAI-format requests and routes to inference nodes."""

import http.server
import hashlib
import json
import os
from collections import Counter
from dataclasses import dataclass
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
from typing import Any

_STUB_HIDDEN_DIM = 64
_STUB_DTYPE = "bfloat16"
_WIRE_VERSION = "2"
_DTYPE_SIZES = {
    "float16": 2,
    "bfloat16": 2,
    "float32": 4,
}


@dataclass
class _BinaryActivation:
    body: bytes
    shape: list[int]
    dtype: str
    session: str
    chunk_index: int
    chunk_total: int
    encoding: str | None
    headers: dict[str, str]


class _GatewayHTTPServer(http.server.HTTPServer):
    def __init__(
        self,
        addr,
        handler,
        inference_route: list[str] | None,
        tracker_url: str | None,
        model_presets: list[str] | None,
        contracts: Any | None,
        minimum_stake: int,
        cost_per_layer_token_lamport: int,
    ):
        super().__init__(addr, handler)
        self.inference_route = inference_route
        self.tracker_url = tracker_url
        self.model_presets = model_presets
        self.contracts = contracts
        self.minimum_stake = minimum_stake
        self.cost_per_layer_token_lamport = cost_per_layer_token_lamport
        self.last_binary_chunk_responses: list[_BinaryActivation] = []
        self.request_count: int = 0


class _GatewayHandler(http.server.BaseHTTPRequestHandler):
    def log_message(self, fmt, *args):  # noqa: suppress request logs in tests
        pass

    def do_GET(self):
        parsed = urllib.parse.urlparse(self.path)
        if parsed.path == "/v1/health":
            self._handle_health()
        elif parsed.path == "/v1/models":
            self._handle_models()
        elif parsed.path == "/v1/models/available":
            self._handle_available_models()
        elif parsed.path == "/v1/wallet/balance":
            self._handle_wallet_balance()
        elif parsed.path == "/v1/estimate_cost":
            self._handle_estimate_cost(parsed)
        else:
            self.send_response(404)
            self.end_headers()

    def do_POST(self):
        parsed = urllib.parse.urlparse(self.path)
        if parsed.path == "/v1/chat/completions":
            self._handle_chat_completions()
        elif parsed.path == "/v1/wallet/top_up":
            self._handle_wallet_top_up()
        elif parsed.path == "/v1/request":
            self._handle_meshnet_request()
        else:
            self.send_response(404)
            self.end_headers()

    def _send_json(self, status: int, data: dict) -> None:
        body = json.dumps(data).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _send_json_error(self, status: int, message: str) -> None:
        self._send_json(status, {"error": message})

    def _send_openai_error(self, status: int, message: str, code: str) -> None:
        self._send_json(status, {
            "error": {
                "message": message,
                "type": "invalid_request_error",
                "param": "model",
                "code": code,
            }
        })

    def _read_json_body(self) -> dict | None:
        length = int(self.headers.get("Content-Length", 0))
        try:
            body = json.loads(self.rfile.read(length) or b"{}")
        except (json.JSONDecodeError, ValueError):
            self._send_json_error(400, "invalid JSON body")
            return None
        if not isinstance(body, dict):
            self._send_json_error(400, "JSON body must be an object")
            return None
        return body

    def _handle_health(self) -> None:
        self._send_json(200, {"status": "ok"})

    def _handle_models(self) -> None:
        server: _GatewayHTTPServer = self.server  # type: ignore[assignment]
        created = int(time.time())
        if server.tracker_url is not None:
            try:
                models_resp = _get_json(f"{server.tracker_url}/v1/models")
                self._send_json(200, models_resp)
                return
            except Exception:
                self._send_json_error(503, "tracker unavailable")
                return
        presets = server.model_presets or ["stub-model"]
        data = [
            {"id": p, "object": "model", "created": created, "owned_by": "meshnet"}
            for p in presets
        ]
        self._send_json(200, {"object": "list", "data": data})

    def _handle_available_models(self) -> None:
        server: _GatewayHTTPServer = self.server  # type: ignore[assignment]
        if server.tracker_url is not None:
            try:
                models_resp = _get_json(f"{server.tracker_url}/v1/models")
            except Exception:
                self._send_json_error(503, "tracker unavailable")
                return
            data = [
                {
                    "id": model["id"],
                    "shard_coverage_percentage": float(model.get("shard_coverage_percentage", 0.0)),
                }
                for model in models_resp.get("data", [])
                if isinstance(model, dict) and isinstance(model.get("id"), str)
            ]
            self._send_json(200, {"data": data})
            return
        presets = server.model_presets or ["stub-model"]
        self._send_json(200, {
            "data": [
                {"id": model, "shard_coverage_percentage": 100.0}
                for model in presets
            ]
        })

    def _handle_wallet_balance(self) -> None:
        server: _GatewayHTTPServer = self.server  # type: ignore[assignment]
        api_key = _api_key_from_authorization(self.headers.get("Authorization"))
        if not api_key:
            self._send_json_error(401, "missing API key")
            return
        if server.contracts is None:
            self._send_json(200, {"sol": 0.0, "usdc": 0.0})
            return
        balance = server.contracts.payment.get_balance(api_key)
        self._send_json(200, {
            "sol": _lamports_to_sol(balance.lamports),
            "usdc": balance.usdc_micro / 1_000_000,
        })

    def _handle_wallet_top_up(self) -> None:
        api_key = _api_key_from_authorization(self.headers.get("Authorization"))
        if not api_key:
            self._send_json_error(401, "missing API key")
            return
        body = self._read_json_body()
        if body is None:
            return
        try:
            amount_sol = float(body.get("amount_sol", 0.0))
            amount_usdc = float(body.get("amount_usdc", 0.0))
        except (TypeError, ValueError):
            self._send_json_error(400, "top-up amount must be numeric")
            return
        if amount_sol <= 0 and amount_usdc <= 0:
            self._send_json_error(400, "top-up amount must be positive")
            return
        payment_address = _payment_address_for_api_key(api_key)
        query = urllib.parse.urlencode({
            "amount": amount_sol,
            "spl-token": "USDC" if amount_usdc > 0 else "SOL",
        })
        self._send_json(200, {
            "payment_address": payment_address,
            "qr": f"solana:{payment_address}?{query}",
            "amount_sol": amount_sol,
            "amount_usdc": amount_usdc,
        })

    def _handle_estimate_cost(self, parsed: urllib.parse.ParseResult) -> None:
        server: _GatewayHTTPServer = self.server  # type: ignore[assignment]
        params = urllib.parse.parse_qs(parsed.query)
        model = params.get("model", [""])[0]
        try:
            tokens = int(params.get("tokens", ["0"])[0])
        except ValueError:
            self._send_json_error(400, "tokens must be an integer")
            return
        if not model:
            self._send_json_error(400, "model is required")
            return
        if tokens <= 0:
            self._send_json_error(400, "tokens must be positive")
            return
        layer_count = 32
        if server.tracker_url is None and server.model_presets is not None and model not in server.model_presets:
            self._send_json_error(404, "model not available")
            return
        lamports = layer_count * tokens * server.cost_per_layer_token_lamport
        self._send_json(200, {
            "model": model,
            "tokens": tokens,
            "lamports": lamports,
            "sol": _lamports_to_sol(lamports),
        })

    def _handle_chat_completions(self):
        server: _GatewayHTTPServer = self.server  # type: ignore[assignment]
        # Read raw bytes first so we can proxy them if tracker-nodes are available
        length = int(self.headers.get("Content-Length", 0))
        raw_body = self.rfile.read(length)
        try:
            body = json.loads(raw_body or b"{}")
        except (json.JSONDecodeError, ValueError):
            self._send_json_error(400, "invalid JSON body")
            return
        if not isinstance(body, dict):
            self._send_json_error(400, "JSON body must be an object")
            return

        model = str(body.get("model", "stub-model"))
        tracker_nodes = _get_tracker_nodes(server, model)
        if tracker_nodes:
            # Proxy to a tracker-node (round-robin by request count)
            target = tracker_nodes[server.request_count % len(tracker_nodes)]
            server.request_count += 1
            return self._proxy_to_tracker_node(target, raw_body)

        # Fallback: use existing direct pipeline (backward compat)
        streaming = bool(body.get("stream", False))
        try:
            completion = self._build_completion(body)
        except _ModelUnavailable as exc:
            self._send_openai_error(503, exc.message, "model_not_available")
            return
        except _TrackerUnavailable:
            self._send_json_error(503, "tracker unavailable")
            return
        except (urllib.error.URLError, TimeoutError, KeyError, json.JSONDecodeError, ValueError) as exc:
            self._send_openai_error(503, f"inference pipeline error: {exc}", "pipeline_error")
            return

        if streaming:
            text = completion["choices"][0]["message"]["content"]
            self._send_sse_response(text, completion["model"])
        else:
            self._send_json(200, completion)

    def _proxy_to_tracker_node(self, url: str, body_bytes: bytes) -> None:
        """Forward a raw request body to a tracker-node and stream the response back."""
        target_url = f"{url}/v1/chat/completions"
        req = urllib.request.Request(
            target_url,
            data=body_bytes,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        try:
            with urllib.request.urlopen(req, timeout=30.0) as r:
                content_type = r.headers.get("Content-Type", "application/json")
                resp_body = r.read()
                status = r.status
        except urllib.error.HTTPError as exc:
            resp_body = exc.read()
            self.send_response(exc.code)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(resp_body)))
            self.end_headers()
            self.wfile.write(resp_body)
            return
        except Exception as exc:
            self._send_json_error(503, f"tracker-node unreachable: {exc}")
            return
        self.send_response(status)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(resp_body)))
        self.end_headers()
        self.wfile.write(resp_body)

    def _handle_meshnet_request(self) -> None:
        body = self._read_json_body()
        if body is None:
            return
        try:
            redundancy = int(body.get("redundancy", 1))
        except (TypeError, ValueError):
            self._send_json_error(400, "redundancy must be an integer")
            return
        if redundancy < 1:
            self._send_json_error(400, "redundancy must be at least 1")
            return

        try:
            routes = self._resolve_routes(str(body.get("model", "stub-model")), redundancy)
            responses = [self._build_completion(body, route_payload=route_payload) for route_payload in routes]
        except _ModelUnavailable as exc:
            self._send_openai_error(503, exc.message, "model_not_available")
            return
        except _TrackerUnavailable:
            self._send_json_error(503, "tracker unavailable")
            return
        except (urllib.error.URLError, TimeoutError, KeyError, json.JSONDecodeError, ValueError) as exc:
            self._send_openai_error(503, f"inference pipeline error: {exc}", "pipeline_error")
            return

        majority_response = _majority_response(responses)
        self._send_json(200, {
            "redundancy": redundancy,
            "majority_response": majority_response,
            "responses": responses,
        })

    def _build_completion(self, body: dict, route_payload: dict | None = None) -> dict:
        server: _GatewayHTTPServer = self.server  # type: ignore[assignment]
        if route_payload is None:
            route_payload = self._resolve_route(str(body.get("model", "stub-model")))
        route = route_payload["route"]
        route_nodes = route_payload.get("nodes", [])
        messages = body.get("messages", [])

        prompt = _last_message_content(messages)
        chunk_responses = _run_binary_pipeline(route, prompt)
        _reassemble_binary_responses(chunk_responses)
        server.last_binary_chunk_responses = chunk_responses
        response_prefix = _stub_response_prefix(chunk_responses)
        text = f"{response_prefix} {prompt}"
        model = body.get("model", "stub-model")
        self._record_inference(model, messages, text, route_nodes)
        return _completion_response(str(model), text)

    def _resolve_route(self, model: str) -> dict:
        server: _GatewayHTTPServer = self.server  # type: ignore[assignment]
        if server.tracker_url is not None:
            tracker_path = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(model)}"
            try:
                tracker_resp = _get_json(tracker_path)
                _validate_route_payload(tracker_resp)
                banned_wallet = _banned_route_wallet(server.contracts, tracker_resp.get("nodes", []))
                if banned_wallet is not None:
                    raise _ModelUnavailable(f"model not available: route includes banned wallet {banned_wallet}")
                return tracker_resp
            except urllib.error.HTTPError as exc:
                error_body = _safe_error_body(exc)
                raise _ModelUnavailable(error_body.get("error", "model not available")) from exc
            except (urllib.error.URLError, TimeoutError, json.JSONDecodeError, KeyError, TypeError, ValueError) as exc:
                raise _TrackerUnavailable() from exc
        return {"route": server.inference_route, "nodes": []}

    def _resolve_routes(self, model: str, redundancy: int) -> list[dict]:
        server: _GatewayHTTPServer = self.server  # type: ignore[assignment]
        if server.tracker_url is None:
            return [self._resolve_route(model) for _ in range(redundancy)]
        tracker_path = (
            f"{server.tracker_url}/v1/routes?"
            f"model={urllib.parse.quote(model)}&redundancy={redundancy}"
        )
        try:
            tracker_resp = _get_json(tracker_path)
            routes = tracker_resp["routes"]
            if not isinstance(routes, list) or len(routes) != redundancy:
                raise ValueError("invalid tracker routes")
            for route_payload in routes:
                _validate_route_payload(route_payload)
                banned_wallet = _banned_route_wallet(server.contracts, route_payload.get("nodes", []))
                if banned_wallet is not None:
                    raise _ModelUnavailable(f"model not available: route includes banned wallet {banned_wallet}")
            return routes
        except urllib.error.HTTPError as exc:
            error_body = _safe_error_body(exc)
            raise _ModelUnavailable(error_body.get("error", "model not available")) from exc
        except (urllib.error.URLError, TimeoutError, json.JSONDecodeError, KeyError, TypeError, ValueError) as exc:
            raise _TrackerUnavailable() from exc

    def _record_inference(
        self,
        model: object,
        messages: object,
        text: str,
        route_nodes: object,
    ) -> None:
        server: _GatewayHTTPServer = self.server  # type: ignore[assignment]
        session_id = str(uuid.uuid4())
        if server.contracts is not None and isinstance(route_nodes, list):
            api_key = _api_key_from_authorization(self.headers.get("Authorization"))
            if api_key:
                tokens = max(1, _estimate_token_count(messages, text))
                for node in route_nodes:
                    wallet_address = node.get("wallet_address") if isinstance(node, dict) else None
                    if not wallet_address:
                        continue
                    server.contracts.payment.record_attribution(
                        session_id=session_id,
                        api_key=api_key,
                        node_wallet=wallet_address,
                        layer_start=node["shard_start"],
                        layer_end=node["shard_end"],
                        tokens=tokens,
                        speed_score=float(node.get("score", 1.0)),
                    )
        if server.contracts is not None:
            server.contracts.validation.record_completed_inference(
                session_id=session_id,
                model=str(model),
                messages=messages if isinstance(messages, list) else [],
                observed_output=text,
                route_nodes=route_nodes if isinstance(route_nodes, list) else [],
            )

    def _send_sse_response(self, text: str, model: str) -> None:
        chunk_id = "chatcmpl-stub"
        created = int(time.time())

        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream; charset=utf-8")
        self.send_header("Cache-Control", "no-cache")
        self.send_header("X-Accel-Buffering", "no")
        self.end_headers()

        def _emit(data: str) -> None:
            self.wfile.write(f"data: {data}\n\n".encode())
            self.wfile.flush()

        # Role delta
        _emit(json.dumps({
            "id": chunk_id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model,
            "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
        }))

        # Content delta
        _emit(json.dumps({
            "id": chunk_id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model,
            "choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}],
        }))

        # Stop delta
        _emit(json.dumps({
            "id": chunk_id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model,
            "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
        }))

        # SSE sentinel
        self.wfile.write(b"data: [DONE]\n\n")
        self.wfile.flush()


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 _valid_route_nodes(route_nodes: object) -> bool:
    if not isinstance(route_nodes, list):
        return False
    for node in route_nodes:
        if not isinstance(node, dict):
            return False
        if not isinstance(node.get("endpoint"), str) or not node["endpoint"]:
            return False
        wallet_address = node.get("wallet_address")
        if wallet_address is not None and not isinstance(wallet_address, str):
            return False
        try:
            shard_start = int(node["shard_start"])
            shard_end = int(node["shard_end"])
        except (KeyError, TypeError, ValueError):
            return False
        if shard_start < 0 or shard_end < shard_start:
            return False
        node["shard_start"] = shard_start
        node["shard_end"] = shard_end
    return True


def _validate_route_payload(route_payload: dict) -> None:
    route = route_payload["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")
    route_nodes = route_payload.get("nodes", [])
    if route_nodes and not _valid_route_nodes(route_nodes):
        raise ValueError("invalid tracker route metadata")


def _banned_route_wallet(contracts: Any | None, route_nodes: object) -> str | None:
    if contracts is None or not isinstance(route_nodes, list):
        return None
    for node in route_nodes:
        if not isinstance(node, dict):
            continue
        wallet_address = node.get("wallet_address")
        if isinstance(wallet_address, str) and contracts.registry.get_wallet(wallet_address).banned:
            return wallet_address
    return None


def _api_key_from_authorization(header: str | None) -> str | None:
    if not header:
        return None
    scheme, _, token = header.partition(" ")
    if scheme.lower() != "bearer" or not token.strip():
        return None
    return token.strip()


def _estimate_token_count(messages: object, completion_text: str) -> int:
    """Request-token placeholder until shard nodes return tokenizer usage."""
    return 1


def _lamports_to_sol(lamports: int) -> float:
    return lamports / 1_000_000_000


def _payment_address_for_api_key(api_key: str) -> str:
    alphabet = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"
    value = int.from_bytes(hashlib.sha256(f"meshnet:{api_key}".encode()).digest(), "big")
    chars = []
    while value:
        value, remainder = divmod(value, 58)
        chars.append(alphabet[remainder])
    return "".join(reversed(chars))[:44].rjust(32, "1")


def _completion_response(model: str, text: str) -> dict:
    return {
        "id": "chatcmpl-stub",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model,
        "choices": [{
            "index": 0,
            "message": {
                "role": "assistant",
                "content": text,
            },
            "finish_reason": "stop",
        }],
        "usage": {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
        },
    }


def _majority_response(responses: list[dict]) -> dict:
    if not responses:
        raise ValueError("responses must not be empty")
    contents = [
        response["choices"][0]["message"]["content"]
        for response in responses
    ]
    majority_content, _ = Counter(contents).most_common(1)[0]
    for response in responses:
        if response["choices"][0]["message"]["content"] == majority_content:
            return response
    return responses[0]


class _ModelUnavailable(Exception):
    def __init__(self, message: str) -> None:
        super().__init__(message)
        self.message = message


class _TrackerUnavailable(Exception):
    pass


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")
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())


def _last_message_content(messages: object) -> str:
    if not isinstance(messages, list) or not messages:
        return ""
    last_message = messages[-1]
    if not isinstance(last_message, dict):
        return ""
    content = last_message.get("content", "")
    return content if isinstance(content, str) else str(content)


def _run_binary_pipeline(route: list[str], prompt: str, timeout: float = 5.0) -> list[_BinaryActivation]:
    session = str(uuid.uuid4())
    chunk_token_count = _chunk_token_count()
    total_tokens = max(1, _prompt_token_count(prompt))
    chunk_total = max(1, (total_tokens + chunk_token_count - 1) // chunk_token_count)
    responses: list[_BinaryActivation] = []

    for chunk_index in range(chunk_total):
        remaining_tokens = total_tokens - (chunk_index * chunk_token_count)
        seq_len = min(chunk_token_count, remaining_tokens)
        activation = _BinaryActivation(
            body=_make_stub_binary_activation([1, seq_len, _STUB_HIDDEN_DIM], _STUB_DTYPE),
            shape=[1, seq_len, _STUB_HIDDEN_DIM],
            dtype=_STUB_DTYPE,
            session=session,
            chunk_index=chunk_index,
            chunk_total=chunk_total,
            encoding=_preferred_binary_encoding(),
            headers={},
        )
        for hop_index, node_url in enumerate(route):
            activation = _post_binary_forward(
                f"{node_url}/forward",
                activation,
                hop_index=hop_index,
                timeout=timeout,
            )
        responses.append(activation)
    return responses


def _chunk_token_count() -> int:
    raw_value = os.environ.get("MESHNET_CHUNK_TOKENS", "128")
    try:
        value = int(raw_value)
    except ValueError:
        return 128
    return value if value > 0 else 128


def _prompt_token_count(prompt: str) -> int:
    return len(prompt.split()) if prompt.split() else max(1, len(prompt))


def _preferred_binary_encoding() -> str | None:
    try:
        import zstandard  # noqa: F401
    except ModuleNotFoundError:
        return None
    return "zstd"


def _make_stub_binary_activation(shape: list[int], dtype: str) -> bytes:
    element_count = 1
    for dim in shape:
        element_count *= dim
    return b"\x00" * (element_count * _DTYPE_SIZES[dtype])


def _post_binary_forward(
    url: str,
    activation: _BinaryActivation,
    *,
    hop_index: int,
    timeout: float,
) -> _BinaryActivation:
    request_body = _compress_body(activation.body, activation.encoding)
    headers = {
        "Content-Type": "application/octet-stream",
        "X-Meshnet-Wire": _WIRE_VERSION,
        "X-Meshnet-Shape": ",".join(str(dim) for dim in activation.shape),
        "X-Meshnet-Dtype": activation.dtype,
        "X-Meshnet-Session": activation.session,
        "X-Meshnet-Chunk-Index": str(activation.chunk_index),
        "X-Meshnet-Chunk-Total": str(activation.chunk_total),
        "X-Meshnet-Hop-Index": str(hop_index),
    }
    if activation.encoding:
        headers["X-Meshnet-Encoding"] = activation.encoding
    req = urllib.request.Request(url, data=request_body, headers=headers, method="POST")
    with urllib.request.urlopen(req, timeout=timeout) as response:
        response_body = response.read()
        response_headers = {key.lower(): value for key, value in response.headers.items()}

    encoding = response_headers.get("x-meshnet-encoding")
    raw_body = _decompress_body(response_body, encoding)
    shape = _parse_shape(response_headers["x-meshnet-shape"])
    dtype = response_headers["x-meshnet-dtype"]
    session = response_headers["x-meshnet-session"]
    chunk_index = int(response_headers["x-meshnet-chunk-index"])
    chunk_total = int(response_headers["x-meshnet-chunk-total"])
    if session != activation.session:
        raise ValueError("binary activation response changed session")
    if chunk_index != activation.chunk_index or chunk_total != activation.chunk_total:
        raise ValueError("binary activation response changed chunk metadata")
    if dtype != activation.dtype or shape != activation.shape:
        raise ValueError("binary activation response shape/dtype mismatch")
    _validate_activation_body(raw_body, shape, dtype)
    return _BinaryActivation(
        body=raw_body,
        shape=shape,
        dtype=dtype,
        session=session,
        chunk_index=chunk_index,
        chunk_total=chunk_total,
        encoding=encoding,
        headers=response_headers,
    )


def _reassemble_binary_responses(chunks: list[_BinaryActivation]) -> bytes:
    """Validate and concatenate final route chunk responses into one activation body."""
    if not chunks:
        raise ValueError("binary pipeline returned no chunks")
    first = chunks[0]
    expected_total = first.chunk_total
    if len(chunks) != expected_total:
        raise ValueError("binary pipeline returned incomplete chunk set")
    if [chunk.chunk_index for chunk in chunks] != list(range(expected_total)):
        raise ValueError("binary pipeline returned out-of-order chunks")
    for chunk in chunks:
        if chunk.session != first.session or chunk.dtype != first.dtype:
            raise ValueError("binary chunks do not belong to the same activation")
        if len(chunk.shape) != 3 or chunk.shape[0] != first.shape[0] or chunk.shape[2] != first.shape[2]:
            raise ValueError("binary chunks have incompatible shapes")
    return b"".join(chunk.body for chunk in chunks)


def _compress_body(body: bytes, encoding: str | None) -> bytes:
    if not encoding:
        return body
    if encoding != "zstd":
        raise ValueError("unsupported binary encoding")
    import zstandard as zstd

    try:
        return zstd.ZstdCompressor(level=1).compress(body)
    except zstd.ZstdError as exc:
        raise ValueError("could not zstd-compress binary activation") from exc


def _decompress_body(body: bytes, encoding: str | None) -> bytes:
    if not encoding:
        return body
    if encoding != "zstd":
        raise ValueError("unsupported binary encoding")
    import zstandard as zstd

    try:
        return zstd.ZstdDecompressor().decompress(body)
    except zstd.ZstdError as exc:
        raise ValueError("invalid zstd binary activation") from exc


def _parse_shape(value: str) -> list[int]:
    shape = [int(part) for part in value.split(",")]
    if not shape or any(dim <= 0 for dim in shape):
        raise ValueError("invalid binary activation shape")
    return shape


def _validate_activation_body(body: bytes, shape: list[int], dtype: str) -> None:
    if dtype not in _DTYPE_SIZES:
        raise ValueError("unsupported binary activation dtype")
    if len(body) != len(_make_stub_binary_activation(shape, dtype)):
        raise ValueError("binary activation byte length does not match headers")


def _stub_response_prefix(chunk_responses: list[_BinaryActivation]) -> str:
    if not chunk_responses:
        return "stub response to:"
    return chunk_responses[-1].headers.get("x-meshnet-stub-response-prefix", "stub response to:")


def _get_json(url: str, timeout: float = 5.0) -> dict:
    with urllib.request.urlopen(url, timeout=timeout) as r:
        return json.loads(r.read())


def _get_tracker_nodes(server: _GatewayHTTPServer, model: str) -> list[str]:
    """Return tracker-node endpoint URLs for this model, or empty list on any failure."""
    if server.tracker_url is None:
        return []
    try:
        resp = _get_json(f"{server.tracker_url}/v1/tracker-nodes/{urllib.parse.quote(model)}")
        return [n["endpoint"] for n in resp.get("tracker_nodes", [])]
    except Exception:
        return []


class GatewayServer:
    """HTTP gateway that routes /v1/chat/completions through an ordered inference route.

    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].

    Optional ``model_presets`` lists the model names returned by GET /v1/models when using
    a static inference route. Ignored when using tracker_url (tracker provides the list).
    """

    def __init__(
        self,
        node_url: str | None = None,
        *,
        inference_route: list[str] | None = None,
        tracker_url: str | None = None,
        model_presets: list[str] | None = None,
        contracts: Any | None = None,
        minimum_stake: int = 0,
        cost_per_layer_token_lamport: int = 1,
        host: str = "127.0.0.1",
        port: int = 0,
    ):
        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]
        if inference_route is not None and not inference_route:
            raise ValueError("inference_route must contain at least one node URL")

        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._model_presets: list[str] | None = model_presets
        self._contracts = contracts
        self._minimum_stake = minimum_stake
        self._cost_per_layer_token_lamport = cost_per_layer_token_lamport
        self._host = host
        self._requested_port = port
        self._server: _GatewayHTTPServer | None = None
        self._thread: threading.Thread | None = None
        self.port: int | None = None

    @property
    def last_binary_chunk_responses(self) -> list[_BinaryActivation]:
        """Most recent binary chunk responses returned by the final route node."""
        return self._server.last_binary_chunk_responses if self._server is not None else []

    def start(self) -> int:
        if self._server is not None:
            raise RuntimeError("GatewayServer is already running")

        self._server = _GatewayHTTPServer(
            (self._host, self._requested_port),
            _GatewayHandler,
            self._inference_route,
            self._tracker_url,
            self._model_presets,
            self._contracts,
            self._minimum_stake,
            self._cost_per_layer_token_lamport,
        )
        self.port = self._server.server_address[1]
        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
        self._thread.start()
        return self.port

    def stop(self) -> None:
        if self._server is None:
            return

        self._server.shutdown()
        self._server.server_close()
        if self._thread is not None:
            self._thread.join(timeout=1)
        self._server = None
        self._thread = None
        self.port = None