mcp_server.py

"""
MCP Server for Text Adventure Games

Exposes a text adventure environment (via Jericho) as a set of MCP tools.
Tools available to the agent:
  - execute       : run a game command and get the response
  - snapshot      : get structured JSON game state (observation, score, inventory, valid commands…)
  - session_log   : human-readable summary of recent history
  - world_map     : explored locations and their connections
  - carried_items : list of items currently in the player's possession
"""

import sys
import os
import json

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from fastmcp import FastMCP
from games.zork_env import TextAdventureEnv

INITIAL_GAME = os.environ.get("GAME", "zork1")

mcp = FastMCP("Text Adventure MCP Server")


# =============================================================================
# GameState — wraps the Jericho environment and all derived state
# =============================================================================

class GameState:
    """
    Wraps a single Jericho game session.

    Tracks the raw environment state, a rolling command history, and a
    lightweight world graph (room -> set of "direction -> room" strings)
    built from every successful navigation command.
    """

    def __init__(self, game: str = "zork1") -> None:
        self.game_name = game
        self.env = TextAdventureEnv(game)
        self.env_state = self.env.reset()

        self.cmd_history: list[tuple[str, str]] = []    # (command, response)
        self.world_graph: dict[str, set[str]] = {}      # room -> {"dir -> room", ...}
        self.current_room: str = self._parse_room_header(self.env_state.observation)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _parse_room_header(self, observation: str) -> str:
        """Return the first meaningful line of an observation as the room name."""
        for line in observation.strip().splitlines():
            stripped = line.strip()
            if stripped and not stripped.startswith("[") and not stripped.startswith("+"):
                return stripped
        return "Unknown"

    _NAV_COMMANDS = frozenset([
        "north", "south", "east", "west", "up", "down",
        "enter", "exit", "n", "s", "e", "w", "u", "d",
    ])

    def _is_navigation(self, command: str) -> bool:
        return command.lower().strip() in self._NAV_COMMANDS

    # ------------------------------------------------------------------
    # Core action execution
    # ------------------------------------------------------------------

    def execute_command(self, command: str) -> str:
        """Step the environment with *command* and update derived state."""
        self.env_state = self.env.step(command)
        response = self.env_state.observation

        # Rolling history (capped at 60 entries)
        self.cmd_history.append((command, response))
        if len(self.cmd_history) > 60:
            self.cmd_history = self.cmd_history[-60:]

        # Update world graph on successful navigation
        new_room = self.env_state.location or self._parse_room_header(response)
        if self._is_navigation(command) and new_room != self.current_room:
            self.world_graph.setdefault(self.current_room, set()).add(
                f"{command.lower().strip()} -> {new_room}"
            )
        self.current_room = new_room

        return response

    # ------------------------------------------------------------------
    # Query methods (read-only)
    # ------------------------------------------------------------------

    def query_valid_commands(self) -> list[str]:
        """Ask Jericho for the list of currently valid commands."""
        try:
            return self.env.get_valid_actions()
        except Exception:
            return []

    def snapshot(self) -> dict:
        """Return a structured dict with everything the agent needs."""
        return {
            "observation":    self.env_state.observation,
            "location":       self.env_state.location or self.current_room,
            "score":          self.env_state.score,
            "max_score":      self.env_state.max_score,
            "moves":          self.env_state.moves,
            "done":           self.env_state.done,
            "reward":         self.env_state.reward,
            "inventory":      self.env_state.inventory or [],
            "valid_commands": self.query_valid_commands(),
        }

    def session_log(self) -> str:
        """Human-readable summary of the current session."""
        recent = self.cmd_history[-5:]
        recent_lines = (
            "\n".join(f"  > {cmd}  ->  {resp[:60]}..." for cmd, resp in recent)
            if recent else "  (no commands yet)"
        )
        s = self.env_state
        return (
            f"Room    : {self.current_room}\n"
            f"Score   : {s.score} / {s.max_score}\n"
            f"Moves   : {s.moves}\n"
            f"Game    : {self.game_name}\n"
            f"\nRecent commands:\n{recent_lines}\n"
            f"\nCurrent observation:\n{s.observation}"
        )

    def render_world_map(self) -> str:
        """Render the explored world graph as indented text."""
        if not self.world_graph:
            return "World map: no locations explored yet — move around to build the map."

        lines = ["Explored world:"]
        for room in sorted(self.world_graph):
            lines.append(f"\n  [{room}]")
            for connection in sorted(self.world_graph[room]):
                lines.append(f"      {connection}")
        lines.append(f"\n  [You are here] {self.current_room}")
        return "\n".join(lines)

    def list_carried_items(self) -> str:
        """Return a readable inventory string."""
        raw_items = (
            self.env_state.inventory
            if hasattr(self.env_state, "inventory") and self.env_state.inventory
            else []
        )
        if not raw_items:
            return "Carrying: nothing."

        names: list[str] = []
        for item in raw_items:
            item_str = str(item)
            lower = item_str.lower()
            if "parent" in lower:
                chunk = item_str[: lower.index("parent")].strip()
                names.append(chunk.split(":", 1)[1].strip() if ":" in chunk else chunk)
            elif ":" in item_str:
                names.append(item_str.split(":", 1)[1].strip())
            else:
                names.append(item_str)

        return "Carrying: " + ", ".join(names)


# =============================================================================
# Singleton session
# =============================================================================

_session: GameState | None = None


def get_session() -> GameState:
    """Return the singleton GameState, initialising it on first call."""
    global _session
    if _session is None:
        _session = GameState(INITIAL_GAME)
    return _session


# =============================================================================
# MCP Tools
# =============================================================================

@mcp.tool()
def execute(command: str) -> str:
    """
    Run a game command and return the game's response.

    Args:
        command: A valid game command, e.g. 'north', 'take lamp', 'open mailbox'.

    Returns:
        The game's text response, followed by score / move count.
        Appends 'GAME OVER' when the session ends.
    """
    sess = get_session()
    response = sess.execute_command(command)

    s = sess.env_state
    score_line = (
        f"\n\n+{s.reward} pts  (total: {s.score})"
        if s.reward > 0
        else f"\n\n[Score: {s.score} | Moves: {s.moves}]"
    )
    end_line = "\n\nGAME OVER" if s.done else ""
    return response + score_line + end_line


@mcp.tool()
def snapshot() -> str:
    """
    Return the full structured game state as a JSON string.

    Fields: observation, location, score, max_score, moves, done,
            reward, inventory, valid_commands.

    Prefer this over 'session_log' when you need machine-readable data.
    """
    return json.dumps(get_session().snapshot())


@mcp.tool()
def session_log() -> str:
    """
    Return a human-readable summary of the current session.

    Includes current room, score, move count, recent command history,
    and the current observation.
    """
    return get_session().session_log()


@mcp.tool()
def world_map() -> str:
    """
    Return a text rendering of all explored rooms and their connections.

    Use this to plan routes and avoid revisiting dead ends.
    """
    return get_session().render_world_map()


@mcp.tool()
def carried_items() -> str:
    """
    List the items you are currently carrying.
    """
    return get_session().list_carried_items()


# =============================================================================
# Entry point
# =============================================================================

if __name__ == "__main__":
    mcp.run()