nethack_player_bridge.py

#!/usr/bin/env python3
"""
NetHackPlayer Bridge Script

Manages the Claude Agent SDK session for playing NetHack.
Communicates with the Swift app via JSON lines over stdin/stdout.

Protocol:
  Input (stdin):  {"command": "start_run"} | {"command": "stop_run"} |
                  {"command": "shutdown"} | {"command": "user_message", "text": "..."}
  Output (stdout): JSON lines with type field (session_ready, agent_text, agent_action, etc.)
"""

import asyncio
import atexit
import json
import re
import subprocess
import sys
import signal
from typing import Any

from claude_agent_sdk import (
    ClaudeSDKClient,
    ClaudeAgentOptions,
    AssistantMessage,
    ResultMessage,
    TextBlock,
    ThinkingBlock,
    ToolUseBlock,
    ToolResultBlock,
)

NETHACK_INSTRUCTIONS = """
## ABSOLUTE RULE — ONE RUN ONLY
You are playing ONE single game of NetHack. When the game ends (death, quit, escape, ascension),
you MUST emit the run_summary and STOP. You must NEVER start a new game. NEVER run nethack again.
NEVER send commands to start another session. Your entire purpose is to play ONE run and report the results.
If you find yourself at a point where the game is over, DO NOT CONTINUE. Emit the summary and stop.
Violating this rule is the worst possible outcome.

You are playing NetHack autonomously. Your goal is to survive and explore as deeply as possible.

## Controls
- Use `tmux send-keys -t nethack "KEY"` to send commands
- Use `tmux capture-pane -t nethack -p` to read the current screen
- Movement: hjklyubn (vi keys)
- Always read the screen after each action

## Saved Game Restore
When NetHack loads a saved game, it shows a restore screen with messages like
"Restoring save file..." and possibly "--More--" prompts. The game will NOT accept
movement commands until these are dismissed. After your first screen read:
1. If the turn counter (T:xxx) is > 1, this is a restored game
2. Press space repeatedly (checking the screen each time) until the message line is clear
3. Only then attempt movement — if movement doesn't change the turn counter, there's
   still a prompt to dismiss
4. Look for ANY text on the message line (line 1 or line 22) that might need a response

## Strategy
- IMPORTANT: Always invoke the /nethack skill at the start of a session to load full strategy knowledge
- Always check HP before combat
- Identify items before using them
- Pick up food to prevent starvation

## Player Name
Use "Claude" as the player name when prompted.

## Screen Layout
- Lines 1-21: Dungeon map
- Line 22: Message line
- Lines 23-24: Status lines (HP, AC, Dlvl, etc.)

Always look for "--More--" prompts and press space to continue.

## End-of-Game Handling
When you die or the game ends:
1. Dismiss all "--More--" prompts by pressing space
2. When asked "Do you want your possessions identified?" answer 'y'
3. Read ALL end-game screens (tombstone, inventory list, score)
4. After reading everything, emit a structured JSON block with your run summary in this exact format:

```run_summary
{
  "cause_of_death": "killed by a newt",
  "score": 42,
  "dungeon_level": 1,
  "experience_level": 1,
  "turns": 150,
  "gold": 0,
  "hp_max": 16,
  "player_name": "Claude",
  "player_class": "Valkyrie",
  "player_race": "Human",
  "inventory": ["a +0 long sword", "a +3 small shield"],
  "death_screen_text": "raw text of the final screens"
}
```

Fill in all fields from what you observed on the death/end screens. Use null for any field you couldn't determine.

IMPORTANT: After emitting the run_summary block, your task is COMPLETE. Stop IMMEDIATELY.
- Do NOT start a new game.
- Do NOT send any more commands to tmux.
- Do NOT run nethack again.
- Do NOT create a new tmux session.
- Simply end your response after the summary. You are DONE. The user must manually start the next run.

## Important
- Start by reading the screen with `tmux capture-pane -t nethack -p`
- After every action, read the screen again to see the result
- If you see a menu or prompt, respond to it appropriately
- Never send multiple keys without checking the screen between them
- If movement doesn't advance the turn counter, something is blocking input — press space or Enter to dismiss it
"""

