claudemol 0.1.0__py3-none-any.whl

claudemol/__init__.py

@@ -0,0 +1,32 @@
+"""
+claudemol: PyMOL integration for Claude Code
+
+Connect to PyMOL via socket for AI-assisted molecular visualization.
+"""
+
+from claudemol.connection import (
+    PyMOLConnection,
+    connect_or_launch,
+    launch_pymol,
+    find_pymol_command,
+    check_pymol_installed,
+)
+from claudemol.session import (
+    PyMOLSession,
+    get_session,
+    ensure_running,
+    stop_pymol,
+)
+
+__version__ = "0.1.0"
+__all__ = [
+    "PyMOLConnection",
+    "PyMOLSession",
+    "connect_or_launch",
+    "launch_pymol",
+    "find_pymol_command",
+    "check_pymol_installed",
+    "get_session",
+    "ensure_running",
+    "stop_pymol",
+]

claudemol/cli.py

@@ -0,0 +1,167 @@
+"""
+CLI for claudemol setup and management.
+
+Usage:
+    claudemol setup    # Configure PyMOL to auto-load the socket plugin
+    claudemol status   # Check if PyMOL is running and connected
+    claudemol test     # Test the connection
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+from claudemol.connection import (
+    PyMOLConnection,
+    check_pymol_installed,
+    find_pymol_command,
+    get_plugin_path,
+)
+
+
+def setup_pymol():
+    """Configure PyMOL to auto-load the socket plugin."""
+    plugin_path = get_plugin_path()
+    if not plugin_path.exists():
+        print(f"Error: Plugin not found at {plugin_path}", file=sys.stderr)
+        return 1
+
+    pymolrc_path = Path.home() / ".pymolrc"
+
+    # Check if already configured
+    if pymolrc_path.exists():
+        content = pymolrc_path.read_text()
+        if "claudemol" in content or "claude_socket_plugin" in content:
+            print("PyMOL already configured for claudemol.")
+            print(f"Plugin: {plugin_path}")
+            return 0
+
+    # Add to .pymolrc
+    run_command = f'\n# claudemol: Claude Code integration\nrun {plugin_path}\n'
+
+    if pymolrc_path.exists():
+        with open(pymolrc_path, "a") as f:
+            f.write(run_command)
+        print(f"Added claudemol plugin to existing {pymolrc_path}")
+    else:
+        pymolrc_path.write_text(run_command.lstrip())
+        print(f"Created {pymolrc_path} with claudemol plugin")
+
+    print(f"Plugin path: {plugin_path}")
+    print("\nSetup complete! The plugin will auto-load when you start PyMOL.")
+
+    # Check if PyMOL is installed
+    if not check_pymol_installed():
+        print("\nNote: PyMOL not found in PATH.")
+        print("Install PyMOL with one of:")
+        print("  - pip install pymol-open-source-whl")
+        print("  - brew install pymol (macOS)")
+        print("  - Download from https://pymol.org")
+
+    return 0
+
+
+def check_status():
+    """Check PyMOL connection status."""
+    print("Checking PyMOL status...")
+
+    # Check if PyMOL is installed
+    pymol_cmd = find_pymol_command()
+    if pymol_cmd:
+        print(f"PyMOL found: {' '.join(pymol_cmd)}")
+    else:
+        print("PyMOL not found in PATH")
+        return 1
+
+    # Try to connect
+    conn = PyMOLConnection()
+    try:
+        conn.connect(timeout=2.0)
+        print("Socket connection: OK (port 9880)")
+        conn.disconnect()
+        return 0
+    except ConnectionError:
+        print("Socket connection: Not available")
+        print("  (PyMOL may not be running, or plugin not loaded)")
+        return 1
+
+
+def test_connection():
+    """Test the PyMOL connection with a simple command."""
+    conn = PyMOLConnection()
+    try:
+        conn.connect(timeout=2.0)
+        result = conn.execute("print('claudemol connection test')")
+        print("Connection test: OK")
+        print(f"Response: {result}")
+        conn.disconnect()
+        return 0
+    except ConnectionError as e:
+        print(f"Connection failed: {e}", file=sys.stderr)
+        print("\nMake sure PyMOL is running with the socket plugin.")
+        print("Start PyMOL and run: claude_status")
+        return 1
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)
+        return 1
+
+
+def show_info():
+    """Show claudemol installation info."""
+    plugin_path = get_plugin_path()
+    pymolrc_path = Path.home() / ".pymolrc"
+
+    print("claudemol installation info:")
+    print(f"  Plugin: {plugin_path}")
+    print(f"  Plugin exists: {plugin_path.exists()}")
+    print(f"  .pymolrc: {pymolrc_path}")
+    print(f"  .pymolrc exists: {pymolrc_path.exists()}")
+
+    if pymolrc_path.exists():
+        content = pymolrc_path.read_text()
+        configured = "claudemol" in content or "claude_socket_plugin" in content
+        print(f"  Configured in .pymolrc: {configured}")
+
+    pymol_cmd = find_pymol_command()
+    print(f"  PyMOL command: {' '.join(pymol_cmd) if pymol_cmd else 'not found'}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="claudemol: PyMOL integration for Claude Code",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Commands:
+  setup     Configure PyMOL to auto-load the socket plugin
+  status    Check if PyMOL is running and connected
+  test      Test the connection with a simple command
+  info      Show installation info
+
+For Claude Code skills, install the claudemol-skills plugin:
+  /plugin marketplace add ANaka/ai-mol?path=claude-plugin
+  /plugin install claudemol-skills
+""",
+    )
+    parser.add_argument(
+        "command",
+        nargs="?",
+        choices=["setup", "status", "test", "info"],
+        default="info",
+        help="Command to run",
+    )
+
+    args = parser.parse_args()
+
+    if args.command == "setup":
+        return setup_pymol()
+    elif args.command == "status":
+        return check_status()
+    elif args.command == "test":
+        return test_connection()
+    elif args.command == "info":
+        show_info()
+        return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

