예제 소스

omook_ws_bridge.py

#!/usr/bin/env python3
"""Omook 웹 ↔ vlm_pva_server TCP 브리지.

브라우저 WebSocket(JSON) → 로컬 OMK3 TCP (:9299)

  .venv/bin/python scripts/omook_ws_bridge.py --vlm-port 9299 --ws-port 9322

프로토콜 (JSON):
  → {"type":"move","board":[N*N ints],"size":13|15|19,"ai_side":1|2}
  ← {"type":"move","row":int,"col":int}
  ← {"type":"error","error":"..."}
  → {"type":"ping"} ← {"type":"pong"}
"""

from __future__ import annotations

import argparse
import asyncio
import io
import json
import socket
import struct
import sys
from pathlib import Path
from typing import List, Optional, Tuple

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from dataset_omook import render_board_png  # noqa: E402

try:
    import websockets
except ImportError:
    print("websockets 패키지 필요: .venv/bin/pip install websockets", file=sys.stderr)
    raise SystemExit(1)

try:
    from websockets.asyncio.server import serve as ws_serve
except ImportError:
    from websockets import serve as ws_serve  # type: ignore

MAGIC_OMK3 = b"OMK3"
MAGIC_RPLY = b"RPLY"
VALID_SIZES = (13, 15, 19)


def _u32be(n: int) -> bytes:
    return struct.pack(">I", int(n) & 0xFFFFFFFF)


def _recv_exact(sock: socket.socket, n: int) -> bytes:
    buf = b""
    while len(buf) < n:
        chunk = sock.recv(n - len(buf))
        if not chunk:
            raise EOFError("socket closed")
        buf += chunk
    return buf


def _board_png(board: List[int], size: int) -> bytes:
    img = render_board_png(board, size, cell_px=28, pad=14)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()


def _send_omk3(sock: socket.socket, png: bytes, board: bytes, size: int, ai_side: int) -> None:
    body = bytearray()
    body.append(size & 0xFF)
    body.append(ai_side & 0xFF)
    body += _u32be(len(png))
    body += png
    body += board
    sock.sendall(MAGIC_OMK3 + _u32be(len(body)) + body)


def _recv_reply(sock: socket.socket) -> Tuple[bool, Optional[int], Optional[int], Optional[str]]:
    head = _recv_exact(sock, 4)
    if head != MAGIC_RPLY:
        raise ValueError(f"bad reply header: {head!r}")
    flag = _recv_exact(sock, 1)[0]
    if flag == 1:
        row = _recv_exact(sock, 1)[0]
        col = _recv_exact(sock, 1)[0]
        return True, int(row), int(col), None
    if flag == 0:
        (n,) = struct.unpack(">H", _recv_exact(sock, 2))
        msg = _recv_exact(sock, n).decode("utf-8", errors="replace")
        return False, None, None, msg
    raise ValueError(f"unknown flag: {flag}")


def request_move_sync(
    host: str,
    port: int,
    board: List[int],
    size: int,
    ai_side: int,
    timeout_s: float,
) -> dict:
    if size not in VALID_SIZES:
        return {"type": "error", "error": f"unsupported size: {size}"}
    if len(board) != size * size:
        return {"type": "error", "error": f"board length must be {size * size}"}
    packed = bytes(int(x) & 0xFF for x in board)
    sock = socket.create_connection((host, port), timeout=15)
    sock.settimeout(timeout_s)
    try:
        _send_omk3(sock, _board_png(board, size), packed, size, ai_side)
        ok, row, col, err = _recv_reply(sock)
        if not ok:
            return {"type": "error", "error": err or "vlm error"}
        if row is None or col is None:
            return {"type": "error", "error": "empty move"}
        return {"type": "move", "row": row, "col": col}
    finally:
        try:
            sock.close()
        except Exception:
            pass


async def handler(ws, *, vlm_host: str, vlm_port: int, timeout_s: float):
    async for raw in ws:
        try:
            msg = json.loads(raw)
        except json.JSONDecodeError:
            await ws.send(json.dumps({"type": "error", "error": "invalid json"}))
            continue
        mtype = msg.get("type")
        if mtype == "ping":
            await ws.send(json.dumps({"type": "pong"}))
            continue
        if mtype != "move":
            await ws.send(json.dumps({"type": "error", "error": f"unknown type: {mtype}"}))
            continue
        board = msg.get("board") or []
        size = int(msg.get("size") or 15)
        ai_side = int(msg.get("ai_side") or 2)
        try:
            result = await asyncio.to_thread(
                request_move_sync,
                vlm_host,
                vlm_port,
                list(board),
                size,
                ai_side,
                timeout_s,
            )
        except Exception as e:
            result = {"type": "error", "error": str(e)}
        await ws.send(json.dumps(result, ensure_ascii=False))


async def main_async(args: argparse.Namespace) -> None:
    host = args.ws_host
    print(
        f"[omook-bridge] ws://{host}:{args.ws_port} → VLM {args.vlm_host}:{args.vlm_port}",
        flush=True,
    )
    if host in ("127.0.0.1", "localhost"):
        print(
            "[omook-bridge] WARNING: ws-host 가 localhost 전용입니다. "
            "LAN(예: 192.168.x.x)에서 접속하려면 --ws-host ******** 이 필요합니다.",
            flush=True,
        )

    async def _handler(ws):
        await handler(ws, vlm_host=args.vlm_host, vlm_port=args.vlm_port, timeout_s=args.timeout)

    async with ws_serve(_handler, host, args.ws_port):
        await asyncio.Future()


def main() -> None:
    p = argparse.ArgumentParser(description="Omook WebSocket ↔ OMK3 bridge")
    p.add_argument("--vlm-host", default="127.0.0.1")
    p.add_argument("--vlm-port", type=int, default=9299)
    p.add_argument(
        "--ws-host",
        default="********",
        help="브리지 바인드 주소. 폰/LAN 접속이면 ******** (기본).",
    )
    p.add_argument("--ws-port", type=int, default=9322)
    p.add_argument("--timeout", type=float, default=240.0)
    args = p.parse_args()
    asyncio.run(main_async(args))


if __name__ == "__main__":
    main()