# Key translation table for human-readable action descriptions
KEY_ACTIONS: dict[str, str] = {
    "h": "Moving west",
    "l": "Moving east",
    "j": "Moving south",
    "k": "Moving north",
    "y": "Moving northwest",
    "u": "Moving northeast",
    "b": "Moving southwest",
    "n": "Moving southeast",
    ".": "Waiting",
    "s": "Searching",
    "i": "Opening inventory",
    ",": "Picking up item",
    "e": "Eating",
    "q": "Quaffing",
    "d": "Dropping item",
    ">": "Going downstairs",
    "<": "Going upstairs",
    ";": "Looking around",
    ":": "Looking at floor",
    "@": "Toggling autopickup",
    "#": "Extended command",
    "?": "Getting help",
    " ": "Continuing",
    "Space": "Continuing",
    "Enter": "Confirming",
    "C-m": "Confirming",
}



def emit(msg: dict[str, Any]) -> None:
    """Write a JSON message to stdout."""
    line = json.dumps(msg)
    sys.stdout.write(line + "\n")
    sys.stdout.flush()


def translate_send_keys(command: str) -> str | None:
    """Extract the key from a tmux send-keys command and return a human-readable action."""
    # Match patterns like: tmux send-keys -t nethack "h"
    # or: tmux send-keys -t nethack h
    match = re.search(r'tmux\s+send-keys\s+.*?(?:["\'](.+?)["\']|(\S+))\s*$', command)
    if not match:
        return None
    key = match.group(1) or match.group(2)
    if key in KEY_ACTIONS:
        return KEY_ACTIONS[key]
    return f"Sending command: {key}"


def is_capture_pane(command: str) -> bool:
    """Check if a command is a tmux capture-pane command."""
    return "tmux" in command and "capture-pane" in command


def extract_run_summary(text: str) -> dict | None:
    """Extract a run_summary JSON block from agent text."""
    match = re.search(r'```run_summary\s*\n(.*?)\n```', text, re.DOTALL)
    if not match:
        return None
    try:
        return json.loads(match.group(1))
    except (json.JSONDecodeError, ValueError):
        return None


def strip_run_summary_block(text: str) -> str:
    """Remove the run_summary fenced block from text."""
    return re.sub(r'```run_summary\s*\n.*?\n```', '', text, flags=re.DOTALL).strip()


def create_tmux_session() -> bool:
    """Create a tmux session and start NetHack."""
    try:
        # Kill any existing session
        subprocess.run(
            ["tmux", "kill-session", "-t", "nethack"],
            capture_output=True,
        )

        # Create new session at standard NetHack dimensions
        result = subprocess.run(
            ["tmux", "new-session", "-d", "-s", "nethack", "-x", "80", "-y", "24"],
            capture_output=True,
            text=True,
        )
        if result.returncode != 0:
            emit({"type": "error", "message": f"Failed to create tmux session: {result.stderr}"})
            return False

        # Hide the tmux status bar for a cleaner display
        subprocess.run(
            ["tmux", "set-option", "-t", "nethack", "status", "off"],
            capture_output=True,
        )

        # Start nethack
        subprocess.run(
            ["tmux", "send-keys", "-t", "nethack", "nethack", "C-m"],
            capture_output=True,
        )

        return True
    except FileNotFoundError:
        emit({"type": "error", "message": "tmux not found. Install with: brew install tmux"})
        return False


def kill_tmux_session() -> None:
    """Kill the nethack tmux session."""
    subprocess.run(
        ["tmux", "kill-session", "-t", "nethack"],
        capture_output=True,
    )


def debug(msg: str) -> None:
    """Write debug message to stderr (visible in Xcode console)."""
    sys.stderr.write(f"[bridge] {msg}\n")
    sys.stderr.flush()


