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 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.
    
    Tracks:
    - Action history (for memory tool)
    - Explored locations and connections (for mapping)
    - Actions taken at each location
    - Current score and moves
    """
    
    def __init__(self):
        self.env: TextAdventureEnv = None
        self.state = None
        self.game_name: str = ""
        self.history: list[tuple[str, str]] = []  # (action, observation) pairs
        self.explored_rooms: dict[int, dict[str, int]] = {}  # room_id -> {direction: target_room_id}
        self.room_actions: dict[int, list[str]] = {}  # room_id -> actions taken there
        self.room_names: dict[int, str] = {}  # room_id -> room name (for display)
        self.current_room_id: int = -1
        self.previous_room_id: int = -1
    
    def _get_current_room_info(self) -> tuple[int, str]:
        """Get current room ID and name from Jericho."""
        if self.env and self.env.env:
            try:
                loc = self.env.env.get_player_location()
                room_id = loc.num
                room_name = loc.name
                return room_id, room_name
            except:
                pass
        return -1, "Unknown"
    
    def initialize(self, game: str = "zork1"):
        """Initialize or reset the game."""
        self.game_name = game
        self.env = TextAdventureEnv(game)
        self.state = self.env.reset()
        self.history = []
        self.explored_rooms = {}
        self.room_actions = {}
        self.room_names = {}
        
        # Get initial room
        room_id, room_name = self._get_current_room_info()
        self.current_room_id = room_id
        self.previous_room_id = room_id
        self.room_names[room_id] = room_name
        
        return self.state.observation
    
    def step(self, action: str) -> str:
        """Execute an action and return the result."""
        if self.env is None:
            self.initialize()
        
        # Track action at current room
        if self.current_room_id not in self.room_actions:
            self.room_actions[self.current_room_id] = []
        self.room_actions[self.current_room_id].append(action)
        
        # Execute action
        self.state = self.env.step(action)
        result = self.state.observation
        
        # Track history (keep last 50)
        self.history.append((action, result))
        if len(self.history) > 50:
            self.history = self.history[-50:]
        
        # Get new room info
        self.previous_room_id = self.current_room_id
        new_room_id, new_room_name = self._get_current_room_info()
        
        # Track movement in spatial graph
        direction_actions = ["north", "south", "east", "west", "up", "down",
                            "enter", "exit", "n", "s", "e", "w", "u", "d",
                            "ne", "nw", "se", "sw", "northeast", "northwest",
                            "southeast", "southwest"]
        if action.lower() in direction_actions and new_room_id != self.current_room_id:
            # Build spatial graph
            if self.current_room_id not in self.explored_rooms:
                self.explored_rooms[self.current_room_id] = {}
            self.explored_rooms[self.current_room_id][action.lower()] = new_room_id
        
        # Update current room
        self.current_room_id = new_room_id
        if new_room_id not in self.room_names:
            self.room_names[new_room_id] = new_room_name
        
        return result
    
    def get_score(self) -> int:
        """Get current score."""
        return self.state.score if self.state else 0
    
    def get_moves(self) -> int:
        """Get number of moves taken."""
        return self.state.moves if self.state else 0


# 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 result.
    
    This is the main tool for interacting with the game.
    
    Args:
        action: The command to execute (e.g., "north", "take lamp", "open mailbox")
        
    Returns:
        The game's response to the action with room ID and score info
        
    Valid commands include:
        - Movement: north, south, east, west, up, down, enter, exit
        - Objects: take <item>, drop <item>, open <thing>, examine <thing>
        - Other: look, inventory, read <thing>, turn on lamp
    """
    game = get_game()
    result = game.step(action)
    
    # Add room ID and name
    room_name = game.room_names.get(game.current_room_id, "Unknown")
    location_info = f"\n\n[Room #{game.current_room_id}: {room_name}]"
    
    # Add score and move info
    score_info = f"\n[Score: {game.get_score()} | Moves: {game.get_moves()}]"
    
    # Highlight score gains
    if game.state.reward > 0:
        score_info = f"\n[+{game.state.reward} points! Total: {game.get_score()} | Moves: {game.get_moves()}]"
    
    # Game over detection
    done_info = ""
    if game.state.done:
        done_info = "\n\nGAME OVER"
    
    return result + location_info + score_info + done_info


# TODO: Implement additional tools to help your agent

@mcp.tool()
def get_location_id() -> str:
    """
    Get the current room ID and name.
    
    Returns:
        Room ID (unique identifier) and room name
    """
    game = get_game()
    room_name = game.room_names.get(game.current_room_id, "Unknown")
    return f"Room #{game.current_room_id}: {room_name}"


@mcp.tool()
def memory() -> str:
    """
    Get the current game state summary.
    
    Returns:
        A summary including current location, score, moves, and recent history
    """
    game = get_game()
    
    # Get recent history (last 5 actions)
    recent = game.history[-5:] if game.history else []
    recent_str = "\n".join([f"  > {a} -> {r[:60]}..." for a, r in recent]) if recent else "  (none yet)"
    
    room_name = game.room_names.get(game.current_room_id, "Unknown")
    
    return f"""Current State:
- Room: #{game.current_room_id} ({room_name})
- Score: {game.get_score()} points
- Moves: {game.get_moves()}
- Game: {game.game_name}

Recent Actions:
{recent_str}

Current Observation:
{game.state.observation if game.state else 'No observation yet'}"""


@mcp.tool()
def inventory() -> str:
    """
    Check what the player is carrying.
    
    Returns:
        List of items in the player's inventory
    """
    game = get_game()
    result = game.step("inventory")
    return result


@mcp.tool()
def get_map() -> str:
    """
    Get a map of explored rooms and connections.
    
    Returns:
        A spatial graph of explored rooms and their connections
    """
    game = get_game()
    
    if not game.explored_rooms:
        return "Map: No rooms explored yet. Try moving around!"
    
    lines = ["Spatial Map (Room ID-based):"]
    
    # Show all known rooms
    all_rooms = set(game.explored_rooms.keys()) | {game.current_room_id}
    for room_id in sorted(all_rooms):
        room_name = game.room_names.get(room_id, "Unknown")
        is_current = " (CURRENT)" if room_id == game.current_room_id else ""
        lines.append(f"\n* Room #{room_id}: {room_name}{is_current}")
        
        # Show exits from this room
        if room_id in game.explored_rooms:
            for direction, target_id in sorted(game.explored_rooms[room_id].items()):
                target_name = game.room_names.get(target_id, "?")
                lines.append(f"    {direction} -> Room #{target_id} ({target_name})")
    
    lines.append(f"\n[Total Rooms Discovered] {len(game.room_names)}")
    lines.append(f"[Connections Mapped] {sum(len(exits) for exits in game.explored_rooms.values())}")
    
    return "\n".join(lines)


@mcp.tool()
def get_valid_actions() -> str:
    """
    Get a list of likely valid actions from the current location.
    
    Returns:
        List of actions that might work here
    """
    # This is a hint: Jericho provides get_valid_actions()
    game = get_game()
    if game.env and game.env.env:
        valid = game.env.env.get_valid_actions()
        return "Valid actions: " + ", ".join(valid[:20])
    return "Could not determine valid actions"


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

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