claude-co-engineer/server/claude_co_engineer/ask.py

"""Answers in-game /ask questions by running Claude Code headless with the claude-co-engineer MCP server.

The mod sends {"type": "ask", "id", "player_index", "question"} to FACTORIO_ASK_PORT. We acknowledge it,
run `claude -p` (which uses the CLI's own login, i.e. your subscription), relay which tools Claude is
using as status lines, and send the final answer back through the game's UDP port.
"""

import asyncio
import json
import os
import re
import shutil
import sys
from pathlib import Path

from .game_link import GameError, GameLink

GAME_PORT = int(os.environ.get("FACTORIO_UDP_PORT", "34200"))
ASK_PORT = int(os.environ.get("FACTORIO_ASK_PORT", "34201"))
CLAUDE = os.environ.get("FACTORIO_ASK_CLAUDE") or shutil.which("claude") or "claude"
# Used when the mod doesn't send a model (the in-game "/ask model" setting normally does).
DEFAULT_MODEL = os.environ.get("FACTORIO_ASK_MODEL", "sonnet")
# Wins over the in-game setting, e.g. to try a model the dropdown doesn't list yet.
MODEL_OVERRIDE = os.environ.get("FACTORIO_ASK_MODEL_OVERRIDE")
TIMEOUT = float(os.environ.get("FACTORIO_ASK_TIMEOUT", "300"))
MAX_CONCURRENT = int(os.environ.get("FACTORIO_ASK_MAX_CONCURRENT", "2"))
MAX_QUESTION = 2000  # characters; longer questions are cut off
WORKDIR = Path(os.environ.get("FACTORIO_ASK_WORKDIR", "~/.cache/claude-co-engineer/ask")).expanduser()

SYSTEM_PROMPT = """\
The user typed this question into Factorio's in-game chat with /ask, and your final answer is printed
in that chat. Use the factorio tools to look at their game before answering.
Format for the chat window: plain text, no Markdown (no **, #, backticks or tables), at most about
6 short lines. Factorio rich text works and is encouraged: [gps=x,y] for clickable map locations,
[item=iron-plate], [entity=assembling-machine-2], [fluid=water], [technology=automation] for icons.
When you point at places, also call highlight with chat=false (your answer already has the links).
Advise; don't play the game for them.
"""

STATUS = {
    "get_overview": "looking around",
    "get_player": "checking your inventory",
    "inspect_entity": "inspecting that",
    "find_entities": "searching the area",
    "find_problems": "checking for problems",
    "get_production": "reading production stats",
    "get_power": "checking power",
    "get_trends": "looking at trends",
    "trace_belt": "following the belt",
    "get_research": "checking research",
    "get_recipe": "looking up recipes",
    "calculate_ratios": "calculating ratios",
    "screenshot": "taking a look",
    "highlight": "marking it on your map",
    "run_lua": "digging deeper",
}


def strip_markdown(text: str) -> str:
    text = re.sub(r"\*\*(.+?)\*\*", r"\1", text)
    text = re.sub(r"`([^`]+)`", r"\1", text)
    text = re.sub(r"^#+\s*", "", text, flags=re.MULTILINE)
    return text.strip()