async def run_agent(client: ClaudeSDKClient, interrupt_flag: list[bool], user_message_flag: list[bool]) -> None:
    """Process messages from the Claude Agent SDK client."""
    # Track tool_use IDs for capture-pane commands so we can suppress their results
    capture_pane_ids: set[str] = set()
    # Accumulate all text blocks to detect run_summary at end
    accumulated_text: list[str] = []
    # Flag: once we detect a run_summary, stop processing immediately
    run_summary_emitted = False

    debug("Listening for agent messages")
    msg_count = 0

    try:
        async for message in client.receive_messages():
            msg_count += 1
            msg_type = type(message).__name__
            debug(f"msg#{msg_count} type={msg_type}")
            if isinstance(message, AssistantMessage):
                for block in message.content:
                    block_type = type(block).__name__
                    if isinstance(block, TextBlock):
                        accumulated_text.append(block.text)
                        # Strip any run_summary block before emitting to chat
                        cleaned = strip_run_summary_block(block.text)
                        if cleaned:
                            debug(f"emit agent_text ({len(cleaned)} chars)")
                            emit({"type": "agent_text", "content": cleaned})
                        # Proactive detection: if a run_summary appeared, stop immediately
                        if not run_summary_emitted:
                            full_text = "\n".join(accumulated_text)
                            summary = extract_run_summary(full_text)
                            if summary:
                                debug("Run summary detected mid-run — stopping agent")
                                run_summary_emitted = True
                                emit({"type": "run_summary", **summary})
                                emit({"type": "run_ended", "reason": "completed"})
                                return
                    elif isinstance(block, ThinkingBlock):
                        debug(f"emit agent_thinking ({len(block.thinking)} chars)")
                        emit({"type": "agent_thinking", "content": block.thinking})
                    elif isinstance(block, ToolUseBlock):
                        command_str = block.input.get("command", "") if block.input else ""
                        debug(f"ToolUse: {block.name} cmd={command_str[:60]}")
                        # Detect attempt to restart nethack
                        if command_str and ("nethack" in command_str.lower()) and ("send-keys" in command_str) and not is_capture_pane(command_str):
                            # Agent is trying to type "nethack" into the terminal — block it
                            debug("BLOCKED: Agent attempted to restart nethack via send-keys")
                            emit({"type": "agent_action", "content": "Blocked: attempted to restart game"})
                            emit({"type": "run_ended", "reason": "completed"})
                            return
                        if command_str and is_capture_pane(command_str):
                            capture_pane_ids.add(block.id)
                            emit({"type": "agent_action", "content": "Reading the screen"})
                        elif command_str and "send-keys" in command_str:
                            action = translate_send_keys(command_str)
                            if action:
                                emit({"type": "agent_action", "content": action})
                            else:
                                emit({"type": "agent_action", "content": f"Sending command"})
                        else:
                            # Non-tmux tool use — emit as generic tool_use
                            input_str = json.dumps(block.input) if block.input else ""
                            emit({
                                "type": "agent_tool_use",
                                "name": block.name,
                                "input": input_str,
                            })
                    elif isinstance(block, ToolResultBlock):
                        if block.tool_use_id in capture_pane_ids:
                            # Suppress capture-pane results (visible in terminal panel)
                            capture_pane_ids.discard(block.tool_use_id)
                            continue
                        if block.is_error:
                            content = ""
                            if isinstance(block.content, str):
                                content = block.content
                            elif isinstance(block.content, list):
                                parts = []
                                for item in block.content:
                                    if isinstance(item, dict) and "text" in item:
                                        parts.append(item["text"])
                                content = "\n".join(parts)
                            emit({
                                "type": "agent_tool_result",
                                "content": content,
                                "is_error": True,
                            })
            elif isinstance(message, ResultMessage):
                debug(f"ResultMessage received (is_error={message.is_error})")
                if interrupt_flag[0]:
                    debug("Ignoring ResultMessage from interrupt")
                    interrupt_flag[0] = False
                    continue
                if user_message_flag[0]:
                    debug("User message answered, sending continuation query")
                    user_message_flag[0] = False
                    await client.query(
                        "Continue playing NetHack. Read the screen and make your next moves."
                    )
                    continue
                reason = "completed" if not message.is_error else "error"
                if message.result:
                    reason = message.result
                # Check accumulated text for a run summary
                full_text = "\n".join(accumulated_text)
                summary = extract_run_summary(full_text)
                if summary:
                    debug("Emitting run_summary before run_ended")
                    emit({"type": "run_summary", **summary})
                emit({"type": "run_ended", "reason": reason})
                return
            else:
                debug(f"Skipping message type: {msg_type}")

    except asyncio.CancelledError:
        raise
    except Exception as e:
        debug(f"Error: {e}")
        emit({"type": "error", "message": str(e)})


