- Pass the /ask question to claude on stdin so it can't be parsed as a CLI option. - Block run_lua in multiplayer unless an admin enables the new "Allow Lua queries in multiplayer" map setting. - Ignore UDP packets that don't come from the game's port, in both the MCP server and the ask helper. - Build the screenshot path from our own name instead of the game's reply, and restrict screenshot names in the mod. - Validate ask packets, cap question length and limit concurrent claude runs. - Replace hard-coded home paths in the README and systemd unit. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
93 lines
3.9 KiB
Python
93 lines
3.9 KiB
Python
"""UDP request/response link to the claude-co-engineer mod running inside the game."""
|
|
|
|
import asyncio
|
|
import itertools
|
|
import json
|
|
import socket
|
|
from typing import Any
|
|
|
|
from mcp.server.mcpserver.exceptions import ToolError
|
|
|
|
|
|
class GameError(ToolError):
|
|
"""The game couldn't answer a query. A ToolError so the message reaches the model."""
|
|
|
|
|
|
class _Protocol(asyncio.DatagramProtocol):
|
|
def __init__(self, link: "GameLink"):
|
|
self.link = link
|
|
|
|
def datagram_received(self, data: bytes, addr: Any) -> None:
|
|
# The game replies from its --enable-lua-udp port; ignore anything else on localhost.
|
|
if addr[:2] == self.link.peer:
|
|
self.link._on_packet(data)
|
|
|
|
|
|
class GameLink:
|
|
"""Sends JSON requests to the game's --enable-lua-udp port and reassembles chunked replies.
|
|
|
|
The mod replies to whatever port a request came from, so we bind an ephemeral port.
|
|
"""
|
|
|
|
def __init__(self, game_port: int, host: str = "127.0.0.1", timeout: float = 5.0):
|
|
self.game_addr = (host, game_port)
|
|
self.host = host
|
|
self.timeout = timeout
|
|
self._ids = itertools.count(1)
|
|
self._pending: dict[int, asyncio.Future] = {}
|
|
self._parts: dict[int, dict[int, bytes]] = {}
|
|
self._transport: asyncio.DatagramTransport | None = None
|
|
|
|
async def _ensure_open(self) -> asyncio.DatagramTransport:
|
|
if self._transport is None or self._transport.is_closing():
|
|
loop = asyncio.get_running_loop()
|
|
# The address replies come from, resolved so a hostname like "localhost" matches.
|
|
self.peer = (socket.gethostbyname(self.host), self.game_addr[1])
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
# Large replies arrive as a burst of ~32KB packets; give the kernel room to queue them.
|
|
sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 8 * 1024 * 1024)
|
|
sock.bind((self.host, 0))
|
|
self._transport, _ = await loop.create_datagram_endpoint(lambda: _Protocol(self), sock=sock)
|
|
return self._transport
|
|
|
|
def _on_packet(self, data: bytes) -> None:
|
|
try:
|
|
raw_id, raw_seq, raw_total, fragment = data.split(b"|", 3)
|
|
request_id, seq, total = int(raw_id), int(raw_seq), int(raw_total)
|
|
except ValueError:
|
|
return
|
|
future = self._pending.get(request_id)
|
|
if future is None or future.done():
|
|
return
|
|
parts = self._parts.setdefault(request_id, {})
|
|
parts[seq] = fragment
|
|
if len(parts) < total:
|
|
return
|
|
del self._parts[request_id]
|
|
try:
|
|
future.set_result(json.loads(b"".join(parts[i] for i in range(1, total + 1))))
|
|
except (KeyError, ValueError) as e:
|
|
future.set_exception(GameError(f"malformed reply from game: {e}"))
|
|
|
|
async def request(self, op: str, **args: Any) -> Any:
|
|
transport = await self._ensure_open()
|
|
request_id = next(self._ids)
|
|
future = asyncio.get_running_loop().create_future()
|
|
self._pending[request_id] = future
|
|
payload = {"id": request_id, "op": op, "args": {k: v for k, v in args.items() if v is not None}}
|
|
try:
|
|
transport.sendto(json.dumps(payload).encode(), self.game_addr)
|
|
response = await asyncio.wait_for(future, self.timeout)
|
|
except TimeoutError:
|
|
raise GameError(
|
|
f"No reply from Factorio within {self.timeout:g}s on UDP port {self.game_addr[1]}. "
|
|
"Check that the game was launched with --enable-lua-udp on that port, the claude-co-engineer "
|
|
"mod is enabled, a save is loaded, and the game isn't paused (the Esc menu pauses singleplayer)."
|
|
) from None
|
|
finally:
|
|
self._pending.pop(request_id, None)
|
|
self._parts.pop(request_id, None)
|
|
if not response.get("ok"):
|
|
raise GameError(response.get("error", "unknown error"))
|
|
return response.get("result")
|