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
# =============================================================================

import re
from typing import Optional

class GameManager:
    """
    Manages the text adventure game state.

    Extended tracking:
    - Action history (for memory tool)
    - Explored locations (for mapping)
    - Current score and moves
    - Current location (best-effort, robust across games)
    """

    # Lines that are often NOT room titles across many IF games
    _HEADER_LIKE_PATTERNS = [
        r"^\s*score\s*[:=]\s*\d+",
        r"^\s*moves?\s*[:=]\s*\d+",
        r"^\s*turns?\s*[:=]\s*\d+",
        r"^\s*time\s*[:=]\s*",
        r"^\s*health\s*[:=]\s*\d+",
        r"^\s*location\s*[:=]\s*",
        r"^\s*\[.*\]\s*$",            # bracket-only status lines
        r"^\s*\(.*\)\s*$",            # parenthetical-only lines
        r"^\s*you\s+(are|see|can)\b",  # narrative sentence starters
    ]
    # Movement commands we consider for mapping (Zork-style + abbreviations)
    _MOVE_CMDS = {
        "north", "south", "east", "west", "up", "down", "enter", "exit",
        "n", "s", "e", "w", "u", "d"
    }

    # Common failure phrases when trying to move (best-effort, not perfect)
    _MOVE_FAIL_PHRASES = [
        "you can't go", "you cannot go", "can't go that way", "cannot go that way",
        "you can't go that way", "you cannot go that way",
        "you can't", "you cannot",
        "there is no way", "you can't see any way", "you see no way",
        "blocked", "closed", "won't open", "is locked", "locked",
        "too dark", "pitch black"
    ]

    def _is_movement_action(self, action: str) -> bool:
        """Return True if this action is a movement command we track."""
        a = (action or "").strip().lower()
        return a in self._MOVE_CMDS

    def _move_likely_succeeded(self, old_loc: str, new_loc: str, observation: str) -> bool:
        """
        Decide whether a move likely succeeded.
        Strong signal: location label changed.
        Negative signal: failure phrases in observation.
        """
        if new_loc and old_loc and new_loc != old_loc:
            return True

        text = (observation or "").lower()
        if any(phrase in text for phrase in self._MOVE_FAIL_PHRASES):
            return False

        # If location didn't change and no clear failure phrase, treat as "not sure" → don't add edge
        return False

    def _update_map(self, action: str, old_loc: str, new_loc: str) -> None:
        """Record a directed edge old_loc --action--> new_loc in explored_locations."""
        if not old_loc or not new_loc:
            return
        self.explored_locations.setdefault(old_loc, set()).add(f"{action} -> {new_loc}")


    def __init__(self):
        self.env: TextAdventureEnv = None
        self.state = None
        self.game_name: str = ""

        # Tracking for agent-support tools
        self.history: list[tuple[str, str]] = []
        self.explored_locations: dict[str, set[str]] = {}
        self.current_location: str = "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()

        # Reset tracking
        self.history = []
        self.explored_locations = {}
        self.current_location = self._extract_location(self.state.observation, fallback="Unknown")

        return self.state.observation

    def _extract_location(self, observation: str, fallback: Optional[str] = None) -> str:
        """
        Best-effort location extraction from the observation text.

        Strategy:
        1) Split into lines, skip empties
        2) Skip lines that look like status bars / headers / pure brackets
        3) Prefer a short, title-like line (room name)
        4) If nothing confident, return fallback (usually previous location)
        """
        if not observation:
            return fallback or "Unknown"

        lines = [ln.strip() for ln in observation.splitlines() if ln.strip()]
        if not lines:
            return fallback or "Unknown"

        header_res = [re.compile(pat, re.IGNORECASE) for pat in self._HEADER_LIKE_PATTERNS]

        def looks_like_header(line: str) -> bool:
            return any(rx.search(line) for rx in header_res)

        def looks_like_title(line: str) -> bool:
            # Many room titles are short and not ending with punctuation.
            if len(line) > 60:
                return False
            if line.endswith((".", "!", "?", ";", ":")):
                return False
            # Too many digits usually means a status line.
            if sum(ch.isdigit() for ch in line) >= 3:
                return False
            return True

        # First pass: first "title-like" line that isn't header-like
        for line in lines[:8]:  # only inspect top chunk; titles are usually early
            if looks_like_header(line):
                continue
            if looks_like_title(line):
                return line

        # Second pass: first non-header line
        for line in lines[:8]:
            if not looks_like_header(line):
                return line

        return fallback or "Unknown"

    def step(self, action: str) -> str:
        """Execute an action and return the result."""
        if self.env is None:
            self.initialize()

        # Save old location before action
        old_location = self.current_location

        # Apply action to the real game
        self.state = self.env.step(action)
        obs = self.state.observation

        # Track history (keep last 50)
        self.history.append((action, obs))
        if len(self.history) > 50:
            self.history = self.history[-50:]

        # Extract new location (fallback to old)
        new_location = self._extract_location(obs, fallback=old_location)

        # Update map only if it was a movement attempt AND it likely succeeded
        action_norm = (action or "").strip().lower()
        if self._is_movement_action(action_norm) and self._move_likely_succeeded(old_location, new_location, obs):
            self._update_map(action_norm, old_location, new_location)

        # Finally update current location
        self.current_location = new_location

        return obs


    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
    def _extract_facts(self, observation: str) -> dict:
        """
        Best-effort extraction of useful 'facts' from the current observation text.
        This is intentionally heuristic so it can work across many games.
        """
        obs = observation or ""
        text = obs.strip()
        lower = text.lower()

        # --- Exits mentioned (simple direction scan) ---
        directions = ["north", "south", "east", "west", "up", "down", "in", "out"]
        exits_found = []
        for d in directions:
            # We detect directions as whole words to reduce false matches
            if re.search(rf"\b{re.escape(d)}\b", lower):
                exits_found.append(d)
        exits_found = sorted(set(exits_found))

        # --- Visible things (very light heuristics) ---
        # We look for common IF patterns like "You see ... here." / "There is ... here."
        visible_candidates: list[str] = []

        patterns = [
            r"you see (.+?) here\.",
            r"you can see (.+?) here\.",
            r"there is (.+?) here\.",
            r"there are (.+?) here\.",
            r"you notice (.+?)\.",
        ]
        for pat in patterns:
            for m in re.finditer(pat, lower):
                chunk = m.group(1).strip()
                if chunk:
                    visible_candidates.append(chunk)

        # Clean visible candidates a bit (split simple lists, avoid huge strings)
        visible = []
        for chunk in visible_candidates:
            # Split on commas and "and" to get smaller pieces
            parts = re.split(r",|\band\b", chunk)
            for p in parts:
                item = p.strip(" .;:!?\t")
                if 1 <= len(item) <= 40:
                    visible.append(item)

        # Deduplicate and limit (so memory stays compact)
        visible = sorted(set(visible))[:10]

        return {
            "exits_mentioned": exits_found,
            "visible": visible,
        }

    def get_memory(self) -> str:
        """
        LLM-friendly summary of current game state.
        Format: Facts first, then recent actions, then the raw observation.
        """
        game = self.game_name or "Unknown"
        location = self.current_location or "Unknown"
        score = self.get_score()
        moves = self.get_moves()

        # Recent actions (keep short and anti-loop)
        recent = self.history[-5:] if self.history else []
        if recent:
            recent_lines = []
            for a, r in recent:
                snippet = (r or "").replace("\n", " ").strip()
                if len(snippet) > 80:
                    snippet = snippet[:80] + "..."
                recent_lines.append(f"- {a} -> {snippet}")
            recent_str = "\n".join(recent_lines)
        else:
            recent_str = "(none yet)"

        # Facts extracted from current observation
        obs = self.state.observation if self.state else ""
        facts = self._extract_facts(obs)

        exits_txt = ", ".join(facts["exits_mentioned"]) if facts["exits_mentioned"] else "(none detected)"
        visible_txt = ", ".join(facts["visible"]) if facts["visible"] else "(none detected)"

        return (
            "STATE\n"
            f"Game: {game}\n"
            f"Location: {location}\n"
            f"Score: {score}   Moves: {moves}\n"
            f"Visible (best effort): {visible_txt}\n"
            f"Exits mentioned (best effort): {exits_txt}\n"
            "\n"
            "RECENT\n"
            f"{recent_str}\n"
            "\n"
            "OBSERVATION\n"
            f"{obs}"
        )
    def get_map(self) -> str:
        """
        Return a readable map of explored locations.
        Uses explored_locations built during movement actions.

        Output is stable + compact for LLM use.
        """
        if not self.explored_locations:
            return "MAP\n(no locations recorded yet — try moving with north/south/east/west/etc.)"

        lines = ["MAP", "Explored locations and exits:"]
        for loc in sorted(self.explored_locations.keys()):
            exits = sorted(self.explored_locations[loc])
            lines.append(f"\n* {loc}")
            for e in exits:
                lines.append(f"  - {e}")

        lines.append(f"\n[Current] {self.current_location}")
        return "\n".join(lines)
    def get_inventory(self) -> str:
        """
        Return inventory in a robust way across different games/envs.

        Strategy:
        1) If state.inventory exists and is non-empty -> format it
        2) Otherwise, fall back to issuing the command "inventory"
           through the environment and return that observation
        """
        # 1) Try structured inventory if provided by env
        items = []
        if self.state is not None and hasattr(self.state, "inventory"):
            inv = getattr(self.state, "inventory")
            if inv:
                # Normalize to strings
                try:
                    items = [str(x).strip() for x in inv if str(x).strip()]
                except Exception:
                    items = []

        if items:
            # Keep it simple and safe: just join a cleaned list
            # (Avoid overly aggressive parsing that breaks across games)
            items = sorted(set(items))
            return "INVENTORY\n" + ", ".join(items)

        # 2) Fallback: ask the game directly (does NOT change inventory, just prints it)
        # NOTE: We do not want to record this as agent history/map; this is a server-side query.
        if self.env is None:
            self.initialize()

        try:
            tmp_state = self.env.step("inventory")
            inv_text = tmp_state.observation if tmp_state else "Inventory: (no response)"
        except Exception:
            inv_text = "Inventory: (unable to retrieve)"

        return "INVENTORY\n" + inv_text.strip()


# 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
        
    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()
    
    # TODO: You might want to add action validation here
    # TODO: You might want to include score changes in the response
    
    result = game.step(action)

    # Append score/moves for clearer feedback (LLM-friendly, low noise)
    result += f"\n[Score: {game.get_score()} | Moves: {game.get_moves()}]"
    return result
    
    # Optional: Append score info
    # result += f"\n[Score: {game.get_score()} | Moves: {game.get_moves()}]"
    

@mcp.tool()
def memory() -> str:
    """
    Return an LLM-friendly summary of the current game state.
    """
    game = get_game()
    return game.get_memory()
@mcp.tool()
def get_map() -> str:
    """
    Return a map of explored locations and recorded exits.
    """
    game = get_game()
    return game.get_map()

@mcp.tool()
def inventory() -> str:
    """
    Return the player's inventory in a robust way.
    """
    game = get_game()
    return game.get_inventory()


# TODO: Implement additional tools to help your agent

# @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()
#     # TODO: Return useful state information
#     pass


# @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 locations.
#     
#     Returns:
#         A text representation of explored locations and connections
#     """
#     game = get_game()
#     # TODO: Return map of explored locations
#     pass


# @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()