mcp_server.py

"""
Student MCP Server for Text Adventure Games

This is your MCP server submission. Implement the tools that your agent
will use to play text adventure games.

Required tool:
    play_action(action: str) -> str
        Execute a game command and return the result.

Recommended tools:
    memory() -> str
        Return current game state, score, and recent history.
    
    inventory() -> str  
        Return the player's current inventory.
    
    get_map() -> str
        Return a map of explored locations.

Test your server with:
    fastmcp dev submission_template/mcp_server.py

Then open the MCP Inspector in your browser to test the tools interactively.
"""
import re
import sys
import os

# Add parent directory to path to import games module
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from fastmcp import FastMCP
from games.zork_env import TextAdventureEnv


# =============================================================================
# Create the MCP Server
# =============================================================================

mcp = FastMCP("Student Text Adventure Server")


# =============================================================================
# Game State Management
# =============================================================================

class GameManager:
    """
    Manages the text adventure game state.
    
    TODO: Extend this class to track:
    - Action history (for memory tool)
    - Explored locations (for mapping)
    - Current score and moves
    """
    
    def __init__(self):
        """Initialize game manager state."""
        self.env: TextAdventureEnv = None
        self.state = None
        self.game_name: str = ""

        self.current_location: str = "Unknown"
        self.prev_location: str = "Unknown"
        self.history: list[dict] = []
        self.location_visits: dict[str, int] = {}
        # (location, action) -> list of result summaries
        self.tried: dict[tuple[str, str], list[str]] = {}
        # location -> {direction: destination}
        self.connections: dict[str, dict[str, str]] = {}

    def initialize(self, game: str = "zork1"):
        """Start a new game."""
        self.game_name = game
        self.env = TextAdventureEnv(game)
        self.state = self.env.reset()
        self.history = []
        self.current_location = self._get_location()
        self.prev_location = self.current_location
        self.location_visits = {self.current_location: 1}
        self.tried = {}
        self.connections = {}
        return self.state.observation

    def _get_location(self) -> str:
        """Get clean location from Jericho internal state."""
        if self.state and getattr(self.state, "location", None):
            loc = str(self.state.location)
        else:
            try:
                loc_obj = self.env.env.get_player_location()
                loc = loc_obj.name if hasattr(loc_obj, "name") and loc_obj.name else str(loc_obj)
            except Exception:
                loc = "Unknown"
        # Clean Jericho noise
        loc = re.sub(r'Obj\d+:\s*', '', loc)
        loc = re.sub(r'\s*(Parent|Sibling|Child)\d+.*', '', loc)
        loc = re.sub(r'\s*Attributes\s*\[.*', '', loc)
        loc = re.sub(r'\[.*?\]', '', loc)
        return loc.strip() or "Unknown"

    _DIRECTIONS = {
        "north", "south", "east", "west", "up", "down",
        "northeast", "northwest", "southeast", "southwest",
        "enter", "exit", "in", "out", "n", "s", "e", "w",
        "ne", "nw", "se", "sw", "u", "d",
    }

    def step(self, action: str) -> dict:
        """Execute an action in the game and update state."""
        if self.env is None:
            self.initialize()

        old_location = self.current_location
        old_score = self.state.score if self.state else 0

        self.state = self.env.step(action)
        new_location = self._get_location()
        new_score = self.state.score
        score_change = new_score - old_score
        moved = (new_location != old_location and new_location != "Unknown")

        act_lower = action.lower().strip()
        is_dir = act_lower in self._DIRECTIONS

        if moved:
            self.prev_location = old_location
            self.current_location = new_location
            self.location_visits[new_location] = self.location_visits.get(new_location, 0) + 1
            if is_dir:
                if old_location not in self.connections:
                    self.connections[old_location] = {}
                self.connections[old_location][act_lower] = new_location

        # Record (location, action) -> result
        key = (old_location, act_lower)
        obs_summary = self.state.observation[:150].replace("\n", " ")
        if key not in self.tried:
            self.tried[key] = []
        self.tried[key].append(obs_summary)

        record = {
            "action": action,
            "observation": self.state.observation,
            "old_location": old_location,
            "new_location": self.current_location,
            "moved": moved,
            "score": new_score,
            "score_change": score_change,
            "moves": self.state.moves,
        }
        self.history.append(record)
        return record

    def get_score(self) -> int:
        """Get current score."""
        return self.state.score if self.state else 0

    def get_moves(self) -> int:
        """Get current move count."""
        return self.state.moves if self.state else 0

    def get_inventory(self) -> list[str]:
        """Get current inventory items."""
        if self.env and self.env.env:
            try:
                items = self.env.env.get_inventory()
                return [obj.name for obj in items]
            except Exception:
                return []
        return []

    def get_context(self) -> str:
        """Get a summary of the current game state, including location, score, inventory, known exits, and recent history."""
        parts = []
        loc = self.current_location

        parts.append(f"LOCATION: {loc}")
        parts.append(f"SCORE: {self.get_score()}")
        parts.append(f"MOVES: {self.get_moves()}")
        inv = self.get_inventory()
        parts.append(f"INVENTORY: {', '.join(inv) if inv else 'empty'}")

        # Known working exits
        exits_here = self.connections.get(loc, {})
        if exits_here:
            parts.append(f"KNOWN EXITS: {', '.join(f'{d}->{dest}' for d, dest in exits_here.items())}")

        # Actions tried here (non-direction)
        non_dir_tried = []
        for (place, act), results in self.tried.items():
            if place == loc and act not in self._DIRECTIONS and act != "look":
                non_dir_tried.append(f'"{act}" (x{len(results)})')
        if non_dir_tried:
            parts.append(f"OTHER ACTIONS TRIED HERE: {', '.join(non_dir_tried)}")

        # Directions tried here with their results
        dir_tried = []
        for (place, act), results in self.tried.items():
            if place == loc and act in self._DIRECTIONS:
                last_result = results[-1][:80]
                dir_tried.append(f'"{act}" x{len(results)} → {last_result}')
        if dir_tried:
            parts.append(f"DIRECTIONS TRIED HERE: {', '.join(dir_tried)}")

        # Recent history
        parts.append(f"\nRECENT HISTORY (last 10):")
        for rec in self.history[-10:]:
            tags = []
            if rec["score_change"] > 0:
                tags.append(f"+{rec['score_change']}pts!")
            if rec["score_change"] < 0:
                tags.append(f"{rec['score_change']}pts")
            if rec["moved"]:
                tags.append(f"moved->{rec['new_location']}")
            tag_str = f" [{', '.join(tags)}]" if tags else ""
            obs_short = rec["observation"][:100].replace("\n", " ")
            parts.append(f"  > {rec['action']}{tag_str} => {obs_short}")

        # Full map
        all_locations = set(self.location_visits.keys())
        if len(all_locations) > 1:
            parts.append(f"\nMAP ({len(all_locations)} locations):")
            for place in sorted(all_locations):
                marker = " *** YOU ARE HERE ***" if place == loc else ""
                conns = self.connections.get(place, {})
                conn_str = ", ".join(f"{d}->{dest}" for d, dest in conns.items()) if conns else "none found"
                parts.append(f"  {place}{marker}")
                parts.append(f"    exits: {conn_str}")

        return "\n".join(parts)


