Aether/tests/e2e_hub_relay.py

"""
aether-gateway tunnel runtime harness 端到端测试

测试流程:
1. 启动 aether-gateway 的 tunnel runtime harness（绑定随机端口）
2. 用 websockets 库模拟一个 aether-proxy client 连接到 tunnel runtime
3. Mock proxy 在收到请求帧后返回固定响应帧
4. 通过 /api/internal/tunnel/relay/{node_id} HTTP API 发送请求
5. 验证完整链路: HTTP request -> tunnel runtime -> WS frame -> mock proxy -> WS frame -> tunnel runtime -> HTTP response

运行: uv run python tests/e2e_hub_relay.py
"""

from __future__ import annotations

import asyncio
import gzip
import json
import os
import signal
import struct
import subprocess
import sys
import time

import httpx

# ---------------------------------------------------------------------------
# Protocol constants (mirror apps/aether-gateway/src/tunnel/embedded/protocol.rs)
# ---------------------------------------------------------------------------
HEADER_SIZE = 10

REQUEST_HEADERS = 0x01
REQUEST_BODY = 0x02
RESPONSE_HEADERS = 0x03
RESPONSE_BODY = 0x04
STREAM_END = 0x05
STREAM_ERROR = 0x06
PING = 0x10
PONG = 0x11
GOAWAY = 0x12

FLAG_END_STREAM = 0x01
FLAG_GZIP_COMPRESSED = 0x02


def encode_frame(stream_id: int, msg_type: int, flags: int, payload: bytes) -> bytes:
    header = struct.pack(">I", stream_id) + bytes([msg_type, flags]) + struct.pack(">I", len(payload))
    return header + payload


def parse_frame(data: bytes) -> tuple[int, int, int, bytes] | None:
    if len(data) < HEADER_SIZE:
        return None
    stream_id = struct.unpack(">I", data[0:4])[0]
    msg_type = data[4]
    flags = data[5]
    payload_len = struct.unpack(">I", data[6:10])[0]
    if len(data) < HEADER_SIZE + payload_len:
        return None
    payload = data[HEADER_SIZE : HEADER_SIZE + payload_len]
    if flags & FLAG_GZIP_COMPRESSED:
        payload = gzip.decompress(payload)
    return stream_id, msg_type, flags, payload


