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

"""MCP server exposing read-only questions about a running Factorio game."""

import asyncio
import functools
import json
import os
import uuid
from pathlib import Path
from typing import Annotated, Any, Literal

from mcp.server import MCPServer
from mcp.server.mcpserver import Image
from mcp.types import ToolAnnotations
from pydantic import BaseModel, Field

from .game_link import GameError, GameLink

INSTRUCTIONS = """\
You are Claude Co-Engineer, a Factorio advisor connected to the user's live game (Factorio 2.x, possibly Space Age).
The user is playing themselves and wants help deciding what to do and understanding their factory -
answer questions, diagnose bottlenecks, suggest next steps. Do not try to play the game for them.

Start with get_overview when you lack context. Prefer the dedicated tools; use run_lua only for
questions they can't answer, and only to read state - never change entities, inventories, research
or anything else in the world.

Positions are map tiles: +x is east, +y is SOUTH (north is -y). The player can't see coordinates
by default, so describe locations relative to them ("about 40 tiles north-east of you") using
their position from get_overview/get_player. Hovered = the entity under the user's mouse cursor;
when they say "this" or "here" they likely mean the hovered entity or their position.
Whenever your answer points the user at specific places (a broken machine, an ore patch, a nest),
also call highlight so they're outlined in the game and linked in chat - but don't spam it.
Production rates are per minute. Recipe 'seconds' is crafting time at speed 1.
Parameters: omitted x/y mean the player's position (or the hovered entity where that makes sense);
surface defaults to the player's; window is a statistics period (5s to 50h).
"""

READ_ONLY = ToolAnnotations(readOnlyHint=True, openWorldHint=False)
Window = Literal["5s", "1m", "10m", "1h", "10h", "50h"]

game = GameLink(
    game_port=int(os.environ.get("FACTORIO_UDP_PORT", "34200")),
    host=os.environ.get("FACTORIO_UDP_HOST", "127.0.0.1"),
    timeout=float(os.environ.get("FACTORIO_TIMEOUT", "5")),
)
SCRIPT_OUTPUT = Path(os.environ.get("FACTORIO_SCRIPT_OUTPUT", "~/.factorio/script-output")).expanduser()
mcp = MCPServer("claude-co-engineer", instructions=INSTRUCTIONS)


def tool(**kwargs):
    """mcp.tool that returns dict results as compact JSON text. The SDK's default pretty-printed JSON
    (plus a structured copy) costs about twice the tokens for the same data."""

    def decorate(fn):
        @functools.wraps(fn)
        async def wrapper(*args, **kw):
            result = await fn(*args, **kw)
            return json.dumps(result, separators=(",", ":")) if isinstance(result, dict) else result

        return mcp.tool(structured_output=False, **kwargs)(wrapper)

    return decorate


@tool(annotations=READ_ONLY)
async def get_overview() -> dict[str, Any]:
    """Player, research and broken-machine counts in one call. Start here."""
    status = await game.request("ping")
    try:
        player = await game.request("player")
    except GameError as e:
        player = {"unavailable": str(e)}
    research = await game.request("research")
    problems = await game.request("find_problems", samples=0)
    return {
        "game": {k: status[k] for k in ("tick", "ticks_played", "game_version", "multiplayer", "allow_lua")},
        "hours_played": round(status["ticks_played"] / 60 / 3600, 2),
        "mods": sorted(status["mods"]),
        "player": player,
        "research": {
            # Lua drops nil fields, so anything optional may be absent.
            "current": research.get("current"),
            "current_progress": research.get("current_progress"),
            "queue": research.get("queue") or [],
            "researched": f'{research["researched"]}/{research["total"]}',
            "available_now": [t["name"] for t in research["available_now"]],
        },
        "machine_problems": {
            # An empty Lua table serialises as [], hence the `or {}`.
            name: {status: p["count"] for status, p in (group.get("problems") or {}).items()} | {"total": group["total"]}
            for name, group in (problems.get("machines") or {}).items()
        },
    }


@tool(annotations=READ_ONLY)
async def get_player(
    player: str | None = None,
) -> dict[str, Any]:
    """Position, inventory, cursor, crafting queue, hovered/opened entity."""
    return await game.request("player", player=player)


@tool(annotations=READ_ONLY)
async def inspect_entity(
    x: float | None = None,
    y: float | None = None,
    surface: str | None = None,
) -> dict[str, Any]:
    """One entity in full: status, recipe, speed, contents, fluids, fuel, power. Default: hovered."""
    return await game.request("inspect_entity", x=x, y=y, surface=surface)