claudemol/connection.py

@@ -0,0 +1,249 @@
+"""
+PyMOL Connection Module
+
+Provides functions for Claude Code to communicate with PyMOL via socket.
+"""
+
+import json
+import os
+import shutil
+import socket
+import subprocess
+import time
+from pathlib import Path
+
+DEFAULT_HOST = "localhost"
+DEFAULT_PORT = 9880
+CONNECT_TIMEOUT = 5.0
+RECV_TIMEOUT = 30.0
+
+# Common PyMOL installation paths
+PYMOL_PATHS = [
+    # uv environment (created by /pymol-setup)
+    os.path.expanduser("~/.pymol-env/bin/python"),
+    # macOS app locations
+    "/Applications/PyMOL.app/Contents/MacOS/PyMOL",
+    os.path.expanduser("~/Applications/PyMOL.app/Contents/MacOS/PyMOL"),
+    # Linux system locations
+    "/usr/bin/pymol",
+    "/usr/local/bin/pymol",
+]
+
+
+class PyMOLConnection:
+    def __init__(self, host=DEFAULT_HOST, port=DEFAULT_PORT):
+        self.host = host
+        self.port = port
+        self.socket = None
+
+    def connect(self, timeout=CONNECT_TIMEOUT):
+        """Connect to PyMOL socket server."""
+        if self.socket:
+            return True
+        try:
+            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            self.socket.settimeout(timeout)
+            self.socket.connect((self.host, self.port))
+            self.socket.settimeout(RECV_TIMEOUT)
+            return True
+        except Exception as e:
+            self.socket = None
+            raise ConnectionError(
+                f"Cannot connect to PyMOL on {self.host}:{self.port}: {e}"
+            )
+
+    def disconnect(self):
+        """Disconnect from PyMOL."""
+        if self.socket:
+            try:
+                self.socket.close()
+            except OSError:
+                pass
+            self.socket = None
+
+    def is_connected(self):
+        """Check if connected to PyMOL."""
+        if not self.socket:
+            return False
+        try:
+            self.socket.setblocking(False)
+            try:
+                data = self.socket.recv(1, socket.MSG_PEEK)
+                if data == b"":
+                    self.disconnect()
+                    return False
+            except BlockingIOError:
+                pass
+            finally:
+                self.socket.setblocking(True)
+                self.socket.settimeout(RECV_TIMEOUT)
+            return True
+        except OSError:
+            self.disconnect()
+            return False
+
+    def send_command(self, code):
+        """Send Python code to PyMOL and return result."""
+        if not self.socket:
+            raise ConnectionError("Not connected to PyMOL")
+        try:
+            message = json.dumps({"type": "execute", "code": code})
+            self.socket.sendall(message.encode("utf-8"))
+            response = b""
+            while True:
+                chunk = self.socket.recv(4096)
+                if not chunk:
+                    raise ConnectionError("Connection closed by PyMOL")
+                response += chunk
+                try:
+                    result = json.loads(response.decode("utf-8"))
+                    return result
+                except json.JSONDecodeError:
+                    continue
+        except socket.timeout:
+            raise TimeoutError("PyMOL command timed out")
+        except Exception as e:
+            self.disconnect()
+            raise ConnectionError(f"Communication error: {e}")
+
+    def execute(self, code):
+        """Execute code, reconnecting if necessary. Returns output string or raises."""
+        for attempt in range(3):
+            try:
+                if not self.is_connected():
+                    self.connect()
+                result = self.send_command(code)
+                if result.get("status") == "success":
+                    return result.get("output", "")
+                else:
+                    raise RuntimeError(result.get("error", "Unknown error"))
+            except ConnectionError:
+                if attempt < 2:
+                    time.sleep(0.5)
+                    continue
+                raise
+        raise ConnectionError("Failed to connect after 3 attempts")
+
+
+def find_pymol_command():
+    """
+    Find how to launch PyMOL.
+
+    Returns:
+        List of command arguments (e.g., ["pymol"] or ["python", "-m", "pymol"])
+        or None if PyMOL is not found.
+    """
+    # Check if pymol is in PATH
+    pymol_path = shutil.which("pymol")
+    if pymol_path:
+        return [pymol_path]
+
+    # Check uv environment first (most common for this project)
+    uv_python = os.path.expanduser("~/.pymol-env/bin/python")
+    if os.path.isfile(uv_python) and os.access(uv_python, os.X_OK):
+        # Verify pymol is installed in this environment
+        try:
+            result = subprocess.run(
+                [uv_python, "-c", "import pymol"],
+                capture_output=True,
+                timeout=5,
+            )
+            if result.returncode == 0:
+                return [uv_python, "-m", "pymol"]
+        except (subprocess.TimeoutExpired, FileNotFoundError):
+            pass
+
+    # Check other common paths
+    for path in PYMOL_PATHS:
+        if path.endswith("/python"):
+            continue  # Already checked uv environment above
+        if os.path.isfile(path) and os.access(path, os.X_OK):
+            return [path]
+
+    return None
+
+
+def find_pymol_executable():
+    """Find PyMOL executable path (legacy compatibility)."""
+    cmd = find_pymol_command()
+    return cmd[0] if cmd else None
+
+
+def check_pymol_installed():
+    """Check if pymol command is available."""
+    return find_pymol_command() is not None
+
+
+def get_plugin_path():
+    """Get the path to the socket plugin file."""
+    return Path(__file__).parent / "plugin.py"
+
+
+def launch_pymol(file_path=None, wait_for_socket=True, timeout=10.0):
+    """
+    Launch PyMOL with the Claude socket plugin.
+
+    Args:
+        file_path: Optional file to open (e.g., .pdb, .cif)
+        wait_for_socket: Wait for socket to become available
+        timeout: How long to wait for socket
+
+    Returns:
+        subprocess.Popen process handle
+    """
+    pymol_cmd = find_pymol_command()
+    if not pymol_cmd:
+        raise RuntimeError(
+            "PyMOL not found. Please install PyMOL:\n"
+            "  - Run: ai-mol setup\n"
+            "  - Or: pip install pymol-open-source-whl\n"
+            "  - Or: brew install pymol (macOS)"
+        )
+
+    plugin_path = get_plugin_path()
+    if not plugin_path.exists():
+        raise RuntimeError(f"Plugin not found: {plugin_path}")
+
+    # Build command: base pymol command + optional file + plugin
+    cmd_args = list(pymol_cmd)
+    if file_path:
+        cmd_args.append(str(file_path))
+    cmd_args.extend(["-d", f"run {plugin_path}"])
+
+    process = subprocess.Popen(cmd_args)
+
+    if wait_for_socket:
+        start = time.time()
+        while time.time() - start < timeout:
+            try:
+                conn = PyMOLConnection()
+                conn.connect(timeout=1.0)
+                conn.disconnect()
+                return process
+            except ConnectionError:
+                time.sleep(0.5)
+        raise TimeoutError(f"PyMOL socket not available after {timeout}s")
+
+    return process
+
+
+def connect_or_launch(file_path=None):
+    """
+    Connect to existing PyMOL or launch new instance.
+
+    Returns:
+        (PyMOLConnection, process_or_None)
+    """
+    conn = PyMOLConnection()
+
+    # Try connecting to existing instance
+    try:
+        conn.connect(timeout=1.0)
+        return conn, None
+    except ConnectionError:
+        pass
+
+    # Launch new instance
+    process = launch_pymol(file_path=file_path)
+    conn.connect()
+    return conn, process