def encode_relay_envelope(meta: dict, body: bytes) -> bytes:
    meta_json = json.dumps(meta, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
    return struct.pack("!I", len(meta_json)) + meta_json + body


# ---------------------------------------------------------------------------
# Mock aether-proxy: connects to Hub via WebSocket, handles request frames
# ---------------------------------------------------------------------------
async def mock_proxy(hub_ws_url: str, node_id: str, ready_event: asyncio.Event) -> None:
    """Simulate an aether-proxy node that echoes requests as fixed responses."""
    try:
        import websockets
    except ImportError:
        print("SKIP: websockets package not installed (uv pip install websockets)")
        sys.exit(1)

    headers = {
        "X-Node-ID": node_id,
        "X-Node-Name": f"test-{node_id}",
    }

    async with websockets.connect(
        hub_ws_url,
        additional_headers=headers,
        max_size=64 * 1024 * 1024,
    ) as ws:
        ready_event.set()
        print(f"  [mock-proxy] connected to hub as node_id={node_id}")

        request_meta: dict | None = None
        request_body: bytes = b""

        async for raw_msg in ws:
            if not isinstance(raw_msg, bytes):
                continue

            parsed = parse_frame(raw_msg)
            if parsed is None:
                continue

            stream_id, msg_type, flags, payload = parsed

            if msg_type == PING:
                await ws.send(encode_frame(0, PONG, 0, payload))
                continue

            if msg_type == REQUEST_HEADERS:
                request_meta = json.loads(payload)
                print(f"  [mock-proxy] stream={stream_id} got REQUEST_HEADERS: {request_meta.get('method')} {request_meta.get('url')}")

            elif msg_type == REQUEST_BODY:
                request_body = payload
                is_end = bool(flags & FLAG_END_STREAM)
                print(f"  [mock-proxy] stream={stream_id} got REQUEST_BODY ({len(payload)} bytes, end={is_end})")

                if is_end and request_meta:
                    # Send response: 200 OK with echoed body
                    resp_meta = {
                        "status": 200,
                        "headers": [
                            ["content-type", "application/json"],
                            ["x-test-echo", "true"],
                        ],
                    }
                    resp_meta_json = json.dumps(resp_meta, separators=(",", ":")).encode("utf-8")
                    await ws.send(encode_frame(stream_id, RESPONSE_HEADERS, 0, resp_meta_json))

                    echo_body = json.dumps({
                        "echo": True,
                        "received_method": request_meta.get("method"),
                        "received_url": request_meta.get("url"),
                        "received_body_len": len(request_body),
                    }, separators=(",", ":")).encode("utf-8")
                    await ws.send(encode_frame(stream_id, RESPONSE_BODY, 0, echo_body))
                    await ws.send(encode_frame(stream_id, STREAM_END, 0, b""))
                    print(f"  [mock-proxy] stream={stream_id} sent response (200, {len(echo_body)} bytes)")

                    request_meta = None
                    request_body = b""

            elif msg_type == STREAM_ERROR:
                error_msg = payload.decode("utf-8", errors="replace")
                print(f"  [mock-proxy] stream={stream_id} got STREAM_ERROR: {error_msg}")


# ---------------------------------------------------------------------------
# Test runner
# ---------------------------------------------------------------------------
async def run_test() -> bool:
    hub_port = 18085
    hub_bind = f"127.0.0.1:{hub_port}"
    hub_binary = os.path.join(
        os.path.dirname(__file__),
        "..",
        "target",
        "release",
        "examples",
        "tunnel_runtime_harness",
    )
    hub_binary = os.path.normpath(hub_binary)

    if not os.path.isfile(hub_binary):
        hub_binary = os.path.join(
            os.path.dirname(__file__),
            "..",
            "target",
            "debug",
            "examples",
            "tunnel_runtime_harness",
        )
        hub_binary = os.path.normpath(hub_binary)
        if not os.path.isfile(hub_binary):
            print(f"FAIL: tunnel runtime harness binary not found at {hub_binary}")
            print("  run: cargo build -p aether-gateway --example tunnel_runtime_harness")
            return False

    # Start tunnel runtime harness (with control plane disabled since we don't have the app running)
    print(f"[1/5] Starting tunnel runtime harness on {hub_bind} ...")
    hub_proc = subprocess.Popen(
        [
            hub_binary,
            "--bind", hub_bind,
            "--proxy-idle-timeout", "0",
            "--ping-interval", "30",
            # Use a non-existent app URL -- control plane callbacks will fail silently
            "--app-base-url", "http://127.0.0.1:19999",
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )

    try:
        # Wait for runtime to be ready
        for _ in range(30):
            await asyncio.sleep(0.2)
            try:
                async with httpx.AsyncClient() as client:
                    resp = await client.get(f"http://{hub_bind}/health", timeout=1.0)
                    if resp.status_code == 200:
                        print("  tunnel runtime harness is healthy")
                        break
            except Exception:
                continue
        else:
            print("FAIL: tunnel runtime harness did not start in time")
            return False

        # Check initial stats
        async with httpx.AsyncClient() as client:
            stats = (await client.get(f"http://{hub_bind}/stats")).json()
            print(f"  initial stats: {stats}")
            assert stats["proxy_connections"] == 0
            assert stats["nodes"] == 0

        # Start mock proxy
        node_id = "test-node-e2e"
        proxy_ready = asyncio.Event()
        print(f"\n[2/5] Connecting mock proxy (node_id={node_id}) ...")
        proxy_task = asyncio.create_task(
            mock_proxy(
                f"ws://{hub_bind}/api/internal/proxy-tunnel", node_id, proxy_ready
            )
        )

        try:
            await asyncio.wait_for(proxy_ready.wait(), timeout=5.0)
        except asyncio.TimeoutError:
            print("FAIL: mock proxy did not connect in time")
            return False

        # Give runtime a moment to register
        await asyncio.sleep(0.3)

        async with httpx.AsyncClient() as client:
            stats = (await client.get(f"http://{hub_bind}/stats")).json()
            print(f"  stats after connect: {stats}")
            assert stats["proxy_connections"] == 1, f"expected 1 proxy connection, got {stats['proxy_connections']}"
            assert stats["nodes"] == 1

        # Send request through local relay
        print(f"\n[3/5] Sending request via local relay ...")
        request_body = b'{"model":"test","messages":[]}'
        envelope = encode_relay_envelope(
            {
                "method": "POST",
                "url": "https://api.example.com/v1/chat/completions",
                "headers": {
                    "content-type": "application/json",
                    "authorization": "Bearer sk-test-123",
                },
                "timeout": 30,
            },
            request_body,
        )

        async with httpx.AsyncClient() as client:
            relay_url = f"http://{hub_bind}/api/internal/tunnel/relay/{node_id}"
            resp = await client.post(
                relay_url,
                content=envelope,
                headers={"content-type": "application/vnd.aether.tunnel-envelope"},
                timeout=10.0,
            )

        print(f"  relay response: status={resp.status_code}")
        assert resp.status_code == 200, f"expected 200, got {resp.status_code}: {resp.text}"

        echo = resp.json()
        print(f"  echo body: {echo}")
        assert echo["echo"] is True
        assert echo["received_method"] == "POST"
        assert echo["received_url"] == "https://api.example.com/v1/chat/completions"
        assert echo["received_body_len"] == len(request_body)

        assert resp.headers.get("x-test-echo") == "true"

        # Verify active streams cleaned up
        print(f"\n[4/5] Verifying stream cleanup ...")
        await asyncio.sleep(0.2)
        async with httpx.AsyncClient() as client:
            stats = (await client.get(f"http://{hub_bind}/stats")).json()
            print(f"  stats after request: {stats}")
            assert stats["active_streams"] == 0, f"expected 0 active streams, got {stats['active_streams']}"

        # Test error case: request to non-existent node
        print(f"\n[5/5] Testing error cases ...")
        async with httpx.AsyncClient() as client:
            resp = await client.post(
                f"http://{hub_bind}/api/internal/tunnel/relay/non-existent-node",
                content=encode_relay_envelope(
                    {"method": "GET", "url": "https://example.com", "headers": {}, "timeout": 5},
                    b"",
                ),
                headers={"content-type": "application/vnd.aether.tunnel-envelope"},
                timeout=5.0,
            )
            assert resp.status_code == 503, f"expected 503 for missing node, got {resp.status_code}"
            assert resp.headers.get("x-aether-tunnel-error") == "connect"
            print(f"  missing node: status={resp.status_code}, error={resp.text}")

        # Test: request from non-loopback should be rejected
        # (can't easily test from non-loopback, but verify header is present for valid errors)

        # Cleanup: cancel proxy
        proxy_task.cancel()
        try:
            await proxy_task
        except asyncio.CancelledError:
            pass

        print("\n" + "=" * 50)
        print("ALL TESTS PASSED")
        print("=" * 50)
        return True

    finally:
        hub_proc.send_signal(signal.SIGTERM)
        try:
            hub_proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            hub_proc.kill()
            hub_proc.wait()
        print("\n[cleanup] hub process stopped")


if __name__ == "__main__":
    success = asyncio.run(run_test())
    sys.exit(0 if success else 1)