@tool(annotations=READ_ONLY)
async def find_entities(
    name: str | list[str] | None = None,
    type: str | list[str] | None = None,
    radius: Annotated[float, Field(le=1000)] = 32,
    x: float | None = None,
    y: float | None = None,
    force: str | None = None,
    limit: Annotated[int, Field(description="Nearest matches listed")] = 15,
    surface: str | None = None,
) -> dict[str, Any]:
    """Count entities by name near a point (ore totals for type=resource) and list the nearest."""
    return await game.request(
        "find_entities", name=name, type=type, radius=radius, x=x, y=y, force=force, limit=limit, surface=surface
    )


@tool(annotations=READ_ONLY)
async def find_problems(
    radius: Annotated[float | None, Field(description="Default: whole surface")] = None,
    types: Annotated[list[str] | None, Field(description="Default: machines, labs, power, turrets")] = None,
    samples: Annotated[int, Field(description="Nearest examples per machine+status")] = 3,
    include_healthy: bool = False,
    surface: str | None = None,
) -> dict[str, Any]:
    """Non-working machines grouped by machine and status (no_ingredients, full_output, low_power...)."""
    return await game.request(
        "find_problems", radius=radius, types=types, samples=samples, include_healthy=include_healthy, surface=surface
    )


@tool(annotations=READ_ONLY)
async def get_production(
    window: Window = "10m",
    items: Annotated[list[str] | None, Field(description="Default: top movers")] = None,
    top: int = 40,
    surface: str | None = None,
) -> dict[str, Any]:
    """Produced/consumed/net per minute for items and fluids, plus all-time totals."""
    return await game.request("production", window=window, items=items, top=top, surface=surface)


@tool(annotations=READ_ONLY)
async def trace_belt(
    x: float | None = None,
    y: float | None = None,
    direction: Literal["both", "downstream", "upstream"] = "both",
    max_tiles: Annotated[int, Field(ge=10, le=1500)] = 300,
    section_length: Annotated[int, Field(ge=2, le=100)] = 10,
    surface: str | None = None,
) -> dict[str, Any]:
    """Follow a belt (default: hovered): fill, items and motion per section (backed up vs flowing),
    inserters/drills on and off it, and where it ends."""
    return await game.request(
        "trace_belt", x=x, y=y, direction=direction, max_tiles=max_tiles, section_length=section_length, surface=surface
    )


@tool(annotations=READ_ONLY)
async def get_trends(
    items: Annotated[list[str] | None, Field(description="Default: the biggest movers")] = None,
    window: Window = "1h",
    buckets: Literal[5, 6, 10, 12, 15, 20, 30, 60] = 12,
    top: Annotated[int, Field(ge=1, le=30)] = 8,
    power: bool = True,
    surface: str | None = None,
) -> dict[str, Any]:
    """Production/consumption (and power) over time as short series with rising/falling/stopped labels."""
    return await game.request(
        "trends", items=items, window=window, buckets=buckets, top=top, power=power, surface=surface
    )


@tool(annotations=READ_ONLY)
async def get_power(
    window: Window = "1m",
    top: int = 5,
    surface: str | None = None,
) -> dict[str, Any]:
    """kW consumed and produced per electric network, by entity."""
    return await game.request("power", window=window, top=top, surface=surface)


@tool(annotations=READ_ONLY)
async def get_research(
    technology: str | None = None,
) -> dict[str, Any]:
    """Current research and what's researchable now; with technology: its cost, unlocks and missing chain."""
    return await game.request("research", technology=technology)


@tool(annotations=READ_ONLY)
async def get_recipe(
    name: Annotated[str, Field(description="Recipe, item or fluid")],
) -> dict[str, Any]:
    """Recipes making it (machines, unlocked or not) and recipes using it."""
    return await game.request("recipe", name=name)


@tool(annotations=READ_ONLY)
async def calculate_ratios(
    item: str,
    rate_per_minute: Annotated[float | None, Field(gt=0)] = None,
    rate_per_second: Annotated[float | None, Field(gt=0)] = None,
    machines: Annotated[float | None, Field(gt=0, description="Final-step machine count instead of a rate")] = None,
    inputs: Annotated[list[str] | None, Field(description="Already supplied (e.g. bus plates); not expanded")] = None,
    recipe_overrides: Annotated[dict[str, str] | None, Field(description="item -> recipe")] = None,
    machine_overrides: Annotated[dict[str, str] | None, Field(description="item/recipe -> machine or drill")] = None,
    surface: str | None = None,
) -> dict[str, Any]:
    """Machines, drills, belts and power for a target, from the user's unlocked recipes, machines and
    bonuses, vs. what they've built. Use instead of doing ratio math. Give one of the three targets."""
    if sum(v is not None for v in (rate_per_minute, rate_per_second, machines)) != 1:
        raise GameError("Give exactly one of rate_per_minute, rate_per_second or machines.")
    return await game.request(
        "ratio",
        item=item,
        rate_per_minute=rate_per_minute if rate_per_second is None else rate_per_second * 60,
        machines=machines,
        inputs=inputs,
        recipe_overrides=recipe_overrides,
        machine_overrides=machine_overrides,
        surface=surface,
    )