claudemol/plugin.py

@@ -0,0 +1,186 @@
+"""
+Claude Socket Plugin for PyMOL
+
+Headless socket listener that receives Python commands from Claude Code.
+Auto-starts on load, no UI required.
+
+Usage:
+    run /path/to/plugin.py    # Start listening
+    claude_status             # Check connection status
+    claude_stop               # Stop listener
+    claude_start              # Restart listener
+"""
+
+import socket
+import json
+import threading
+import traceback
+import io
+from contextlib import redirect_stdout
+
+from pymol import cmd
+
+# Global state
+_server = None
+_port = 9880
+
+
+class SocketServer:
+    def __init__(self, host='localhost', port=9880):
+        self.host = host
+        self.port = port
+        self.socket = None
+        self.client = None
+        self.running = False
+        self.thread = None
+
+    def start(self):
+        if self.running:
+            return False
+        self.running = True
+        self.thread = threading.Thread(target=self._run_server, daemon=True)
+        self.thread.start()
+        return True
+
+    def _run_server(self):
+        try:
+            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            self.socket.bind((self.host, self.port))
+            self.socket.listen(5)
+            self.socket.settimeout(1.0)
+
+            print(f"Claude socket listener active on port {self.port}")
+
+            while self.running:
+                try:
+                    new_client, address = self.socket.accept()
+                    if self.client:
+                        try:
+                            self.client.close()
+                        except:
+                            pass
+                    self.client = new_client
+                    self.client.settimeout(1.0)
+                    self._handle_client(address)
+                except socket.timeout:
+                    continue
+                except Exception as e:
+                    if self.running:
+                        print(f"Connection error: {e}")
+        except Exception as e:
+            print(f"Socket server error: {e}")
+            traceback.print_exc()
+        finally:
+            self._cleanup()
+
+    def _handle_client(self, address):
+        buffer = b''
+        while self.running and self.client:
+            try:
+                data = self.client.recv(4096)
+                if not data:
+                    break
+                buffer += data
+                try:
+                    command = json.loads(buffer.decode('utf-8'))
+                    buffer = b''
+                    result = self._execute_command(command)
+                    response = json.dumps(result)
+                    self.client.sendall(response.encode('utf-8'))
+                except json.JSONDecodeError:
+                    continue
+            except socket.timeout:
+                continue
+            except Exception as e:
+                if self.running:
+                    print(f"Client error: {e}")
+                break
+        if self.client:
+            try:
+                self.client.close()
+            except:
+                pass
+            self.client = None
+
+    def _execute_command(self, command):
+        code = command.get("code", "")
+        if not code:
+            return {"status": "error", "error": "No code provided"}
+        try:
+            exec_globals = {"cmd": cmd, "__builtins__": __builtins__}
+            output_buffer = io.StringIO()
+            with redirect_stdout(output_buffer):
+                exec(code, exec_globals)
+            output = output_buffer.getvalue()
+            if '_result' in exec_globals:
+                output = str(exec_globals['_result'])
+            return {"status": "success", "output": output or "OK"}
+        except Exception as e:
+            return {"status": "error", "error": str(e)}
+
+    def _cleanup(self):
+        if self.client:
+            try:
+                self.client.close()
+            except:
+                pass
+        if self.socket:
+            try:
+                self.socket.close()
+            except:
+                pass
+        self.socket = None
+        self.client = None
+        self.running = False
+
+    def stop(self):
+        self.running = False
+        if self.thread:
+            self.thread.join(2.0)
+        self._cleanup()
+
+    @property
+    def is_running(self):
+        return self.running and self.thread and self.thread.is_alive()
+
+
+def claude_status():
+    """Print Claude socket listener status."""
+    global _server
+    if _server and _server.is_running:
+        connected = "connected" if _server.client else "waiting"
+        print(f"Claude socket listener: running on port {_port} ({connected})")
+    else:
+        print("Claude socket listener: not running")
+
+
+def claude_stop():
+    """Stop the Claude socket listener."""
+    global _server
+    if _server:
+        _server.stop()
+        _server = None
+        print("Claude socket listener stopped")
+    else:
+        print("Claude socket listener was not running")
+
+
+def claude_start(port=9880):
+    """Start the Claude socket listener."""
+    global _server, _port
+    if _server and _server.is_running:
+        print(f"Claude socket listener already running on port {_port}")
+        return
+    _port = port
+    _server = SocketServer(port=port)
+    _server.start()
+
+
+# Register commands with PyMOL
+cmd.extend("claude_status", claude_status)
+cmd.extend("claude_stop", claude_stop)
+cmd.extend("claude_start", claude_start)
+
+# Auto-start on load
+claude_start()

