devbox-images/devbox/base-config/claude/peer-mcp-bridge-claude.py

#!/usr/bin/env python3
# MCP stdio server exposing one tool ("ask") that hands a prompt to a peer
# workspace's Claude Code over TCP and returns its final result — the
# Claude Code counterpart to kilo/peer-mcp-bridge.py.
#
# Claude Code has no live agent-protocol session to drive the way `kilo
# acp` gives peer-mcp-bridge.py (no session/new, no streamed tool_call
# updates, no mid-run permission callback to answer). Instead this talks to
# claude-peer-listen.sh's socat listener, which runs a single headless
# `claude -p --permission-mode bypassPermissions` invocation per connection
# via claude-peer-runner.py and returns one `--output-format json` object
# once that run finishes. So this bridge is simpler than the kilo one: send
# the token, send the length-prefixed prompt, read the response until the
# peer closes its end, and parse the one JSON object out of it.
#
# The prompt is length-prefixed (4-byte big-endian byte count, then that
# many bytes), NOT terminated by half-closing the connection
# (`shutdown(SHUT_WR)`) the way an earlier version of this did. Confirmed
# live against real workspace containers: that half-close made socat tear
# the connection down (or stop relaying output — not fully diagnosed
# which) the instant the far side's response took more than an instant to
# produce, regardless of what was actually running there — reproduced with
# a plain `sleep N; echo` script that never touched Claude Code at all. A
# real `claude -p` run doing actual work is never "an instant," so
# half-closing here would silently and consistently return nothing.
import argparse
import json
import socket
import struct
import sys

TOOL_NAME = "ask"


def read_stdin_messages():
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            yield json.loads(line)
        except json.JSONDecodeError:
            continue


def write_message(obj):
    sys.stdout.write(json.dumps(obj) + "\n")
    sys.stdout.flush()


def send_response(msg_id, result=None, error=None):
    msg = {"jsonrpc": "2.0", "id": msg_id}
    if error is not None:
        msg["error"] = error
    else:
        msg["result"] = result
    write_message(msg)


class PeerConnectionError(Exception):
    pass


class ClaudePeer:
    def __init__(self, host, port, timeout, token):
        self.host = host
        self.port = port
        self.timeout = timeout
        self.token = token

    def ask(self, prompt_text):
        try:
            sock = socket.create_connection((self.host, self.port), timeout=15)
        except OSError as e:
            raise PeerConnectionError(f"could not reach peer Claude bridge {self.host}:{self.port}: {e}")

        with sock:
            sock.settimeout(self.timeout)
            sock.sendall((self.token + "\n").encode())
            prompt_bytes = prompt_text.encode()
            sock.sendall(struct.pack("!I", len(prompt_bytes)))
            sock.sendall(prompt_bytes)

            chunks = []
            try:
                while True:
                    chunk = sock.recv(65536)
                    if not chunk:
                        break
                    chunks.append(chunk)
            except socket.timeout:
                raise TimeoutError(f"peer did not finish within {self.timeout}s")

            raw = b"".join(chunks).decode(errors="replace").strip()
            if not raw:
                raise PeerConnectionError("peer closed the connection with no response (bad token?)")

            try:
                payload = json.loads(raw)
            except json.JSONDecodeError:
                return raw  # fall back to raw text rather than hide a malformed-but-useful reply

            if payload.get("is_error"):
                raise PeerConnectionError(payload.get("result", "peer reported an error"))
            return payload.get("result") or "(peer produced no text response)"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", required=True)
    ap.add_argument("--port", type=int, default=4197)
    ap.add_argument("--peer-name", required=True)
    ap.add_argument("--timeout", type=int, default=900)
    args = ap.parse_args()

    # Read from the environment (same as kilo's peer-mcp-bridge.py) rather
    # than a CLI flag, so the shared secret never ends up in the rendered
    # MCP config's command array or in `ps` output.
    import os
    token = os.environ.get("PEER_BRIDGE_TOKEN", "")
    if not token:
        print(json.dumps({"jsonrpc": "2.0", "error": {"code": -32000, "message": "PEER_BRIDGE_TOKEN not set"}}))
        sys.exit(1)

    peer = ClaudePeer(args.host, args.port, args.timeout, token)

    for msg in read_stdin_messages():
        method = msg.get("method")
        msg_id = msg.get("id")

        if method == "initialize":
            client_version = msg.get("params", {}).get("protocolVersion", "2024-11-05")
            send_response(msg_id, result={
                "protocolVersion": client_version,
                "capabilities": {"tools": {}},
                "serverInfo": {"name": f"claude-peer-bridge-{args.peer_name}", "version": "1.0.0"},
            })
        elif method == "notifications/initialized":
            pass
        elif method == "tools/list":
            send_response(msg_id, result={"tools": [{
                "name": TOOL_NAME,
                "description": (
                    f"Ask the '{args.peer_name}' workspace's Claude Code agent a question about "
                    "its project, or hand it a task to work on. It runs with full autonomy in its "
                    "own project directory (reads/edits/writes files, runs bash) and returns its "
                    "final response once done. Each call starts a fresh session on the peer's "
                    "side — include any needed context in the prompt itself."
                ),
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "prompt": {"type": "string", "description": "The question or task for the peer agent."},
                    },
                    "required": ["prompt"],
                },
            }]})
        elif method == "tools/call":
            params = msg.get("params", {})
            if params.get("name") != TOOL_NAME:
                send_response(msg_id, result={
                    "content": [{"type": "text", "text": f"Unknown tool: {params.get('name')}"}],
                    "isError": True,
                })
                continue
            prompt_text = params.get("arguments", {}).get("prompt", "")
            try:
                reply = peer.ask(prompt_text)
                send_response(msg_id, result={"content": [{"type": "text", "text": reply}], "isError": False})
            except Exception as e:
                send_response(msg_id, result={
                    "content": [{"type": "text", "text": f"Peer request failed: {e}"}],
                    "isError": True,
                })
        elif msg_id is not None:
            send_response(msg_id, error={"code": -32601, "message": f"Method not found: {method}"})
        # else: unrecognized notification, ignore


if __name__ == "__main__":
    main()