#!/usr/bin/env python3
import re
import os
import socket
import json
import logging
import time
from pathlib import Path
from contextlib import asynccontextmanager
from typing import Optional, Dict, Any, AsyncIterator
from mcp.server.fastmcp import FastMCP, Context
##############################################################################
# SCREENSHOT CONFIGURATION
##############################################################################
# Use home directory instead of /tmp to avoid snap sandboxing issues
SCREENSHOT_DIR = Path.home() / "pymol_screenshots"
SCREENSHOT_DIR.mkdir(exist_ok=True)
_screenshot_counter = 0
_max_screenshots = 20 # Default, user-configurable
def get_screenshot_path() -> Path:
"""Get next screenshot path, auto-cleanup old ones if over limit."""
global _screenshot_counter
_screenshot_counter += 1
# Auto-cleanup: delete oldest if over limit
existing = sorted(SCREENSHOT_DIR.glob("view_*.png"))
while len(existing) >= _max_screenshots:
try:
existing[0].unlink()
except OSError:
pass
existing.pop(0)
return SCREENSHOT_DIR / f"view_{_screenshot_counter:04d}.png"
##############################################################################
# PYMOL COMMAND DEFINITIONS AND ERROR PATTERNS (Provided by user)
##############################################################################
PYMOL_COMMANDS = {
# MOLECULAR VISUALIZATION
"show": {
"description": "Shows a representation for the specified selection",
"pattern": r"^show\s+([\w.]+)(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "representation", "required": True, "options": [
"lines", "sticks", "spheres", "surface", "mesh", "dots",
"ribbon", "cartoon", "labels", "nonbonded", "nb_spheres",
"ellipsoids", "volume", "slice", "extent", "dots_as_spheres",
"cell", "cgo", "everything", "dashes", "angles", "dihedrals",
"licorice", "spheres", "putty"
]},
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"hide": {
"description": "Hides a representation for the specified selection",
"pattern": r"^hide\s+([\w.]+)(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "representation", "required": True, "options": [
"lines", "sticks", "spheres", "surface", "mesh", "dots",
"ribbon", "cartoon", "labels", "nonbonded", "nb_spheres",
"ellipsoids", "volume", "slice", "extent", "dots_as_spheres",
"cell", "cgo", "everything", "dashes", "angles", "dihedrals",
"licorice", "spheres", "putty"
]},
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"color": {
"description": "Sets the color for the specified selection",
"pattern": r"^color\s+([\w.]+)(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "color", "required": True},
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"as": {
"description": "Shows one representation while hiding all others for the specified selection",
"pattern": r"^as\s+([\w.]+)(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "representation", "required": True},
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"set": {
"description": "Sets a PyMOL setting to a specified value",
"pattern": r"^set\s+([\w.]+)(?:\s*,\s*([^,]+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "setting", "required": True},
{"name": "value", "required": True},
{"name": "selection", "required": False}
],
"check_selection": False
},
"cartoon": {
"description": "Sets the cartoon type for the specified selection",
"pattern": r"^cartoon\s+([\w.]+)(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "type", "required": True, "options": [
"automatic", "loop", "rectangle", "oval", "tube", "arrow", "dumbbell", "putty"
]},
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"spectrum": {
"description": "Colors selection in a spectrum",
"pattern": r"^spectrum\s+([^,]+)(?:\s*,\s*([^,]+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "expression", "required": True},
{"name": "palette", "required": False, "default": "rainbow"},
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"label": {
"description": "Adds labels to atoms in the selection",
"pattern": r"^label\s+([^,]+)(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "selection", "required": True},
{"name": "expression", "required": False, "default": "name"}
],
"check_selection": True
},
"distance": {
"description": "Measures the distance between two selections",
"pattern": r"^distance(?:\s+([^,]+))?(?:\s*,\s*([^,]+))?(?:\s*,\s*([^,]+))?$",
"parameters": [
{"name": "name", "required": False},
{"name": "selection1", "required": False, "default": "(pk1)"},
{"name": "selection2", "required": False, "default": "(pk2)"}
],
"check_selection": True
},
"angle": {
"description": "Measures the angle between three selections",
"pattern": r"^angle(?:\s+([^,]+))?(?:\s*,\s*([^,]+))?(?:\s*,\s*([^,]+))?(?:\s*,\s*([^,]+))?$",
"parameters": [
{"name": "name", "required": False},
{"name": "selection1", "required": False, "default": "(pk1)"},
{"name": "selection2", "required": False, "default": "(pk2)"},
{"name": "selection3", "required": False, "default": "(pk3)"}
],
"check_selection": True
},
"dihedral": {
"description": "Measures the dihedral angle between four selections",
"pattern": r"^dihedral(?:\s+([^,]+))?(?:\s*,\s*([^,]+))?(?:\s*,\s*([^,]+))?(?:\s*,\s*([^,]+))?(?:\s*,\s*([^,]+))?$",
"parameters": [
{"name": "name", "required": False},
{"name": "selection1", "required": False, "default": "(pk1)"},
{"name": "selection2", "required": False, "default": "(pk2)"},
{"name": "selection3", "required": False, "default": "(pk3)"},
{"name": "selection4", "required": False, "default": "(pk4)"}
],
"check_selection": True
},
# VIEWING OPERATIONS
"center": {
"description": "Centers the view on a selection",
"pattern": r"^center(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"orient": {
"description": "Orients the view to align with principal axes of the selection",
"pattern": r"^orient(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"zoom": {
"description": "Zooms the view on a selection",
"pattern": r"^zoom(?:\s+([^,]+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"},
{"name": "buffer", "required": False, "default": "5"}
],
"check_selection": True
},
"reset": {
"description": "Resets the view, optionally resetting an object's matrix",
"pattern": r"^reset(?:\s+(.+))?$",
"parameters": [
{"name": "object", "required": False}
],
"check_selection": False
},
"turn": {
"description": "Rotates the camera around an axis",
"pattern": r"^turn\s+([xyz])(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "axis", "required": True, "options": ["x", "y", "z"]},
{"name": "angle", "required": False, "default": "90"}
],
"check_selection": False
},
"move": {
"description": "Moves the camera along an axis",
"pattern": r"^move\s+([xyz])(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "axis", "required": True, "options": ["x", "y", "z"]},
{"name": "distance", "required": False, "default": "1"}
],
"check_selection": False
},
"clip": {
"description": "Adjusts the clipping planes",
"pattern": r"^clip\s+([\w.]+)(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "mode", "required": True, "options": ["near", "far", "slab", "atoms", "near_slab", "far_slab"]},
{"name": "distance", "required": False, "default": "1"}
],
"check_selection": False
},
# FILE OPERATIONS
"load": {
"description": "Loads a file into PyMOL",
"pattern": r"^load\s+([^,]+)(?:\s*,\s*([^,]+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "filename", "required": True},
{"name": "object", "required": False},
{"name": "options", "required": False}
],
"check_selection": False
},
"fetch": {
"description": "Fetches a structure from a database (e.g., PDB)",
"pattern": r"^fetch\s+([^,]+)(?:\s*,\s*([^,]+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "code", "required": True},
{"name": "name", "required": False},
{"name": "options", "required": False}
],
"check_selection": False
},
"save": {
"description": "Saves data to a file",
"pattern": r"^save\s+([^,]+)(?:\s*,\s*(.+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "filename", "required": True},
{"name": "selection", "required": False, "default": "all"},
{"name": "state", "required": False, "default": "-1"}
],
"check_selection": True
},
"png": {
"description": "Saves a PNG image",
"pattern": r"^png\s+([^,]+)(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "filename", "required": True},
{"name": "options", "required": False}
],
"check_selection": False
},
# SELECTION OPERATIONS
"select": {
"description": "Creates a named selection",
"pattern": r"^select\s+([^,]+)(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "name", "required": True},
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": False
},
"deselect": {
"description": "Clears the current selection",
"pattern": r"^deselect$",
"parameters": [],
"check_selection": False
},
# OBJECT MANIPULATION
"create": {
"description": "Creates a new object from a selection",
"pattern": r"^create\s+([^,]+)(?:\s*,\s*(.+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "name", "required": True},
{"name": "selection", "required": False, "default": "all"},
{"name": "source_state", "required": False, "default": "1"}
],
"check_selection": True
},
"extract": {
"description": "Extracts a selection to a new object",
"pattern": r"^extract\s+([^,]+)(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "name", "required": True},
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"delete": {
"description": "Deletes objects or selections",
"pattern": r"^delete\s+(.+)$",
"parameters": [
{"name": "name", "required": True}
],
"check_selection": False
},
"remove": {
"description": "Removes atoms in a selection",
"pattern": r"^remove\s+(.+)$",
"parameters": [
{"name": "selection", "required": True}
],
"check_selection": True
},
"align": {
"description": "Aligns one selection to another",
"pattern": r"^align\s+([^,]+)(?:\s*,\s*([^,]+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "mobile", "required": True},
{"name": "target", "required": False, "default": "all"},
{"name": "options", "required": False}
],
"check_selection": True
},
"super": {
"description": "Superimposes one selection onto another",
"pattern": r"^super\s+([^,]+)(?:\s*,\s*([^,]+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "mobile", "required": True},
{"name": "target", "required": False, "default": "all"},
{"name": "options", "required": False}
],
"check_selection": True
},
"intra_fit": {
"description": "Fits all states within an object",
"pattern": r"^intra_fit\s+(.+)$",
"parameters": [
{"name": "selection", "required": True}
],
"check_selection": True
},
"intra_rms": {
"description": "Calculates RMSD between states within an object",
"pattern": r"^intra_rms\s+(.+)$",
"parameters": [
{"name": "selection", "required": True}
],
"check_selection": True
},
# UTILITY AND MODIFICATION
"alter": {
"description": "Alters atomic properties in a selection",
"pattern": r"^alter\s+([^,]+)(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "selection", "required": True},
{"name": "expression", "required": True}
],
"check_selection": True
},
"alter_state": {
"description": "Alters atomic coordinates in a state",
"pattern": r"^alter_state\s+([^,]+)(?:\s*,\s*([^,]+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "state", "required": True},
{"name": "selection", "required": True},
{"name": "expression", "required": True}
],
"check_selection": True
},
"h_add": {
"description": "Adds hydrogens to a selection",
"pattern": r"^h_add(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"h_fill": {
"description": "Adds hydrogens and adjusts valences",
"pattern": r"^h_fill(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"bond": {
"description": "Creates a bond between two atoms",
"pattern": r"^bond\s+([^,]+)(?:\s*,\s*([^,]+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "atom1", "required": True},
{"name": "atom2", "required": True},
{"name": "order", "required": False, "default": "1"}
],
"check_selection": True
},
"unbond": {
"description": "Removes a bond between two atoms",
"pattern": r"^unbond\s+([^,]+)(?:\s*,\s*([^,]+))?$",
"parameters": [
{"name": "atom1", "required": True},
{"name": "atom2", "required": True}
],
"check_selection": True
},
"rebuild": {
"description": "Regenerates all displayed geometry",
"pattern": r"^rebuild(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": False
},
"refresh": {
"description": "Refreshes the display",
"pattern": r"^refresh$",
"parameters": [],
"check_selection": False
},
# UTILITY FUNCTIONS
"util.cbc": {
"description": "Colors by chain (Color By Chain)",
"pattern": r"^util\.cbc(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"util.cbaw": {
"description": "Colors by atom, white carbons (Color By Atom, White)",
"pattern": r"^util\.cbaw(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"util.cbag": {
"description": "Colors by atom, green carbons (Color By Atom, Green)",
"pattern": r"^util\.cbag(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"util.cbac": {
"description": "Colors by atom, cyan carbons (Color By Atom, Cyan)",
"pattern": r"^util\.cbac(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"util.cbam": {
"description": "Colors by atom, magenta carbons (Color By Atom, Magenta)",
"pattern": r"^util\.cbam(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"util.cbay": {
"description": "Colors by atom, yellow carbons (Color By Atom, Yellow)",
"pattern": r"^util\.cbay(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"util.cbas": {
"description": "Colors by atom, salmon carbons (Color By Atom, Salmon)",
"pattern": r"^util\.cbas(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"util.cbab": {
"description": "Colors by atom, slate carbons (Color By Atom, slateBLue)",
"pattern": r"^util\.cbab(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"util.cbao": {
"description": "Colors by atom, orange carbons (Color By Atom, Orange)",
"pattern": r"^util\.cbao(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"util.cbap": {
"description": "Colors by atom, purple carbons (Color By Atom, Purple)",
"pattern": r"^util\.cbap(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"util.cbak": {
"description": "Colors by atom, pink carbons (Color By Atom, pinK)",
"pattern": r"^util\.cbak(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"util.chainbow": {
"description": "Colors chains in rainbow gradient (CHAINs in rainBOW)",
"pattern": r"^util\.chainbow(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"util.rainbow": {
"description": "Colors residues in rainbow from N to C terminus",
"pattern": r"^util\.rainbow(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"util.ss": {
"description": "Colors by secondary structure",
"pattern": r"^util\.ss(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"util.color_by_element": {
"description": "Colors atoms by their element",
"pattern": r"^util\.color_by_element(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"util.color_secondary": {
"description": "Colors secondary structure elements",
"pattern": r"^util\.color_secondary(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
# MOLECULAR DYNAMICS AND ANALYSIS
"spheroid": {
"description": "Displays atoms as smooth spheres",
"pattern": r"^spheroid(?:\s+(.+))?$",
"parameters": [
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"isomesh": {
"description": "Creates a mesh isosurface",
"pattern": r"^isomesh\s+([^,]+)(?:\s*,\s*([^,]+))?(?:\s*,\s*([^,]+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "name", "required": True},
{"name": "map_object", "required": True},
{"name": "level", "required": True},
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"isosurface": {
"description": "Creates a solid isosurface",
"pattern": r"^isosurface\s+([^,]+)(?:\s*,\s*([^,]+))?(?:\s*,\s*([^,]+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "name", "required": True},
{"name": "map_object", "required": True},
{"name": "level", "required": True},
{"name": "selection", "required": False, "default": "all"}
],
"check_selection": True
},
"sculpt_activate": {
"description": "Activates sculpting mode for an object",
"pattern": r"^sculpt_activate\s+(.+)$",
"parameters": [
{"name": "object", "required": True}
],
"check_selection": False
},
"sculpt_deactivate": {
"description": "Deactivates sculpting mode for an object",
"pattern": r"^sculpt_deactivate\s+(.+)$",
"parameters": [
{"name": "object", "required": True}
],
"check_selection": False
},
"sculpt_iterate": {
"description": "Performs sculpting iterations",
"pattern": r"^sculpt_iterate\s+([^,]+)(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "iterations", "required": True},
{"name": "object", "required": False, "default": "all"}
],
"check_selection": False
},
# SCENES AND MOVIES
"scene": {
"description": "Manages scenes for later recall",
"pattern": r"^scene\s+([^,]+)(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "key", "required": True},
{"name": "action", "required": False, "default": "recall"}
],
"check_selection": False
},
"scene_order": {
"description": "Sets the order of scenes",
"pattern": r"^scene_order\s+(.+)$",
"parameters": [
{"name": "scene_list", "required": True}
],
"check_selection": False
},
"mset": {
"description": "Defines a sequence of states for movie playback",
"pattern": r"^mset\s+(.+)$",
"parameters": [
{"name": "specification", "required": True}
],
"check_selection": False
},
"mplay": {
"description": "Starts playing the movie",
"pattern": r"^mplay$",
"parameters": [],
"check_selection": False
},
"mstop": {
"description": "Stops the movie",
"pattern": r"^mstop$",
"parameters": [],
"check_selection": False
},
"frame": {
"description": "Sets or queries the current frame",
"pattern": r"^frame(?:\s+(.+))?$",
"parameters": [
{"name": "frame_number", "required": False}
],
"check_selection": False
},
"forward": {
"description": "Advances one frame",
"pattern": r"^forward$",
"parameters": [],
"check_selection": False
},
"backward": {
"description": "Goes back one frame",
"pattern": r"^backward$",
"parameters": [],
"check_selection": False
},
"rock": {
"description": "Toggles a rocking animation",
"pattern": r"^rock$",
"parameters": [],
"check_selection": False
},
# RENDERING
"ray": {
"description": "Performs ray-tracing",
"pattern": r"^ray(?:\s+([^,]+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "width", "required": False},
{"name": "height", "required": False}
],
"check_selection": False
},
"draw": {
"description": "Uses OpenGL renderer (faster but lower quality)",
"pattern": r"^draw(?:\s+([^,]+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "width", "required": False},
{"name": "height", "required": False}
],
"check_selection": False
},
"mpng": {
"description": "Saves a series of PNG images for movie frames",
"pattern": r"^mpng\s+(.+)$",
"parameters": [
{"name": "prefix", "required": True}
],
"check_selection": False
},
# CRYSTALLOGRAPHY
"symexp": {
"description": "Generates symmetry-related copies",
"pattern": r"^symexp\s+([^,]+)(?:\s*,\s*([^,]+))?(?:\s*,\s*([^,]+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "prefix", "required": True},
{"name": "selection", "required": True},
{"name": "cutoff", "required": False, "default": "20"},
{"name": "segi", "required": False}
],
"check_selection": True
},
"symexp": {
"description": "Generates symmetry-related copies",
"pattern": r"^symexp\s+([^,]+)(?:\s*,\s*([^,]+))?(?:\s*,\s*([^,]+))?(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "prefix", "required": True},
{"name": "selection", "required": True},
{"name": "cutoff", "required": False, "default": "20"},
{"name": "segi", "required": False}
],
"check_selection": True
},
"set_symmetry": {
"description": "Sets symmetry parameters for an object",
"pattern": r"^set_symmetry\s+([^,]+)(?:\s*,\s*([^,]+))?(?:\s*,\s*([^,]+))?(?:\s*,\s*([^,]+))?(?:\s*,\s*([^,]+))?(?:\s*,\s*([^,]+))?(?:\s*,\s*([^,]+))?$",
"parameters": [
{"name": "selection", "required": True},
{"name": "a", "required": True},
{"name": "b", "required": True},
{"name": "c", "required": True},
{"name": "alpha", "required": True},
{"name": "beta", "required": True},
{"name": "gamma", "required": True}
],
"check_selection": True
},
# OTHER
"fab": {
"description": "Creates a peptide chain from a sequence",
"pattern": r"^fab\s+([^,]+)(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "sequence", "required": True},
{"name": "options", "required": False}
],
"check_selection": False
},
"fragment": {
"description": "Loads a molecular fragment",
"pattern": r"^fragment\s+(.+)$",
"parameters": [
{"name": "name", "required": True}
],
"check_selection": False
},
"full_screen": {
"description": "Toggles fullscreen mode",
"pattern": r"^full_screen$",
"parameters": [],
"check_selection": False
},
"viewport": {
"description": "Sets the viewport size",
"pattern": r"^viewport\s+([^,]+)(?:\s*,\s*(.+))?$",
"parameters": [
{"name": "width", "required": True},
{"name": "height", "required": True}
],
"check_selection": False
},
"cd": {
"description": "Changes the current directory",
"pattern": r"^cd\s+(.+)$",
"parameters": [
{"name": "path", "required": True}
],
"check_selection": False
},
"pwd": {
"description": "Prints the current directory",
"pattern": r"^pwd$",
"parameters": [],
"check_selection": False
},
"ls": {
"description": "Lists files in the current directory",
"pattern": r"^ls(?:\s+(.+))?$",
"parameters": [
{"name": "path", "required": False}
],
"check_selection": False
},
"system": {
"description": "Executes a system command",
"pattern": r"^system\s+(.+)$",
"parameters": [
{"name": "command", "required": True}
],
"check_selection": False
},
"help": {
"description": "Shows help for a command",
"pattern": r"^help(?:\s+(.+))?$",
"parameters": [
{"name": "command", "required": False}
],
"check_selection": False
}
}
ERROR_PATTERNS = {
"SYNTAX_ERROR": [
r"Syntax error",
r"invalid syntax",
r"Unknown command"
],
"SELECTION_ERROR": [
r"Invalid selection",
r"No atoms selected",
r"Selection not found",
r"Selection \S+ doesn't exist"
],
"OBJECT_NOT_FOUND": [
r"object \S+ not found",
r"Object \S+ does not exist",
r"Unable to find object named \S+"
],
"ATOM_NOT_FOUND": [
r"No atoms matched",
r"No atoms in selection",
r"Atom not found"
],
"FILE_ERROR": [
r"Unable to open file",
r"No such file",
r"Permission denied",
r"Error reading file",
r"Error writing file"
],
"CONNECTION_ERROR": [
r"Connection refused",
r"Network error",
r"Timeout",
r"Failed to fetch"
],
"PARAMETER_ERROR": [
r"Incorrect number of parameters",
r"Invalid parameter",
r"Parameter out of range"
]
}
##############################################################################
# LOGGING
##############################################################################
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("PyMOLMCPServer")
##############################################################################
# PYMOL SOCKET CONNECTION
##############################################################################
class PyMOLConnection:
def __init__(self, host: str = 'localhost', port: int = 9876):
self.host = host
self.port = port
self.sock: Optional[socket.socket] = None
def connect(self) -> bool:
if self.sock:
return True
try:
self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.sock.connect((self.host, self.port))
logger.info(f"Connected to PyMOL at {self.host}:{self.port}")
return True
except Exception as e:
logger.error(f"Connection error: {e}")
self.sock = None
return False
def disconnect(self) -> None:
if self.sock:
try:
self.sock.close()
except Exception as e:
logger.error(f"Disconnect error: {e}")
finally:
self.sock = None
def send_command(self, code: str) -> Dict[str, Any]:
"""
Sends Python code to PyMOL via the socket plugin and returns a JSON response:
{ "status": "success" or "error",
"result": {
"executed": bool,
"output": str or None,
"error": str or None
},
"message": "error message string if any" }
"""
if not self.sock and not self.connect():
raise ConnectionError("Not connected to PyMOL")
data = {"type": "pymol_command", "code": code}
try:
self.sock.sendall(json.dumps(data).encode('utf-8'))
self.sock.settimeout(10.0)
chunks = []
while True:
chunk = self.sock.recv(4096)
if not chunk:
break
chunks.append(chunk)
buffer = b''.join(chunks)
try:
response = json.loads(buffer.decode('utf-8'))
return response
except json.JSONDecodeError:
continue
if chunks:
buffer = b''.join(chunks)
return json.loads(buffer.decode('utf-8'))
raise ConnectionError("No response from PyMOL")
except socket.timeout:
self.sock = None
raise TimeoutError("PyMOL response timed out")
except Exception as e:
self.sock = None
raise RuntimeError(f"PyMOL command error: {e}")
_global_connection: Optional['PyMOLConnection'] = None
def get_pymol_connection() -> PyMOLConnection:
global _global_connection
if _global_connection is not None:
try:
# Test if connection is alive
_global_connection.send_command("pass")
return _global_connection
except:
try:
_global_connection.disconnect()
except:
pass
_global_connection = None
if _global_connection is None:
conn = PyMOLConnection()
if not conn.connect():
raise RuntimeError("Could not connect to PyMOL socket.")
_global_connection = conn
return _global_connection
##############################################################################
# PARSING USER INPUT TO PYMOL COMMANDS
##############################################################################
def parse_pymol_input(input_text: str) -> str:
"""
Attempts to match the user input against known PYMOL_COMMANDS patterns.
If matched, extracts parameters and builds the final Python code for PyMOL.
Raises ValueError if no command matches or if there's a parameter error.
"""
text_stripped = input_text.strip()
for cmd_name, cmd_info in PYMOL_COMMANDS.items():
pattern = re.compile(cmd_info["pattern"], re.IGNORECASE)
match = pattern.match(text_stripped)
if match:
groups = match.groups()
# Extract parameter definitions
params_def = cmd_info["parameters"]
param_values = {}
for i, param_def in enumerate(params_def):
param_name = param_def["name"]
required = param_def.get("required", False)
default_val = param_def.get("default", None)
options = param_def.get("options", [])
# Attempt to fetch from match group
value = None
if i < len(groups) and groups[i] is not None:
value = groups[i].strip()
elif required and default_val is None:
raise ValueError(f"Missing required parameter '{param_name}' for command {cmd_name}")
elif value is None and default_val is not None:
value = default_val
if options and value and value not in options:
raise ValueError(f"Parameter '{param_name}' must be one of {options}")
param_values[param_name] = value if value is not None else ""
# (Optional) If check_selection is True, we could do extra checks
# But for simplicity, just build PyMOL code
return build_pymol_code(cmd_name, param_values)
raise ValueError("No recognized PyMOL command pattern matched this input.")
def build_pymol_code(command_name: str, param_values: Dict[str, Any]) -> str:
"""
Translates a recognized command plus parameters into Python code for PyMOL.
This is a naive approach that constructs a single cmd.* invocation.
Modify as needed for more complex logic.
"""
# Example approach: "cmd.show('sticks', 'sele')"
# For 'show' -> "cmd.show('sticks', 'all')"
# This is simplified. Real approach might be more advanced.
if command_name == "help":
# We can do a special return for help
cmd_obj = param_values.get("command") or ""
if cmd_obj and cmd_obj in PYMOL_COMMANDS:
return f"print('Help for {cmd_obj}: {PYMOL_COMMANDS[cmd_obj]['description']}')"
return "print('List of PyMOL commands...')"
# Generic pattern (the user can adapt this to each command's syntax)
py_code = []
py_code.append("from pymol import cmd")
if command_name in ["util.cbc", "util.cbaw", "util.cbag", "util.cbac", "util.cbam",
"util.cbay", "util.cbas", "util.cbab", "util.cbao", "util.cbap",
"util.cbak", "util.chainbow", "util.rainbow", "util.ss",
"util.color_by_element", "util.color_secondary"]:
# handle util.* style calls
# e.g. "import util" doesn't exist in PyMOL by default. It's cmd.util.* typically
selection = param_values.get("selection","all")
# e.g. "cmd.util.chainbow('all')"
# but realistically "util.chainbow(...)" might be "cmd.do('util.chainbow all')"
call_code = f"cmd.do('{command_name} {selection}')"
py_code.append(call_code)
return "; ".join(py_code)
# For typical direct 'cmd' calls:
# We'll do a simple switch
if command_name == "show":
representation = param_values["representation"]
selection = param_values["selection"]
py_code.append(f"cmd.show('{representation}', '{selection}')")
elif command_name == "hide":
representation = param_values["representation"]
selection = param_values["selection"]
py_code.append(f"cmd.hide('{representation}', '{selection}')")
elif command_name == "color":
color_val = param_values["color"]
selection = param_values["selection"]
py_code.append(f"cmd.color('{color_val}', '{selection}')")
elif command_name == "turn":
axis = param_values["axis"]
angle = param_values["angle"]
py_code.append(f"cmd.turn('{axis}', {angle})")
elif command_name == "move":
axis = param_values["axis"]
distance = param_values["distance"]
py_code.append(f"cmd.move('{axis}', {distance})")
elif command_name == "png":
filename = param_values["filename"]
options = param_values.get("options", "")
if options.strip():
py_code.append(f"cmd.png('{filename}', {options})")
else:
py_code.append(f"cmd.png('{filename}')")
else:
# fallback, naive approach: "cmd.do('original command')"
# build the original command as a string
raw_cmd = command_name
# We skip the prefix if it's something like "util."
for k,v in param_values.items():
if v.strip():
raw_cmd += f" {v}"
py_code.append(f"cmd.do('{raw_cmd}')")
return "; ".join(py_code)
def analyze_pymol_output(output_text: str) -> Optional[str]:
"""
Attempts to map known error patterns in the PyMOL output to a user-friendly error.
Returns None if no known error patterns are matched.
"""
lower_out = output_text.lower()
for error_label, patterns in ERROR_PATTERNS.items():
for p in patterns:
if re.search(p.lower(), lower_out):
return f"{error_label} detected: {p}"
return None
##############################################################################
# MCP SERVER SETUP
##############################################################################
@asynccontextmanager
async def server_lifespan(server: FastMCP) -> AsyncIterator[dict]:
try:
logger.info("Starting PyMOL MCP server (with command parsing).")
try:
get_pymol_connection()
except Exception as e:
logger.warning(f"Initial PyMOL connection failure: {e}")
yield {}
finally:
global _global_connection
if _global_connection:
_global_connection.disconnect()
_global_connection = None
logger.info("PyMOL MCP server shut down.")
mcp = FastMCP("PyMOLMCPServer", lifespan=server_lifespan)
##############################################################################
# MCP TOOL: parse_and_execute
##############################################################################
def generate_natural_response(user_input: str) -> str:
"""
Generate a natural language description of what the PyMOL command did.
"""
text = user_input.strip().lower()
parts = text.split()
cmd_name = parts[0] if parts else ""
# Parse the command for natural language generation
if cmd_name == "fetch":
code = parts[1].upper() if len(parts) > 1 else "structure"
return f"Fetched {code} from the PDB"
elif cmd_name == "load":
filename = parts[1] if len(parts) > 1 else "file"
return f"Loaded {filename}"
elif cmd_name == "show":
rep = parts[1] if len(parts) > 1 else "representation"
sel = parts[2] if len(parts) > 2 else "all"
return f"Showing {rep} for {sel}"
elif cmd_name == "hide":
rep = parts[1] if len(parts) > 1 else "representation"
sel = parts[2] if len(parts) > 2 else "all"
return f"Hid {rep} for {sel}"
elif cmd_name == "color":
color = parts[1] if len(parts) > 1 else "color"
sel = parts[2] if len(parts) > 2 else "all"
return f"Colored {sel} {color}"
elif cmd_name == "as":
rep = parts[1] if len(parts) > 1 else "representation"
sel = parts[2] if len(parts) > 2 else "all"
return f"Set {sel} to {rep} representation"
elif cmd_name == "select":
name = parts[1] if len(parts) > 1 else "selection"
return f"Created selection '{name}'"
elif cmd_name == "zoom":
sel = parts[1] if len(parts) > 1 else "all"
return f"Zoomed to {sel}"
elif cmd_name == "center":
sel = parts[1] if len(parts) > 1 else "all"
return f"Centered view on {sel}"
elif cmd_name == "orient":
sel = parts[1] if len(parts) > 1 else "all"
return f"Oriented view to {sel}"
elif cmd_name == "reset":
return "Reset the view"
elif cmd_name == "delete":
obj = parts[1] if len(parts) > 1 else "object"
return f"Deleted {obj}"
elif cmd_name == "remove":
sel = parts[1] if len(parts) > 1 else "selection"
return f"Removed atoms in {sel}"
elif cmd_name == "cartoon":
style = parts[1] if len(parts) > 1 else "style"
return f"Set cartoon style to {style}"
elif cmd_name == "spectrum":
return "Applied spectrum coloring"
elif cmd_name == "distance":
return "Measured distance"
elif cmd_name == "angle":
return "Measured angle"
elif cmd_name == "align":
mobile = parts[1] if len(parts) > 1 else "mobile"
target = parts[2] if len(parts) > 2 else "target"
return f"Aligned {mobile} to {target}"
elif cmd_name == "super":
mobile = parts[1] if len(parts) > 1 else "mobile"
target = parts[2] if len(parts) > 2 else "target"
return f"Superimposed {mobile} onto {target}"
elif cmd_name == "set":
setting = parts[1] if len(parts) > 1 else "setting"
value = parts[2] if len(parts) > 2 else "value"
return f"Set {setting} to {value}"
elif cmd_name == "turn":
axis = parts[1] if len(parts) > 1 else "axis"
angle = parts[2] if len(parts) > 2 else "90"
return f"Rotated view {angle}° around {axis}-axis"
elif cmd_name == "move":
axis = parts[1] if len(parts) > 1 else "axis"
dist = parts[2] if len(parts) > 2 else "1"
return f"Moved camera {dist} along {axis}-axis"
elif cmd_name == "ray":
return "Ray-traced the scene"
elif cmd_name == "png":
filename = parts[1] if len(parts) > 1 else "image"
return f"Saved image to {filename}"
elif cmd_name == "save":
filename = parts[1] if len(parts) > 1 else "file"
return f"Saved to {filename}"
elif cmd_name.startswith("util."):
return f"Applied {cmd_name} coloring"
else:
return f"Executed: {user_input}"
@mcp.tool()
def parse_and_execute(ctx: Context, user_input: str) -> str:
"""
Parses a text command against PYMOL_COMMANDS, builds PyMOL code,
executes it, and analyzes any error patterns in the output.
"""
try:
code = parse_pymol_input(user_input)
except ValueError as ve:
return f"Could not parse command: {ve}"
except Exception as e:
return f"Parsing error: {e}"
try:
conn = get_pymol_connection()
response = conn.send_command(code)
status = response.get("status", "error")
if status == "success":
res = response.get("result", {})
out = res.get("output","") if isinstance(res, dict) else ""
# If there's output, check known error patterns
check_err = analyze_pymol_output(out)
if check_err:
return f"Warning: {check_err}"
# Return natural language description
return generate_natural_response(user_input)
else:
msg = response.get("message","Unknown error")
check_err = analyze_pymol_output(msg)
if check_err:
return f"Failed: {check_err}"
return f"Error: {msg}"
except Exception as e:
return f"Connection error: {e}"
##############################################################################
# SCREENSHOT TOOLS
##############################################################################
@mcp.tool()
def take_screenshot(ctx: Context, width: int = 200, height: int = 200) -> str:
"""
Capture screenshot of current PyMOL view at specified resolution.
Args:
width: Image width in pixels (default 200)
height: Image height in pixels (default 200)
Returns:
Path to saved PNG file. Use Claude Code's Read tool to view it.
"""
try:
conn = get_pymol_connection()
except Exception as e:
return f"Connection error: {e}"
output_path = get_screenshot_path()
# Ensure directory exists before PyMOL tries to write
SCREENSHOT_DIR.mkdir(parents=True, exist_ok=True)
# Delete any existing file to avoid returning stale screenshots
if output_path.exists():
try:
output_path.unlink()
except OSError:
pass
# Record time before screenshot to verify file is newly created
time_before = time.time()
# Use draw to render at specific resolution without resizing viewport
code = f"from pymol import cmd; cmd.draw({width}, {height}); cmd.png('{output_path}')"
try:
response = conn.send_command(code)
# Check response status
status = response.get("status", "unknown")
if status != "success":
return f"PyMOL error: {response.get('message', response)}"
# Give PyMOL time to write the file
for _ in range(10):
time.sleep(0.2)
if output_path.exists() and output_path.stat().st_size > 0:
# Verify file was created after we sent the command
if output_path.stat().st_mtime >= time_before:
break
if output_path.exists() and output_path.stat().st_size > 0:
if output_path.stat().st_mtime >= time_before:
return str(output_path)
else:
return f"Failed: screenshot file was not updated (stale file detected)"
else:
return f"Failed: file not created. Path: {output_path}, Response: {response}"
except Exception as e:
return f"Screenshot failed: {e}"
@mcp.tool()
def take_all_views(ctx: Context, width: int = 200, height: int = 200) -> str:
"""
Capture screenshots from all 6 orthogonal views (90-degree angles).
Views captured: front, back, left, right, top, bottom.
Args:
width: Image width in pixels (default 200)
height: Image height in pixels (default 200)
Returns:
Paths to all saved PNG files.
"""
from PIL import Image
try:
conn = get_pymol_connection()
except Exception as e:
return f"Connection error: {e}"
SCREENSHOT_DIR.mkdir(parents=True, exist_ok=True)
# Define the 6 orthogonal views with their rotations from reset position
views = [
("front", "cmd.reset(); cmd.orient()"),
("back", "cmd.reset(); cmd.orient(); cmd.turn('y', 180)"),
("right", "cmd.reset(); cmd.orient(); cmd.turn('y', -90)"),
("left", "cmd.reset(); cmd.orient(); cmd.turn('y', 90)"),
("top", "cmd.reset(); cmd.orient(); cmd.turn('x', -90)"),
("bottom", "cmd.reset(); cmd.orient(); cmd.turn('x', 90)"),
]
saved_paths = []
images = []
for view_name, rotation_code in views:
output_path = SCREENSHOT_DIR / f"view_{view_name}.png"
# Delete existing file
if output_path.exists():
try:
output_path.unlink()
except OSError:
pass
time_before = time.time()
# Rotate to view, render, and save
code = f"from pymol import cmd; {rotation_code}; cmd.draw({width}, {height}); cmd.png('{output_path}')"
try:
response = conn.send_command(code)
# Wait for file to be written
for _ in range(10):
time.sleep(0.2)
if output_path.exists() and output_path.stat().st_size > 0:
if output_path.stat().st_mtime >= time_before:
break
if output_path.exists() and output_path.stat().st_size > 0:
saved_paths.append(str(output_path))
images.append((view_name, Image.open(output_path)))
else:
saved_paths.append(f"{view_name}: FAILED")
except Exception as e:
saved_paths.append(f"{view_name}: ERROR - {e}")
# Create composite image (3 columns x 2 rows)
if len(images) == 6:
composite_path = SCREENSHOT_DIR / "view_composite.png"
composite = Image.new('RGBA', (width * 3, height * 2), (0, 0, 0, 255))
positions = [(0, 0), (1, 0), (2, 0), (0, 1), (1, 1), (2, 1)]
for (name, img), (col, row) in zip(images, positions):
composite.paste(img, (col * width, row * height))
composite.save(composite_path)
return str(composite_path)
return "\n".join(saved_paths)
@mcp.tool()
def list_screenshots(ctx: Context) -> str:
"""List all captured PyMOL screenshots in the current session."""
screenshots = sorted(SCREENSHOT_DIR.glob("view_*.png"))
if not screenshots:
return f"No screenshots captured yet.\nDirectory: {SCREENSHOT_DIR}"
result = f"Screenshots ({len(screenshots)}/{_max_screenshots} max):\n"
for ss in screenshots:
try:
stat = ss.stat()
result += f" {ss} ({stat.st_size // 1024}KB)\n"
except OSError:
result += f" {ss} (file error)\n"
return result
@mcp.tool()
def clear_screenshots(ctx: Context) -> str:
"""Delete all captured PyMOL screenshots."""
global _screenshot_counter
count = 0
for f in SCREENSHOT_DIR.glob("view_*.png"):
try:
f.unlink()
count += 1
except OSError:
pass
_screenshot_counter = 0
return f"Deleted {count} screenshots from {SCREENSHOT_DIR}"
@mcp.tool()
def configure_screenshots(ctx: Context, max_count: int = 20) -> str:
"""
Configure screenshot auto-cleanup threshold.
Args:
max_count: Maximum screenshots to keep (oldest auto-deleted when exceeded)
"""
global _max_screenshots
if max_count < 1:
return "Error: max_count must be at least 1"
if max_count > 100:
return "Error: max_count cannot exceed 100"
_max_screenshots = max_count
return f"Screenshot limit set to {max_count}. Oldest will be auto-deleted when exceeded."
##############################################################################
# ENTRY POINT
##############################################################################
def main():
mcp.run()
if __name__ == "__main__":
main()