# Global game manager
_game = GameManager()


def get_game() -> GameManager:
    """Get or initialize the game manager."""
    global _game
    if _game.env is None:
        # Get game from environment variable (set by evaluator)
        game = os.environ.get("GAME", "zork1")
        _game.initialize(game)
    return _game


# =============================================================================
# MCP Tools - IMPLEMENT THESE
# =============================================================================

@mcp.tool()
def play_action(action: str) -> str:
    """Execute a game command and return the observation."""
    game = get_game()
    record = game.step(action)

    response = record["observation"]
    if record["score_change"] > 0:
        response += f"\n*** SCORED {record['score_change']} POINTS! ***"
    if record["score_change"] < 0:
        response += f"\n*** LOST {abs(record['score_change'])} POINTS ***"

    response += f"\n[Score:{record['score']} Moves:{record['moves']} Location:{record['new_location']}]"

    if record["moved"]:
        response += f"\n[Moved from {record['old_location']} to {record['new_location']}]"

    return response


@mcp.tool()
def memory() -> str:
    """Get full game context: location, score, inventory, tried actions, map."""
    return get_game().get_context()


@mcp.tool()
def get_score() -> str:
    """Return current game score."""
    return str(get_game().get_score())


@mcp.tool()
def inventory() -> str:
    """
    Check what the player is carrying.

    Returns:
        List of items in the player's inventory
    """
    game = get_game()
    items = game.get_inventory()
    if items:
        return f"You are carrying: {', '.join(items)}"
    else:
        return "You are not carrying anything."


# =============================================================================
# Run the server
# =============================================================================

if __name__ == "__main__":
    # This runs the server with stdio transport (for MCP clients)
    mcp.run()