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

# 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

# Get game from environment variable (default: zork1)
INITIAL_GAME = os.environ.get("GAME", "lostpig")


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

mcp = FastMCP("Student Text Adventure Server")


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

class LocationLog:
    """Tracks actions, outcomes, and promising leads for a single location."""
    def __init__(self, name: str):
        self.name = name                                # Location name (e.g., "Kitchen")
        self.visit_count: int = 0                       # How many times we've been here
        self.actions_taken: list[tuple[str, str]] = []  # (action, short_outcome)
        self.exits_known: list[str] = []                # List of known exits from this location (e.g., "north -> Kitchen")
        self.promising_actions: list[str] = []          # Actions worth revisiting at this location


class GameManager:
    """
    Manages the text adventure game state with rich location tracking.
    """
    
    def __init__(self):
        self.env: TextAdventureEnv = None
        self.state = None
        self.game_name: str = ""
        self.history: list[tuple[str, str]] = []            # list of (action, result) for recent actions
        self.explored_locations: dict[str, set[str]] = {}   # location name -> set of exits (e.g., "north -> Kitchen")
        self.location_logs: dict[str, LocationLog] = {}     # location name -> log of actions and outcomes at that location
        self.previous_player_location: str = ""             # Jericho internal location object
        self.current_player_location: str = ""              # Jericho internal location object
        self.global_action_count: int = 0
        self.score_history: list[int] = []

    def _get_jericho_location(self):
        """Get the internal Jericho player location object for comparison."""
        try:
            res = self.env.env.get_player_location()
            match = re.search(r"Obj\d+: (.*) Parent\d+", res.name)
            if match:
                return match.group(1)
            # Fallback: return the full location string
            return res.name
        except Exception:
            return None

    def has_moved(self) -> bool:
        """
        Determine if we moved to a new location.
        Compares current player location object to the previous one.
        """
        current_loc = self.current_player_location
        previous_loc = self.previous_player_location

        changed = not (current_loc == previous_loc)
        return changed

    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_locations = {}
        self.location_logs = {}
        self.previous_player_location = ""
        self.current_player_location = self._get_jericho_location()
        self.global_action_count = 0
        self.score_history = [0]
        self._ensure_location_log(self.current_player_location)
        self.location_logs[self.current_player_location].visit_count += 1
        return self.state.observation
    
    def _ensure_location_log(self, location: str):
        """Ensure a LocationLog exists for the given location."""
        if location not in self.location_logs:
            self.location_logs[location] = LocationLog(location)

    def take_action(self, action: str) -> str:
        """Execute an action and return the result."""
        if self.env is None:
            self.initialize()
        
        self.previous_player_location = self._get_jericho_location()    # Store previous location before taking action
        self.state = self.env.step(action)                              # Execute the action in the game environment
        self.current_player_location = self._get_jericho_location()     # Store current location after taking action

        result = self.state.observation                                 # Get the observation/result of the action
        self.global_action_count += 1
        self.score_history.append(self.state.score)
        
        # Track history
        self.history.append((action, result))
        if len(self.history) > 50:
            self.history = self.history[-50:]

        moved = self.has_moved()
        is_new_place = False
            
        # New place! Update explored locations map
        if self.current_player_location not in self.explored_locations:
            self.explored_locations[self.current_player_location] = set()
            is_new_place = True
        # Add exit from previous location to current location
        if moved:
            self.explored_locations[self.previous_player_location].add(f"{action} -> {self.current_player_location}")
        
        # Update location log
        self._ensure_location_log(self.current_player_location)
        current_loc_log = self.location_logs[self.current_player_location]
        if moved:
            current_loc_log.visit_count += 1
        
        # Log this action and a short outcome in the previous location's log
        prev_loc_log = self.location_logs.get(self.previous_player_location)
        if prev_loc_log is not None:
            short_outcome = result[:120].replace('\n', ' ')
            prev_loc_log.actions_taken.append((action, short_outcome))
            # Keep log manageable
            if len(prev_loc_log.actions_taken) > 30:
                prev_loc_log.actions_taken = prev_loc_log.actions_taken[-30:]
        
        return result, is_new_place

    def get_memory(self) -> str:
        """Get a summary of current game state."""
        recent = self.history[-5:] if self.history else []
        recent_str = "\n".join([f"  > {a} -> {r[:60]}..." for a, r in recent]) if recent else "  (none yet)"
        
        # Add location-specific info
        loc_log = self.location_logs.get(self.current_player_location)
        loc_info = ""
        if loc_log:
            loc_info = f"\nThis location visited {loc_log.visit_count} time(s)."
            if loc_log.actions_taken:
                loc_info += f"\nActions tried at this location: {len(loc_log.actions_taken)}"
                recent_here = loc_log.actions_taken[-5:]
                loc_info += "\nRecent actions at this location:"
                for act, out in recent_here:
                    loc_info += f"\n  > {act} -> {out[:50]}..."
            if loc_log.promising_actions:
                loc_info += f"\nPromising actions at this location: {', '.join(loc_log.promising_actions[:10])}"

        return f"""[CURRENT LOCATION: {self.current_player_location}]
Location info:
{loc_info}

Recent Actions:
{recent_str}

Current Observation:
{self.state.observation}"""
    

    def get_map(self) -> str:
        """Get a map of explored locations."""
        if not self.explored_locations:
            return "Map: No locations explored yet. Try moving around!"
        lines = [f"[CURRENT LOCATION: {self.current_player_location}]"]
        lines.append("EXPLORED LOCATIONS AND EXITS:")
        for loc, exits in sorted(self.explored_locations.items()):
            visit_info = ""
            if loc in self.location_logs:
                visit_info = f" (visited {self.location_logs[loc].visit_count}x, {len(self.location_logs[loc].actions_taken)} actions tried)"
            lines.append(f"\n* {loc}{visit_info}")
            if exits:
                for exit_info in sorted(exits):
                    lines.append(f"    -> {exit_info}")
            else:
                lines.append("    -> No exits mapped yet")

        # Add detailed log for current location
        location = self.current_player_location
        loc_log = self.location_logs.get(location)
        if loc_log:
            lines.append(f"\n- INFORMATION FOR CURRENT LOCATION: {location}")
            lines.append(f"   * Visited: {loc_log.visit_count} time(s)")
            lines.append(f"   * Actions tried here: {len(loc_log.actions_taken)}")
            
            if loc_log.actions_taken:
                lines.append(f"   * Action history at this location {location}:")
                for act, out in loc_log.actions_taken[-10:]:
                    lines.append(f"      > {act} -> {out[:80]}")

        return "\n".join(lines)

    def get_inventory(self) -> str:
        """Get current inventory."""
        if self.env is None:
            return "Game not initialized"
        inv_state = self.env.step("inventory")
        lines = [f"[CURRENT LOCATION: {self.current_player_location}]"]
        lines.append(inv_state.observation)
        return "\n".join(lines)
    
    def get_location_log(self) -> str:
        """Get detailed log for a specific location."""
        location = self.current_player_location
        loc_log = self.location_logs.get(location)
        if not loc_log:
            return f"No log for location: {location}"
        
        lines = [f"[CURRENT LOCATION: {location}]"]
        lines.append(f"Visited: {loc_log.visit_count} time(s)")
        lines.append(f"Actions tried: {len(loc_log.actions_taken)}")
        
        if loc_log.actions_taken:
            lines.append("\nAction history:")
            for act, out in loc_log.actions_taken[-10:]:
                lines.append(f"  > {act} -> {out[:80]}")
        
        if loc_log.exits_known:
            lines.append(f"Known exits: {', '.join(loc_log.exits_known)}")
        
        return "\n".join(lines)