async def _read_finished_jpeg(path: Path, timeout: float = 10.0) -> bytes:
    """The game writes screenshots asynchronously after the tick; wait for a complete JPEG."""
    loop = asyncio.get_running_loop()
    deadline = loop.time() + timeout
    while loop.time() < deadline:
        if path.exists():
            data = path.read_bytes()
            if data.endswith(b"\xff\xd9"):  # JPEG end-of-image marker
                return data
        await asyncio.sleep(0.1)
    raise GameError(
        f"The game accepted the screenshot request but {path} never appeared. Screenshots need a rendering "
        "client (a headless server can't take them), and FACTORIO_SCRIPT_OUTPUT must point at the game's "
        "script-output folder."
    )


@tool(annotations=READ_ONLY)
async def screenshot(
    x: float | None = None,
    y: float | None = None,
    width_tiles: Annotated[float, Field(ge=8, le=1000, description="~30 detail, 60 work area, 200+ whole build")] = 60,
    alt_mode: bool = True,
    daylight: bool = True,
    surface: str | None = None,
) -> list[str | Image]:
    """1280x960 off-screen render of the map for layout questions (~1.7k tokens); returns the area covered."""
    name = f"shot-{uuid.uuid4().hex[:12]}"
    meta = await game.request(
        "screenshot", x=x, y=y, width_tiles=width_tiles, alt_mode=alt_mode, daylight=daylight, surface=surface, name=name
    )
    path = SCRIPT_OUTPUT / meta.pop("path")
    try:
        data = await _read_finished_jpeg(path)
    finally:
        path.unlink(missing_ok=True)
    meta["note"] = "+x is east (right), +y is south (down). Pixel (px, py) is at map (left_top.x + px/pixels_per_tile, left_top.y + py/pixels_per_tile)."
    return [json.dumps(meta, separators=(",", ":")), Image(data=data, format="jpeg")]


class Marker(BaseModel):
    x: float
    y: float
    label: str | None = None
    color: Literal["yellow", "red", "green", "blue", "orange", "cyan", "purple", "white"] = "yellow"
    radius: float | None = Field(default=None, description="Ring an area instead of outlining the building")


VISUAL_ONLY = ToolAnnotations(readOnlyHint=False, destructiveHint=False, idempotentHint=False, openWorldHint=False)


@tool(annotations=VISUAL_ONLY)
async def highlight(
    markers: Annotated[list[Marker], Field(min_length=1, max_length=20)],
    message: str | None = None,
    duration_seconds: Annotated[float, Field(ge=0, le=3600, description="0 = until cleared")] = 60,
    guide_lines: bool = True,
    chat: Annotated[bool, Field(description="Post [gps] links in chat")] = True,
) -> dict[str, Any]:
    """Outline buildings / ring areas in the user's game with labels and guide lines. Visual only."""
    return await game.request(
        "highlight",
        markers=[m.model_dump(exclude_none=True) for m in markers],
        message=message,
        duration_seconds=duration_seconds,
        guide_lines=guide_lines,
        chat=chat,
    )


@tool(annotations=VISUAL_ONLY)
async def clear_highlights() -> dict[str, Any]:
    """Remove all markers drawn by highlight."""
    return await game.request("clear_highlights")


@tool(annotations=ToolAnnotations(readOnlyHint=False, destructiveHint=False, openWorldHint=False))
async def run_lua(
    code: Annotated[str, Field(description="Expression, or statements with return")],
) -> dict[str, Any]:
    """READ-ONLY Lua query for what other tools can't answer. Globals: player, surface, force, game,
    prototypes, defines; print() is captured. Never modify game state."""
    return await game.request("run_lua", code=code)


def _compact_schema(node: Any) -> None:
    """Drops pydantic's auto titles and null-unions from advertised schemas. Arguments are still
    validated by the SDK's own model, which accepts null for optional parameters."""
    if isinstance(node, dict):
        node.pop("title", None)
        if "default" in node and node["default"] is None:
            del node["default"]
        if "anyOf" in node:
            options = [o for o in node["anyOf"] if o != {"type": "null"}]
            if len(options) == 1:
                del node["anyOf"]
                node.update(options[0])
            else:
                node["anyOf"] = options
        for value in node.values():
            _compact_schema(value)
    elif isinstance(node, list):
        for value in node:
            _compact_schema(value)


# Private SDK attribute: the registered tools whose `parameters` is what tools/list advertises.
for _tool in mcp._tool_manager._tools.values():
    _compact_schema(_tool.parameters)


def main() -> None:
    mcp.run()


if __name__ == "__main__":
    main()