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
import json

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

mcp = FastMCP("Student Text Adventure Server")


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

class GameManager:
    """Manages the text adventure game state and builds a dynamic map."""
    
    def __init__(self):
        self.env: TextAdventureEnv = None
        self.state = None
        self.game_name: str = ""
        
        # --- Map Tracking Variables ---
        self.visited = set()
        self.connections = {}  # Format: { "Room A": { "north": "Room B" } }
        self.unexplored = {}   # Format: { "Room A": set(["east", "west"]) }
        
        # Standard Z-machine directions
        self.directions = {
            'north', 'south', 'east', 'west', 
            'ne', 'nw', 'se', 'sw', 
            'up', 'down', 'in', 'out', 'enter', 'exit',
            'n', 's', 'e', 'w', 'u', 'd'
        }
        
        # Dictionary to normalize shortcuts (n -> north)
        self.dir_map = {
            'n': 'north', 's': 'south', 'e': 'east', 'w': 'west',
            'u': 'up', 'd': 'down'
        }


    def _normalize_dir(self, direction: str) -> str:
        d = direction.lower().strip()
        return self.dir_map.get(d, d)

    def _update_map(self, previous_loc: str, action: str, current_loc: str):
        """Builds the graph as the player moves around."""
        if not current_loc:
            # Still record the edge from previous room as a dead-end/blocked path
            if previous_loc and action:
                norm_action = self._normalize_dir(action)
                if norm_action in self.directions and previous_loc in self.unexplored:
                    self.unexplored[previous_loc].discard(norm_action)
            return

        # 1. Initialize the new room if we haven't seen it
        if current_loc not in self.connections:
            self.connections[current_loc] = {}
            self.visited.add(current_loc)
            
            # Find which directions are possible from this new room
            if self.env and self.env.env:
                valid_actions = self.env.env.get_valid_actions(use_parallel=False)
                # Filter out standard directions
                valid_dirs = set(self._normalize_dir(a) for a in valid_actions if a.lower() in self.directions)
                self.unexplored[current_loc] = valid_dirs

        # 2. Record the traversal edge if we just moved
        norm_action = self._normalize_dir(action)
        if previous_loc and previous_loc != current_loc and norm_action in self.directions:
            self.connections[previous_loc][norm_action] = current_loc
            
            # Remove this direction from unexplored for the previous room
            if previous_loc in self.unexplored:
                self.unexplored[previous_loc].discard(norm_action)

    def initialize(self, game: str = "zork1"):
        self.game_name = game
        self.env = TextAdventureEnv(game)
        self.state = self.env.reset()
        
        # Map the starting room
        start_loc = self.env.env.get_player_location()
        if start_loc:
            self._update_map(None, "", start_loc.name)
            
        return self.state.observation
    
    def step(self, action: str) -> str:
        if self.env is None:
            self.initialize()
            
        loc_before_obj = self.env.env.get_player_location()
        loc_before = loc_before_obj.name if loc_before_obj else None
        
        self.state = self.env.step(action)
        
        loc_after_obj = self.env.env.get_player_location()
        loc_after = loc_after_obj.name if loc_after_obj else None
        
        # Update our Map Graph!
        self._update_map(loc_before, action, loc_after)
        
        return self.state.observation
    
    def get_score(self) -> int: return self.state.score if self.state else 0
    def get_moves(self) -> int: 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."""
    game = get_game()
    
    # Get location BEFORE action
    loc_before = game.env.env.get_player_location().name if game.env.env.get_player_location() else ""
    
    # Execute action
    result = game.step(action)
    
    # Get location AFTER action
    loc_after = game.env.env.get_player_location().name if game.env.env.get_player_location() else ""
    
    # If the player moved, append a fresh "look" observation automatically
    # if loc_before != loc_after and action.lower() not in ['look', 'l']:
    #     look_result = game.step("look")
    #     result += f"\n\n[You moved to a new area]\n{look_result}"
        
    return result

# TODO: Implement additional tools to help your agent
@mcp.tool()
def game_state() -> str:
    """
    Returns the current state of the game: Score, Moves, Location, and Inventory.
    """
    game = get_game()
    if not game.env.env:
        return "Game not initialized."
    
    inventory = game.env.env.get_inventory()
    inv_names = [obj.name for obj in inventory] if inventory else ["Empty"]
    
    location = game.env.env.get_player_location()
    loc_name = location.name if location else "Unknown"

    state = {
        "location": loc_name,
        "score": game.get_score(),
        "moves": game.get_moves(),
        "inventory": inv_names
    }
    return json.dumps(state)


@mcp.tool()
def inventory() -> str:
    """
    Check what the player is carrying.
    
    Returns:
        List of items in the player's inventory
    """
    game = get_game()
    if not game.env.env: return "Game not initialized."
    
    inventory_objects = game.env.env.get_inventory()
    if not inventory_objects:
        return "Your inventory is empty."
    
    items = [obj.name for obj in inventory_objects]
    return f"Inventory: {', '.join(items)}"

@mcp.tool()
def memory() -> str:
    """
    Get the current game state summary: Location, Score, and Moves.
    Use this to orient yourself.
    """
    game = get_game()
    if not game.env.env: return "Game not initialized."
    
    location = game.env.env.get_player_location()
    loc_name = location.name if location else "Unknown Location"
    
    return json.dumps({
        "location": loc_name,
        "score": game.get_score(),
        "moves": game.get_moves(),
        "max_score": game.env.env.get_max_score()
    })

@mcp.tool()
def get_map() -> str:
    """
    Get a map of explored locations AND the paths to reach unexplored exits.
    Use this to figure out where to go next to discover new areas.
    """
    game = get_game()
    if not game.env or not game.env.env:
        return "Game not initialized."
        
    current_loc_obj = game.env.env.get_player_location()
    if not current_loc_obj:
        return "Cannot determine your current location."
        
    current_loc = current_loc_obj.name
    
    # 1. List Visited Locations
    visited_str = ", ".join(sorted(list(game.visited)))
    
    # 2. Find paths to Unexplored Exits using BFS (Breadth-First Search)
    queue = [(current_loc, [])]  # Queue of (Room Name, [Path of Actions])
    visited_bfs = set([current_loc])
    
    paths_to_unexplored = []
    
    # Rooms that have at least one unexplored exit
    rooms_with_unexplored = {r: exits for r, exits in game.unexplored.items() if len(exits) > 0}
    
    while queue:
        curr, path = queue.pop(0)
        
        # If this room has unexplored exits, log the path to get here!
        if curr in rooms_with_unexplored:
            unexp_str = ", ".join(rooms_with_unexplored[curr])
            if not path:
                paths_to_unexplored.append(f"Right here in '{curr}', you haven't tried: {unexp_str}")
            else:
                path_str = " -> ".join(path)
                paths_to_unexplored.append(f"To explore '{curr}' ({unexp_str}), walk: {path_str}")
                
        # Traverse known connections
        for direction, neighbor in game.connections.get(curr, {}).items():
            if neighbor not in visited_bfs:
                visited_bfs.add(neighbor)
                queue.append((neighbor, path + [direction]))
                
    # 3. Format the final beautiful output
    output = f"CURRENT LOCATION: {current_loc}\n\n"
    output += f"VISITED LOCATIONS ({len(game.visited)} total):\n{visited_str}\n\n"
    output += f"UNEXPLORED PATHS & HOW TO GET THERE:\n"

    if paths_to_unexplored:
        output += "\n".join(f"- {p}" for p in paths_to_unexplored)
    else:
        output += "- No known unexplored paths! You have fully explored everything."
        
    return output

@mcp.tool()
def get_valid_actions() -> str:
    """
    Get a list of guaranteed valid actions for the current game state.
    Use this when you are stuck or don't know what verbs the game understands.
    """
    game = get_game()
    if not game.env.env:
        return "Game not initialized."
    
    # Jericho extracts valid actions based on the object tree
    valid_actions = game.env.env.get_valid_actions(use_object_tree=True, use_parallel=False)
    
    # Filter out boring/meta actions if necessary, or just return them
    return json.dumps({
        "valid_actions": valid_actions[:30] # Limit to top 30 to save context window
    })

@mcp.tool()
def inspect_surroundings() -> str:
    """
    Scans the room and returns a list of interactive objects physically present.
    """
    game = get_game()
    if not game.env.env: return "Game not initialized."
    
    player_loc = game.env.env.get_player_location()
    if not player_loc: return "Cannot determine location."
    
    # Traverse the object tree: get the first child of the room, then iterate siblings
    objects_in_room = []
    child_num = player_loc.child
    
    while child_num != 0:
        obj = game.env.env.get_object(child_num)
        if obj and obj.num != game.env.env.get_player_object().num: # Don't list the player
            objects_in_room.append(obj.name)
        child_num = obj.sibling if obj else 0
        
    if not objects_in_room:
        return "No notable interactive objects found here."
        
    return f"Interactive objects in this room: {', '.join(objects_in_room)}"


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

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