async def main() -> None:
    """Main loop: read commands from stdin, dispatch actions."""
    running = True
    client: ClaudeSDKClient | None = None
    agent_task: asyncio.Task | None = None
    interrupt_flag: list[bool] = [False]
    user_message_flag: list[bool] = [False]

    loop = asyncio.get_event_loop()

    # Handle graceful shutdown
    def handle_signal(sig, frame):
        nonlocal running
        running = False
        if agent_task and not agent_task.done():
            agent_task.cancel()
        kill_tmux_session()

    signal.signal(signal.SIGTERM, handle_signal)
    signal.signal(signal.SIGINT, handle_signal)

    reader = asyncio.StreamReader()
    protocol = asyncio.StreamReaderProtocol(reader)
    await loop.connect_read_pipe(lambda: protocol, sys.stdin)

    while running:
        try:
            line = await reader.readline()
            if not line:
                break

            line_str = line.decode().strip()
            if not line_str:
                continue

            try:
                cmd = json.loads(line_str)
            except json.JSONDecodeError:
                emit({"type": "error", "message": f"Invalid JSON: {line_str}"})
                continue

            command = cmd.get("command")

            if command == "start_run":
                if agent_task and not agent_task.done():
                    emit({"type": "error", "message": "A run is already in progress"})
                    continue

                if not create_tmux_session():
                    continue

                # Give nethack a moment to start
                await asyncio.sleep(1.5)

                emit({"type": "session_ready", "session_name": "nethack"})

                # Create and connect the client
                options = ClaudeAgentOptions(
                    system_prompt={
                        "type": "preset",
                        "preset": "claude_code",
                        "append": NETHACK_INSTRUCTIONS,
                    },
                    allowed_tools=["Bash", "Read", "Glob", "Grep", "Skill"],
                    disallowed_tools=["Write", "Edit", "WebFetch", "WebSearch"],
                    permission_mode="bypassPermissions",
                    setting_sources=["user"],
                    hooks={},
                )

                client = ClaudeSDKClient(options=options)

                prompt = (
                    "Start playing NetHack. First, invoke the /nethack skill "
                    "to load your full strategy knowledge. Then read the current "
                    "screen with `tmux capture-pane -t nethack -p` and begin playing. "
                    "Continue playing until your character dies or I stop you. "
                    "When your character dies, follow the End-of-Game Handling "
                    "instructions to emit a run_summary block, then STOP completely. "
                    "Do NOT start a new game."
                )

                debug("Connecting and starting agent")
                await client.connect()
                agent_task = asyncio.create_task(run_agent(client, interrupt_flag, user_message_flag))
                await asyncio.sleep(0.1)
                await client.query(prompt)

            elif command == "stop_run":
                if agent_task and not agent_task.done():
                    agent_task.cancel()
                    try:
                        await agent_task
                    except asyncio.CancelledError:
                        pass
                if client:
                    await client.disconnect()
                    client = None
                kill_tmux_session()
                emit({"type": "run_ended", "reason": "stopped"})

            elif command == "user_message":
                text = cmd.get("text", "")
                if not text:
                    continue
                if client:
                    debug(f"User message: {text[:60]}")
                    interrupt_flag[0] = True
                    user_message_flag[0] = True
                    try:
                        await client.interrupt()
                    except Exception:
                        pass
                    await client.query(text)
                else:
                    emit({"type": "error", "message": "No active session for user message"})

            elif command == "shutdown":
                if agent_task and not agent_task.done():
                    agent_task.cancel()
                    try:
                        await agent_task
                    except asyncio.CancelledError:
                        pass
                if client:
                    await client.disconnect()
                    client = None
                kill_tmux_session()
                running = False

        except asyncio.CancelledError:
            break
        except Exception as e:
            emit({"type": "error", "message": f"Bridge error: {str(e)}"})

    # Final cleanup
    if client:
        try:
            await client.disconnect()
        except Exception:
            pass


if __name__ == "__main__":
    # Safety net: ensure tmux session is cleaned up on any exit
    atexit.register(kill_tmux_session)
    asyncio.run(main())