class AskServer(asyncio.DatagramProtocol):
    def __init__(self):
        self.game = GameLink(GAME_PORT)
        self.sessions: dict[int, str] = {}  # player_index -> Claude session id
        self.locks: dict[int, asyncio.Lock] = {}
        self.running = asyncio.Semaphore(MAX_CONCURRENT)  # claude processes across all players
        self.tasks: set[asyncio.Task] = set()  # keeps running asks referenced until they finish
        WORKDIR.mkdir(parents=True, exist_ok=True)
        self.mcp_config = WORKDIR / "mcp.json"
        self.mcp_config.write_text(json.dumps({
            "mcpServers": {
                "factorio": {
                    "command": sys.executable,
                    "args": ["-m", "claude_co_engineer.server"],
                    "env": {"FACTORIO_UDP_PORT": str(GAME_PORT)},
                }
            }
        }))

    def datagram_received(self, data: bytes, addr) -> None:
        # The mod sends from the game's --enable-lua-udp port; other local senders are ignored.
        if addr[:2] != ("127.0.0.1", GAME_PORT):
            return
        try:
            message = json.loads(data)
        except ValueError:
            return
        if not isinstance(message, dict) or not isinstance(player := message.get("player_index"), int):
            return
        if message.get("type") == "reset":
            self.sessions.pop(player, None)
            log(f"player {player}: conversation reset")
        elif message.get("type") == "ask":
            ask_id, question = message.get("id"), message.get("question")
            if not isinstance(ask_id, int) or not isinstance(question, str) or not question.strip():
                return
            task = asyncio.ensure_future(self.handle_ask(ask_id, player, question[:MAX_QUESTION], message.get("model")))
            self.tasks.add(task)
            task.add_done_callback(self.tasks.discard)

    async def send(self, op: str, **args) -> None:
        try:
            await self.game.request(op, **args)
        except GameError as e:
            log(f"couldn't reach the game for {op}: {e}")

    async def handle_ask(self, ask_id: int, player: int, question: str, model: str | None) -> None:
        model = MODEL_OVERRIDE or model or DEFAULT_MODEL
        if not isinstance(model, str) or not re.fullmatch(r"[\w.-]+", model):
            model = DEFAULT_MODEL
        log(f"player {player} asks #{ask_id} ({model}): {question}")
        await self.send("ask_ack", ask_id=ask_id)
        async with self.locks.setdefault(player, asyncio.Lock()), self.running:
            try:
                answer, is_error = await asyncio.wait_for(self.run_claude(ask_id, player, question, model), TIMEOUT)
            except TimeoutError:
                answer, is_error = f"Gave up after {TIMEOUT:g}s without an answer.", True
            except Exception as e:  # noqa: BLE001 - report anything to the player rather than dying
                answer, is_error = f"The ask helper failed: {e}", True
        log(f"answer #{ask_id} ({'error' if is_error else 'ok'}): {answer}")
        await self.send("ask_answer", ask_id=ask_id, player_index=player, text=answer, error=is_error)

    async def run_claude(self, ask_id: int, player: int, question: str, model: str) -> tuple[str, bool]:
        # The question goes in on stdin: as an argument, one starting with "-" would be parsed as a CLI option.
        args = [
            CLAUDE, "-p",
            "--model", model,
            "--mcp-config", str(self.mcp_config), "--strict-mcp-config",
            "--tools", "",  # no shell, file or web tools; only the game
            "--allowedTools", "mcp__factorio",
            "--append-system-prompt", SYSTEM_PROMPT,
            "--output-format", "stream-json", "--verbose",
        ]
        if player in self.sessions:
            args += ["--resume", self.sessions[player]]

        process = await asyncio.create_subprocess_exec(
            *args, cwd=WORKDIR, stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, limit=16 * 1024 * 1024,
        )
        process.stdin.write(question.encode())
        process.stdin.close()
        stderr_task = asyncio.ensure_future(process.stderr.read())
        result, last_status = None, None
        try:
            async for line in process.stdout:
                try:
                    event = json.loads(line)
                except ValueError:
                    continue
                if event.get("type") == "assistant":
                    for block in event["message"].get("content", []):
                        if block.get("type") == "tool_use":
                            status = STATUS.get(block["name"].removeprefix("mcp__factorio__"), "working")
                            if status != last_status:
                                last_status = status
                                await self.send("ask_status", ask_id=ask_id, text=status + "...")
                elif event.get("type") == "result":
                    result = event
            await process.wait()
        finally:
            if process.returncode is None:
                process.kill()

        if result is None:
            stderr = (await stderr_task).decode(errors="replace").strip()
            return f"Claude exited without an answer: {stderr[-300:] or f'exit code {process.returncode}'}", True
        if result.get("session_id"):
            self.sessions[player] = result["session_id"]
        text = strip_markdown(result.get("result") or "")
        if "Not logged in" in text or "/login" in text:
            return "The claude CLI isn't logged in. Run 'claude auth login' in a terminal, then ask again.", True
        return text or "(no answer)", bool(result.get("is_error"))


def log(message: str) -> None:
    print(message, file=sys.stderr, flush=True)


async def serve() -> None:
    loop = asyncio.get_running_loop()
    await loop.create_datagram_endpoint(AskServer, local_addr=("127.0.0.1", ASK_PORT))
    log(f"claude-co-engineer-ask listening on 127.0.0.1:{ASK_PORT}, game on {GAME_PORT}, claude = {CLAUDE}")
    await asyncio.Event().wait()


def main() -> None:
    try:
        asyncio.run(serve())
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()