Update MCP server + exploration agent

mcp_server.py

@@ -40,170 +40,294 @@ from games.zork_env import TextAdventureEnv
 
 mcp = FastMCP("Student Text Adventure Server")
 
+MOVE_ACTIONS = {
+    "north", "south", "east", "west", "up", "down", "enter", "exit",
+    "n", "s", "e", "w", "u", "d"
+}
+
+
+def _first_line(obs: str) -> str:
+    obs = (obs or "").strip().replace("\r", "")
+    if not obs:
+        return "Unknown"
+    return obs.splitlines()[0].strip() or "Unknown"
 
-# =============================================================================
-# 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):
-        self.env: TextAdventureEnv = None
-        self.state = None
+        self.env: Optional[TextAdventureEnv] = None
+        self.state: Any = None
         self.game_name: str = ""
-        # TODO: Add more state tracking
-        # self.history: list[tuple[str, str]] = []
-        # self.explored_locations: dict[str, set[str]] = {}
-        # self.current_location: str = ""
-    
-    def initialize(self, game: str = "zork1"):
-        """Initialize or reset the game."""
+
+        self.history: list[dict[str, Any]] = []  # {action, obs_short, reward, score}
+        self.explored: dict[str, dict[str, str]] = {}  # loc_key -> {move_action -> loc_key}
+        self.current_location_key: str = "Unknown"
+
+    def initialize(self, game: str = "zork1") -> str:
         self.game_name = game
         self.env = TextAdventureEnv(game)
         self.state = self.env.reset()
-        # TODO: Reset your state tracking here
+
+        self.history = []
+        self.explored = {}
+        self.current_location_key = self._get_location_key(self.state.observation)
+
         return self.state.observation
-    
-    def step(self, action: str) -> str:
-        """Execute an action and return the result."""
+
+    def _jericho_env(self) -> Optional[Any]:
+        # TextAdventureEnv often stores jericho env as .env
         if self.env is None:
-            self.initialize()
-        
+            return None
+        inner = getattr(self.env, "env", None)
+        return inner
+
+    def _get_location_key(self, observation: str) -> str:
+        """
+        Prefer Jericho stable location id if available; fallback to header line.
+        """
+        inner = self._jericho_env()
+        # Some Jericho wrappers expose get_player_location() -> int
+        if inner is not None and hasattr(inner, "get_player_location"):
+            try:
+                loc_id = inner.get_player_location()
+                return f"id:{loc_id}"
+            except Exception:
+                pass
+
+        # Some expose a world-state hash / state id
+        if inner is not None and hasattr(inner, "get_world_state_hash"):
+            try:
+                h = inner.get_world_state_hash()
+                return f"hash:{h}"
+            except Exception:
+                pass
+
+        return _first_line(observation)[:80]
+
+    def step(self, action: str) -> str:
+        if self.env is None or self.state is None:
+            game = os.environ.get("GAME", "zork1")
+            self.initialize(game)
+
+        prev_loc = self.current_location_key
+        prev_score = int(getattr(self.state, "score", 0))
+
         self.state = self.env.step(action)
-        
-        # TODO: Update your state tracking here
-        # self.history.append((action, self.state.observation))
-        # Update location tracking, etc.
-        
-        return self.state.observation
-    
+        obs = self.state.observation
+
+        new_loc = self._get_location_key(obs)
+        self.current_location_key = new_loc
+
+        reward = int(getattr(self.state, "reward", 0))
+        score = int(getattr(self.state, "score", 0))
+
+        # Map updates on movement
+        a = (action or "").strip().lower()
+        if a in MOVE_ACTIONS:
+            self.explored.setdefault(prev_loc, {})
+            # Only record if movement changed location key
+            if new_loc != prev_loc:
+                self.explored[prev_loc][a] = new_loc
+
+        # History (keep bounded)
+        obs_short = self._summarize(obs)
+        self.history.append(
+            {
+                "action": a,
+                "obs": obs_short,
+                "reward": reward,
+                "score": score,
+                "delta": score - prev_score,
+            }
+        )
+        if len(self.history) > 80:
+            self.history = self.history[-80:]
+
+        return obs
+
+    def _summarize(self, obs: str) -> str:
+        txt = (obs or "").strip().replace("\r", "")
+        lines = [ln.strip() for ln in txt.splitlines() if ln.strip()]
+        s = " ".join(lines[:2])
+        s = re.sub(r"\s+", " ", s)
+        return s[:140]
+
     def get_score(self) -> int:
-        """Get current score."""
-        return self.state.score if self.state else 0
-    
+        return int(getattr(self.state, "score", 0)) if self.state else 0
+
     def get_moves(self) -> int:
-        """Get number of moves taken."""
-        return self.state.moves if self.state else 0
+        return int(getattr(self.state, "moves", 0)) if self.state else 0
+
+    def is_done(self) -> bool:
+        return bool(getattr(self.state, "done", False)) if self.state else False
+
+    def get_inventory_text(self) -> str:
+        """
+        Try to read inventory without consuming a move.
+        Fallback: call jericho inventory methods if available.
+        """
+        if self.state is None:
+            return "Inventory: (unknown)"
+
+        # If state has inventory field
+        inv = getattr(self.state, "inventory", None)
+        if inv:
+            return "Inventory: " + self._format_inventory(inv)
+
+        inner = self._jericho_env()
+        if inner is not None and hasattr(inner, "get_inventory"):
+            try:
+                inv2 = inner.get_inventory()
+                return "Inventory: " + self._format_inventory(inv2)
+            except Exception:
+                pass
+
+        # LAST resort: do not step("inventory") because it may cost moves; but if nothing else works:
+        return "Inventory: (unavailable)"
+
+    def _format_inventory(self, inv: Any) -> str:
+        if inv is None:
+            return "(empty)"
+        # inv may be list of objects / strings
+        try:
+            items = list(inv)
+        except Exception:
+            return str(inv)
+
+        cleaned = []
+        for it in items:
+            s = str(it)
+            s = re.sub(r"\s+", " ", s).strip()
+            # Heuristic cleanup
+            if ":" in s:
+                s = s.split(":", 1)[-1].strip()
+            cleaned.append(s)
+
+        if not cleaned:
+            return "(empty)"
+        return ", ".join(cleaned[:30])
+
+    def get_valid_actions_list(self) -> list[str]:
+        inner = self._jericho_env()
+        if inner is not None and hasattr(inner, "get_valid_actions"):
+            try:
+                acts = inner.get_valid_actions()
+                return [str(a) for a in acts]
+            except Exception:
+                return []
+        return []
 
 
-# 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)
+    if _game.env is None or _game.state is None:
         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
+    Execute a game command and return the result, with appended score/moves info.
     """
     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)
-    
-    # Optional: Append score info
-    # result += f"\n[Score: {game.get_score()} | Moves: {game.get_moves()}]"
-    
-    return result
-
-
-# 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"
 
+    action_clean = (action or "").strip()
+    if not action_clean:
+        action_clean = "look"
+
+    obs = game.step(action_clean)
+
+    score = game.get_score()
+    moves = game.get_moves()
+    reward = int(getattr(game.state, "reward", 0))
+    done = game.is_done()
+
+    suffix = f"\n\n[Score: {score} | Moves: {moves}]"
+    if reward > 0:
+        suffix = f"\n\n+{reward} points! (Total: {score})"
+    if done:
+        suffix += "\n\nGAME OVER"
+
+    return obs + suffix
+
+
+@mcp.tool()
+def memory() -> str:
+    """
+    Return current game summary: location, score, moves, inventory, recent history, observation.
+    """
+    game = get_game()
+    loc = game.current_location_key
+    score = game.get_score()
+    moves = game.get_moves()
+    inv = game.get_inventory_text()
+
+    recent = game.history[-6:]
+    if recent:
+        hist_lines = "\n".join(
+            [f"  > {h['action']} -> {h['obs']} (Δ{h['delta']})" for h in recent]
+        )
+    else:
+        hist_lines = "  (none)"
+
+    return (
+        "Current State:\n"
+        f"LocationID: {loc}\n"
+        f"Score: {score}\n"
+        f"Moves: {moves}\n"
+        f"{inv}\n\n"
+        "Recent Actions:\n"
+        f"{hist_lines}\n\n"
+        "Current Observation:\n"
+        f"{game.state.observation if game.state else ''}"
+    )
+
+
+@mcp.tool()
+def inventory() -> str:
+    """
+    Return inventory text (non-consuming when possible).
+    """
+    game = get_game()
+    return game.get_inventory_text()
+
+
+@mcp.tool()
+def get_map() -> str:
+    """
+    Return explored location graph (from movement actions).
+    """
+    game = get_game()
+    if not game.explored:
+        return "Map: (empty) Try moving around."
+
+    lines = ["Explored Map:"]
+    for loc, exits in sorted(game.explored.items(), key=lambda x: x[0]):
+        lines.append(f"\n* {loc}")
+        for a, dst in sorted(exits.items(), key=lambda x: x[0]):
+            lines.append(f"    {a} -> {dst}")
+
+    lines.append(f"\n[Current] {game.current_location_key}")
+    return "\n".join(lines)
+
+
+@mcp.tool()
+def get_valid_actions() -> str:
+    """
+    Return Jericho get_valid_actions() list (trimmed).
+    Output is JSON for easier parsing.
+    """
+    game = get_game()
+    acts = game.get_valid_actions_list()
+    # Trim to keep messages small
+    acts = acts[:120] if acts else []
+    return json.dumps(acts)
 
-# =============================================================================
-# Run the server
-# =============================================================================
 
 if __name__ == "__main__":
-    # This runs the server with stdio transport (for MCP clients)
-    mcp.run()
+    mcp.run()