# Global game manager
_game = GameManager()


def get_game() -> GameManager:
    """Get or initialize the game manager."""
    global _game
    if _game.env is None:
        game = os.environ.get("GAME", "lostpig")
        _game.initialize(game)
    return _game


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

@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
        
    Valid commands include:
        - Movement: north, south, east, west, northeast, northwest, southeast, southwest, 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, is_new_place = game.take_action(action)
    
    # Add score info
    score_info = f"\n\n[Score: {game.state.score} | Moves: {game.state.moves}]"
    
    if game.state.reward > 0:
        score_info = f"\n\n+{game.state.reward} points! (Total: {game.state.score})"
    
    # Indicate if we moved to a new location
    location_info = ""
    if is_new_place:
        location_info = f"[NEW LOCATION: {game.current_player_location}]\n"
    else:
        location_info = f"[CURRENT LOCATION: {game.current_player_location}]\n"
    
    done_info = ""
    if game.state.done:
        done_info = "\n\nGAME OVER"
    
    return location_info + result + score_info + done_info


@mcp.tool()
def memory() -> str:
    """
    Get the current game state summary.
    
    Returns:
        A summary including current location (number of visits, actions tried, promising actions), 
        recent actions and current observation
    """
    return get_game().get_memory()


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


@mcp.tool()
def get_map() -> str:
    """
    Get a map of explored locations, connections and exits.
    Useful for navigation and avoiding getting lost.
    
    Returns:
        A text representation of explored locations and connections
    """
    return get_game().get_map()


# @mcp.tool()
# def get_valid_actions() -> str:
#     """
#     Get a list of valid actions from the current game state using the game engine.
#     Useful when entering a new location to understand what's possible.
    
#     Returns:
#         A list of valid actions that the game engine considers possible
#     """
#     game = get_game()
#     valid = game.get_valid_actions_list()
#     if valid:
#         return "Valid actions: " + ", ".join(valid)
#     return "Could not determine valid actions. Try: look, inventory, examine, north, south, east, west, up, down, take, drop, open, close, read"


@mcp.tool()
def location_log() -> str:
    """
    Shows what actions were tried and their outcomes at the current location, along with any promising actions to try.
    
    
    Returns:
        A detailed log of the current location, including visit count, actions taken and their outcomes, and promising leads.
    """
    game = get_game()
    return game.get_location_log()



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

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