integration/Plugins/UEWingman/ue-wingman.py

#!/usr/bin/env python3
"""
Human-friendly MCP test client.

Usage: ue-wingman.py <command> [key=value ...]

Values starting with '[' or '{' are parsed as JSON.
"""

import sys
import json
import socket

HOST = "localhost"
PORT = 9851
TIMEOUT = 120


def main():
    args = sys.argv[1:]
    if not args:
        print("Usage: ue-wingman.py <command> [key=value ...]")
        sys.exit(1)

    msg = {"command": args[0]}
    for arg in args[1:]:
        key, _, value = arg.partition("=")
        if value and value[0] in ('[', '{'):
            try:
                value = json.loads(value)
            except json.JSONDecodeError as e:
                print(f"Bad JSON in {key}: {e.msg}")
                sys.exit(1)
        msg[key] = value

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(TIMEOUT)
    try:
        sock.connect((HOST, PORT))
    except (ConnectionRefusedError, socket.timeout, OSError) as e:
        print(f"Cannot connect to {HOST}:{PORT} — is the editor running?")
        sys.exit(1)

    sock.sendall((json.dumps(msg) + "\0").encode())

    result = b""
    while True:
        chunk = sock.recv(65536)
        if not chunk:
            break
        result += chunk
        if b"\0" in result:
            break

    sock.close()
    result = result[:result.index(b"\0")].decode() if b"\0" in result else result.decode()

    try:
        parsed = json.loads(result)
    except json.JSONDecodeError:
        print("Error: response is not valid JSON.")
        sys.exit(1)

    if not isinstance(parsed, list):
        print("Error: response is not a list of content blocks.")
        sys.exit(1)
    for block in parsed:
        if not (isinstance(block, dict)
                and block.get("type") == "text"
                and isinstance(block.get("text"), str)):
            print("Error: response contains a non-text block.")
            sys.exit(1)

    print("\n---\n".join(block["text"] for block in parsed))

if __name__ == "__main__":
    main()