claudemol/session.py

@@ -0,0 +1,295 @@
+"""
+PyMOL Session Manager
+
+Provides reliable session lifecycle management:
+- Launch PyMOL with plugin
+- Health checks
+- Graceful and forced termination
+- Crash detection and recovery
+"""
+
+import os
+import signal
+import subprocess
+import time
+from pathlib import Path
+
+from claudemol.connection import (
+    DEFAULT_HOST,
+    DEFAULT_PORT,
+    PyMOLConnection,
+    find_pymol_command,
+    get_plugin_path,
+)
+
+
+class PyMOLSession:
+    """
+    Manages a PyMOL session with health monitoring and recovery.
+
+    Usage:
+        session = PyMOLSession()
+        session.start()
+
+        # Use the connection
+        result = session.execute("cmd.fetch('1ubq')")
+
+        # Check health
+        if not session.is_healthy():
+            session.recover()
+
+        # Clean up
+        session.stop()
+    """
+
+    def __init__(self, host=DEFAULT_HOST, port=DEFAULT_PORT):
+        self.host = host
+        self.port = port
+        self.process = None
+        self.connection = None
+        self._we_launched = False  # Track if we started PyMOL
+
+    @property
+    def is_running(self):
+        """Check if we have a PyMOL process that's still alive."""
+        if self.process is None:
+            return False
+        return self.process.poll() is None
+
+    @property
+    def is_connected(self):
+        """Check if we have a socket connection."""
+        return self.connection is not None and self.connection.is_connected()
+
+    def is_healthy(self):
+        """
+        Check if PyMOL is responsive by executing a trivial command.
+
+        Returns True if we can communicate with PyMOL.
+        """
+        if not self.is_connected:
+            return False
+        try:
+            result = self.connection.execute("print('ping')")
+            return "ping" in result
+        except Exception:
+            return False
+
+    def start(self, timeout=15.0):
+        """
+        Start PyMOL or connect to existing instance.
+
+        Args:
+            timeout: How long to wait for PyMOL to be ready
+
+        Returns:
+            True if connected successfully
+        """
+        self.connection = PyMOLConnection(self.host, self.port)
+
+        # Try connecting to existing instance first
+        try:
+            self.connection.connect(timeout=2.0)
+            self._we_launched = False
+            return True
+        except ConnectionError:
+            pass
+
+        # Launch new instance
+        pymol_cmd = find_pymol_command()
+        if not pymol_cmd:
+            raise RuntimeError(
+                "PyMOL not found. Run: claudemol setup"
+            )
+
+        # Check if plugin is configured in pymolrc (don't double-load)
+        pymolrc_path = Path.home() / ".pymolrc"
+        plugin_in_pymolrc = False
+        if pymolrc_path.exists():
+            content = pymolrc_path.read_text()
+            if "claude_socket_plugin" in content or "claudemol" in content:
+                plugin_in_pymolrc = True
+
+        # Build command - only add plugin if not in pymolrc
+        cmd_args = list(pymol_cmd)
+        if not plugin_in_pymolrc:
+            plugin_path = get_plugin_path()
+            if not plugin_path.exists():
+                raise RuntimeError(f"Plugin not found: {plugin_path}")
+            cmd_args += ["-d", f"run {plugin_path}"]
+
+        # Launch PyMOL
+        self.process = subprocess.Popen(
+            cmd_args,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+        self._we_launched = True
+
+        # Wait for socket to become available
+        start = time.time()
+        while time.time() - start < timeout:
+            try:
+                self.connection.connect(timeout=1.0)
+                return True
+            except ConnectionError:
+                if not self.is_running:
+                    # Process died during startup
+                    stdout, stderr = self.process.communicate(timeout=1)
+                    raise RuntimeError(
+                        f"PyMOL exited during startup.\n"
+                        f"stdout: {stdout.decode()}\n"
+                        f"stderr: {stderr.decode()}"
+                    )
+                time.sleep(0.5)
+
+        # Timeout - kill the process
+        self._kill_process()
+        raise TimeoutError(f"PyMOL socket not available after {timeout}s")
+
+    def stop(self, graceful_timeout=5.0):
+        """
+        Stop PyMOL session.
+
+        Args:
+            graceful_timeout: Time to wait for graceful shutdown before force kill
+        """
+        # Disconnect socket
+        if self.connection:
+            self.connection.disconnect()
+            self.connection = None
+
+        # Only kill process if we launched it
+        if self._we_launched and self.process:
+            self._kill_process(graceful_timeout)
+
+    def _kill_process(self, graceful_timeout=5.0):
+        """Kill the PyMOL process."""
+        if not self.process:
+            return
+
+        # Try graceful termination first
+        try:
+            self.process.terminate()
+            self.process.wait(timeout=graceful_timeout)
+        except subprocess.TimeoutExpired:
+            # Force kill
+            self.process.kill()
+            self.process.wait(timeout=2.0)
+        except Exception:
+            pass
+
+        self.process = None
+        self._we_launched = False
+
+    def recover(self, timeout=15.0):
+        """
+        Recover from a crashed or unresponsive PyMOL.
+
+        This will:
+        1. Kill any stale process
+        2. Disconnect stale socket
+        3. Start fresh
+
+        Returns:
+            True if recovery successful
+        """
+        # Clean up
+        if self.connection:
+            self.connection.disconnect()
+            self.connection = None
+
+        if self._we_launched:
+            self._kill_process(graceful_timeout=2.0)
+
+        # Also try to kill any orphaned PyMOL processes on our port
+        self._kill_processes_on_port()
+
+        # Brief pause to let OS clean up
+        time.sleep(0.5)
+
+        # Start fresh
+        return self.start(timeout=timeout)
+
+    def _kill_processes_on_port(self):
+        """Kill any processes listening on our port (Linux/macOS)."""
+        try:
+            # Find PIDs using our port
+            result = subprocess.run(
+                ["lsof", "-ti", f":{self.port}"],
+                capture_output=True,
+                text=True,
+                timeout=5,
+            )
+            if result.stdout.strip():
+                pids = result.stdout.strip().split("\n")
+                for pid in pids:
+                    try:
+                        os.kill(int(pid), signal.SIGKILL)
+                    except (ValueError, ProcessLookupError, PermissionError):
+                        pass
+        except (subprocess.TimeoutExpired, FileNotFoundError):
+            pass
+
+    def execute(self, code, auto_recover=True):
+        """
+        Execute code in PyMOL with optional auto-recovery.
+
+        Args:
+            code: Python code to execute in PyMOL
+            auto_recover: If True, attempt recovery on failure
+
+        Returns:
+            Output from PyMOL
+        """
+        try:
+            if not self.is_connected:
+                if auto_recover:
+                    self.recover()
+                else:
+                    raise ConnectionError("Not connected to PyMOL")
+
+            return self.connection.execute(code)
+
+        except (ConnectionError, TimeoutError):
+            if auto_recover:
+                self.recover()
+                return self.connection.execute(code)
+            raise
+
+    def __enter__(self):
+        """Context manager entry."""
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        self.stop()
+
+
+# Convenience: global session instance
+_session = None
+
+
+def get_session():
+    """Get or create the global PyMOL session."""
+    global _session
+    if _session is None:
+        _session = PyMOLSession()
+    return _session
+
+
+def ensure_running():
+    """Ensure PyMOL is running and connected."""
+    session = get_session()
+    if not session.is_healthy():
+        session.start()
+    return session
+
+
+def stop_pymol():
+    """Stop the global PyMOL session."""
+    global _session
+    if _session:
+        _session.stop()
+        _session = None

claudemol/view.py

@@ -0,0 +1,137 @@
+"""
+Visual feedback helper for PyMOL.
+
+Provides a simple way to execute commands and save/view the result.
+
+Usage:
+    from ai_mol.view import pymol_view
+
+    # Execute commands and save a snapshot
+    path = pymol_view("cmd.fetch('1ubq'); cmd.show('cartoon')", name="ubq_cartoon")
+
+    # Or with auto-naming
+    path = pymol_view("cmd.color('red', 'all')")
+"""
+
+import json
+import socket
+import time
+from datetime import datetime
+from pathlib import Path
+
+DEFAULT_HOST = "localhost"
+DEFAULT_PORT = 9880
+SCRATCH_DIR = Path.home() / ".ai-mol" / "scratch"
+
+
+def ensure_scratch_dir():
+    """Ensure the scratch directory exists."""
+    SCRATCH_DIR.mkdir(parents=True, exist_ok=True)
+    return SCRATCH_DIR
+
+
+def send_command(code: str, host: str = DEFAULT_HOST, port: int = DEFAULT_PORT, timeout: float = 120.0) -> dict:
+    """Send a command to PyMOL and return the result."""
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    s.settimeout(timeout)
+    try:
+        s.connect((host, port))
+        msg = json.dumps({"type": "execute", "code": code})
+        s.sendall(msg.encode())
+
+        response = b""
+        while True:
+            chunk = s.recv(4096)
+            if not chunk:
+                break
+            response += chunk
+            try:
+                return json.loads(response.decode())
+            except json.JSONDecodeError:
+                continue
+    finally:
+        s.close()
+    return {"status": "error", "error": "No response received"}
+
+
+def generate_filename(name: str | None = None, extension: str = "png") -> Path:
+    """Generate a unique filename in the scratch directory."""
+    ensure_scratch_dir()
+
+    if name:
+        # Clean the name
+        clean_name = "".join(c if c.isalnum() or c in "-_" else "_" for c in name)
+        return SCRATCH_DIR / f"{clean_name}.{extension}"
+    else:
+        # Auto-generate with timestamp
+        timestamp = datetime.now().strftime("%H%M%S")
+        return SCRATCH_DIR / f"view_{timestamp}.{extension}"
+
+
+def pymol_view(
+    commands: str,
+    name: str | None = None,
+    width: int = 800,
+    height: int = 600,
+    ray: bool = False,
+    port: int = DEFAULT_PORT,
+) -> str:
+    """
+    Execute PyMOL commands and save a snapshot.
+
+    Args:
+        commands: PyMOL Python commands to execute (e.g., "cmd.fetch('1ubq')")
+        name: Optional name for the output file (auto-generated if None)
+        width: Image width in pixels
+        height: Image height in pixels
+        ray: Whether to ray-trace (slower but prettier)
+        port: PyMOL socket port
+
+    Returns:
+        Path to the saved image file
+    """
+    output_path = generate_filename(name)
+
+    # Build the full command sequence
+    # IMPORTANT: cmd.png() with width/height parameters causes a PyMOL bug where
+    # the view matrix Z-distance becomes corrupted after multiple cycles of
+    # reinitialize + fetch + png. The Z-distance grows exponentially from ~120 to
+    # ~50000+ after just 3-4 cycles, causing the view to zoom out into invisibility.
+    #
+    # The fix: ALWAYS use cmd.ray(width, height) before cmd.png(path) WITHOUT
+    # dimensions in the png call. The ray() command renders to an offscreen buffer
+    # without touching the viewport, preventing the corruption.
+
+    render_cmd = f"cmd.ray({width}, {height})"
+    png_cmd = f'cmd.png(r"{output_path}")'
+
+    full_code = f"""
+{commands}
+
+# Render and save the image
+{render_cmd}
+{png_cmd}
+print(f"Saved: {output_path}")
+"""
+
+    result = send_command(full_code, port=port, timeout=120.0)
+
+    if result.get("status") == "success":
+        # Wait briefly for file to be written
+        time.sleep(0.2)
+        if output_path.exists():
+            return str(output_path)
+        else:
+            raise RuntimeError(f"Image was not saved to {output_path}")
+    else:
+        raise RuntimeError(f"PyMOL error: {result.get('error', 'Unknown error')}")
+
+
+def quick_view(port: int = DEFAULT_PORT) -> str:
+    """
+    Take a quick snapshot of the current PyMOL view.
+
+    Returns:
+        Path to the saved image
+    """
+    return pymol_view("pass", name=None, port=port)

claudemol-0.1.0.dist-info/METADATA

@@ -0,0 +1,158 @@
+Metadata-Version: 2.4
+Name: claudemol
+Version: 0.1.0
+Summary: PyMOL integration for Claude Code - control molecular visualization via natural language
+Project-URL: Homepage, https://github.com/ANaka/ai-mol
+Project-URL: Repository, https://github.com/ANaka/ai-mol
+Author: Alex Naka
+License: MIT
+License-File: LICENSE
+Keywords: ai,claude,molecular-visualization,pymol,structural-biology
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Classifier: Topic :: Scientific/Engineering :: Visualization
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+
+# ai-mol: Control PyMOL with Claude Code
+
+Control PyMOL through natural language using Claude Code. This integration enables conversational structural biology, molecular visualization, and analysis.
+
+https://github.com/user-attachments/assets/687f43dc-d45e-477e-ac2b-7438e175cb36
+
+## Features
+
+- **Natural language control**: Tell Claude what you want to visualize and it executes PyMOL commands
+- **Direct socket communication**: Claude Code talks directly to PyMOL (no intermediary server)
+- **Full PyMOL access**: Manipulate representations, colors, views, perform measurements, alignments, and more
+- **Skill-based workflows**: Built-in skills for common tasks like binding site visualization and publication figures
+
+## Architecture
+
+```
+Claude Code → TCP Socket (port 9880) → PyMOL Plugin → cmd.* execution
+```
+
+## Quick Start
+
+### Prerequisites
+
+- PyMOL installed on your system
+- [Claude Code](https://docs.anthropic.com/en/docs/claude-code) CLI installed
+- Python 3.10+
+
+### Installation
+
+1. **Clone the repository:**
+
+   ```bash
+   git clone https://github.com/ANaka/ai-mol
+   cd ai-mol
+   ```
+
+2. **Set up the PyMOL plugin:**
+
+   Add this line to your `~/.pymolrc` (create it if it doesn't exist):
+
+   ```python
+   run /path/to/ai-mol/claude_socket_plugin.py
+   ```
+
+   Replace `/path/to/ai-mol` with the actual path where you cloned the repository.
+
+3. **Start using it:**
+
+   Open Claude Code in the `ai-mol` directory and say:
+
+   > "Open PyMOL and load structure 1UBQ"
+
+   Claude will launch PyMOL (with the socket listener active) and load the structure.
+
+## Usage
+
+### Starting a Session
+
+Simply ask Claude to open PyMOL or load a structure:
+
+- "Open PyMOL"
+- "Load PDB 4HHB and show as cartoon"
+- "Fetch 1UBQ from the PDB"
+
+Claude will launch PyMOL if it's not already running.
+
+### Example Commands
+
+- "Color the protein by secondary structure"
+- "Show the binding site residues within 5Å of the ligand as sticks"
+- "Align these two structures and calculate RMSD"
+- "Create a publication-quality figure with ray tracing"
+- "Make a 360° rotation movie"
+
+### PyMOL Console Commands
+
+Check or control the socket listener from PyMOL's command line:
+
+```
+claude_status   # Check if listener is running
+claude_stop     # Stop the listener
+claude_start    # Start the listener
+```
+
+### Available Skills
+
+Claude Code has built-in skills for common workflows:
+
+- **pymol-fundamentals** - Basic visualization, selections, coloring
+- **protein-structure-basics** - Secondary structure, B-factor, representations
+- **binding-site-visualization** - Protein-ligand interactions
+- **structure-alignment-analysis** - Comparing and aligning structures
+- **antibody-visualization** - CDR loops, epitopes, Fab structures
+- **publication-figures** - High-quality figure export
+- **movie-creation** - Animations and rotations
+
+## Troubleshooting
+
+### Connection Issues
+
+- **"Could not connect to PyMOL"**: Make sure PyMOL is running and the plugin is loaded
+- **Check listener status**: Run `claude_status` in PyMOL's command line
+- **Restart listener**: Run `claude_stop` then `claude_start` in PyMOL
+
+### Plugin Not Loading
+
+- Verify the path in your `~/.pymolrc` is correct
+- Check PyMOL's output for any error messages on startup
+- Try running `run /path/to/claude_socket_plugin.py` manually in PyMOL
+
+### First-Time Setup Help
+
+Run the `/pymol-setup` skill in Claude Code for guided setup assistance.
+
+## Configuration
+
+The default socket port is **9880**. Both the plugin and Claude Code connection module use this port.
+
+Key files:
+- `claude_socket_plugin.py` - PyMOL plugin (headless, auto-loads via pymolrc)
+- `pymol_connection.py` - Python module for socket communication
+- `.claude/skills/` - Claude Code skills for PyMOL workflows
+
+## Limitations
+
+- PyMOL and Claude Code must run on the same machine (localhost connection)
+- One active connection at a time
+- Some complex multi-step operations may need guidance
+
+## Contributing
+
+Contributions welcome! This project aims to build comprehensive skills for Claude-PyMOL interaction. If you discover useful patterns or workflows, consider adding them as skills.
+
+## License
+
+MIT License - see LICENSE file for details.

claudemol-0.1.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+claudemol/__init__.py,sha256=RZDmQTql01kKApSDdMC3qy07Bz71CE1f6vyOJ_l7wHk,618
+claudemol/cli.py,sha256=J6cKHGbkZX_gYAcvMVlN3NXkIEJmMSjAeEL4R7GO-ww,5008
+claudemol/connection.py,sha256=PG-iqqJ8kYUrA9JQieigHNSzjdI-cab6c0SKUpZkxTg,7516
+claudemol/plugin.py,sha256=xRqpd9Ncn2vDZnmZsKHcFrhISJng38FAZv_0Thy-axc,5513
+claudemol/session.py,sha256=BtaIx9UalvnnlmQ1Zr9HFAlvt_0HzrAnT952HPpFG4w,8258
+claudemol/view.py,sha256=-qe4xu_U_AuVtZpgj_XK3GqnUgq_8ClraHpJjFd58lU,4097
+claudemol-0.1.0.dist-info/METADATA,sha256=kTIUkA5cu-dB3N4bC_qfLZidrEGXPEPAukAr6i7aU2o,5104
+claudemol-0.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+claudemol-0.1.0.dist-info/entry_points.txt,sha256=zHSqaPcqaXWoGEavppdq_zTrMARmjjrTqLHa5KEBchw,49
+claudemol-0.1.0.dist-info/licenses/LICENSE,sha256=sZWsSm99w9tlHsmCnA7_6KzEAqQaD7nSQlgYesvvWMQ,1075
+claudemol-0.1.0.dist-info/RECORD,,

claudemol-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any

claudemol-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+claudemol = claudemol.cli:main

claudemol-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Vishnu Rajan Tejus
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.