mcp-stata 1.2.2__py3-none-any.whl → 1.6.2__py3-none-any.whl

mcp_stata/discovery.py

@@ -4,11 +4,30 @@ import glob
 import logging
 import shutil
 
-from typing import Tuple, Optional, List
+from typing import Tuple, List
 
 logger = logging.getLogger("mcp_stata.discovery")
 
 
+def _normalize_env_path(raw: str) -> str:
+    """Strip quotes/whitespace and expand variables for STATA_PATH."""
+    cleaned = raw.strip()
+    if (cleaned.startswith("\"") and cleaned.endswith("\"")) or (
+        cleaned.startswith("'") and cleaned.endswith("'")
+    ):
+        cleaned = cleaned[1:-1].strip()
+    return os.path.expandvars(os.path.expanduser(cleaned))
+
+
+def _is_executable(path: str, system: str) -> bool:
+    if not os.path.exists(path):
+        return False
+    if system == "Windows":
+        # On Windows, check if it's a file and has .exe extension
+        return os.path.isfile(path) and path.lower().endswith('.exe')
+    return os.access(path, os.X_OK)
+
+
 def _dedupe_preserve(items: List[tuple]) -> List[tuple]:
     seen = set()
     unique = []
@@ -27,9 +46,60 @@ def find_stata_path() -> Tuple[str, str]:
     """
     system = platform.system()
 
-    # 1. Check Environment Variable
+    windows_binaries = [
+        ("StataMP-64.exe", "mp"),
+        ("StataMP.exe", "mp"),
+        ("StataSE-64.exe", "se"),
+        ("StataSE.exe", "se"),
+        ("Stata-64.exe", "be"),
+        ("Stata.exe", "be"),
+    ]
+
+    linux_binaries = [
+        ("stata-mp", "mp"),
+        ("stata-se", "se"),
+        ("stata-ic", "be"),
+        ("stata", "be"),
+        ("xstata-mp", "mp"),
+        ("xstata-se", "se"),
+        ("xstata-ic", "be"),
+        ("xstata", "be"),
+    ]
+
+    # 1. Check Environment Variable (supports quoted values and directory targets)
     if os.environ.get("STATA_PATH"):
-        path = os.environ["STATA_PATH"]
+        raw_path = os.environ["STATA_PATH"]
+        path = _normalize_env_path(raw_path)
+        logger.info("Using STATA_PATH override (normalized): %s", path)
+
+        # If a directory is provided, try standard binaries for the platform
+        if os.path.isdir(path):
+            search_set = []
+            if system == "Windows":
+                search_set = windows_binaries
+            elif system == "Linux":
+                search_set = linux_binaries
+            elif system == "Darwin":
+                search_set = [
+                    ("Contents/MacOS/stata-mp", "mp"),
+                    ("Contents/MacOS/stata-se", "se"),
+                    ("Contents/MacOS/stata", "be"),
+                    ("stata-mp", "mp"),
+                    ("stata-se", "se"),
+                    ("stata", "be"),
+                ]
+
+            for binary, edition in search_set:
+                candidate = os.path.join(path, binary)
+                if _is_executable(candidate, system):
+                    logger.info("Found Stata via STATA_PATH directory: %s (%s)", candidate, edition)
+                    return candidate, edition
+
+            raise FileNotFoundError(
+                f"STATA_PATH points to directory '{path}', but no Stata executable was found within. "
+                "Point STATA_PATH directly to the Stata binary (e.g., C:\\Program Files\\Stata18\\StataMP-64.exe)."
+            )
+
         edition = "be"
         lower_path = path.lower()
         if "mp" in lower_path:
@@ -44,7 +114,7 @@ def find_stata_path() -> Tuple[str, str]:
                 "Update STATA_PATH to your Stata binary (e.g., "
                 "/Applications/StataNow/StataMP.app/Contents/MacOS/stata-mp or /usr/local/stata18/stata-mp)."
             )
-        if not os.access(path, os.X_OK):
+        if not _is_executable(path, system):
             raise PermissionError(
                 f"STATA_PATH points to '{path}', but it is not executable. "
                 "Ensure this is the Stata binary, not the .app directory."
@@ -84,29 +154,13 @@ def find_stata_path() -> Tuple[str, str]:
 
         for base_dir in base_dirs:
             for stata_dir in glob.glob(os.path.join(base_dir, "Stata*")):
-                for exe, edition in [
-                    ("StataMP-64.exe", "mp"),
-                    ("StataMP.exe", "mp"),
-                    ("StataSE-64.exe", "se"),
-                    ("StataSE.exe", "se"),
-                    ("Stata-64.exe", "be"),
-                    ("Stata.exe", "be"),
-                ]:
+                for exe, edition in windows_binaries:
                     full_path = os.path.join(stata_dir, exe)
                     if os.path.exists(full_path):
                         candidates.append((full_path, edition))
 
     elif system == "Linux":
-        linux_binaries = [
-            ("stata-mp", "mp"),
-            ("stata-se", "se"),
-            ("stata-ic", "be"),
-            ("stata", "be"),
-            ("xstata-mp", "mp"),
-            ("xstata-se", "se"),
-            ("xstata-ic", "be"),
-            ("xstata", "be"),
-        ]
+        home_base = os.environ.get("HOME") or os.path.expanduser("~")
 
         # 2a. Try binaries available on PATH first
         for binary, edition in linux_binaries:
@@ -118,8 +172,8 @@ def find_stata_path() -> Tuple[str, str]:
         linux_roots = [
             "/usr/local",
             "/opt",
-            os.path.expanduser("~/stata"),
-            os.path.expanduser("~/Stata"),
+            os.path.join(home_base, "stata"),
+            os.path.join(home_base, "Stata"),
         ]
 
         for root in linux_roots:
@@ -143,13 +197,14 @@ def find_stata_path() -> Tuple[str, str]:
                         if os.path.exists(full_path):
                             candidates.append((full_path, edition))
 
+
     candidates = _dedupe_preserve(candidates)
 
     for path, edition in candidates:
         if not os.path.exists(path):
             logger.warning("Discovered candidate missing on disk: %s", path)
             continue
-        if not os.access(path, os.X_OK):
+        if not _is_executable(path, system):
             logger.warning("Discovered candidate is not executable: %s", path)
             continue
         logger.info("Auto-discovered Stata at %s (%s)", path, edition)
@@ -160,3 +215,19 @@ def find_stata_path() -> Tuple[str, str]:
         "Set STATA_PATH to your Stata executable (e.g., "
         "/Applications/StataNow/StataMP.app/Contents/MacOS/stata-mp, /usr/local/stata18/stata-mp, or C:\\Program Files\\Stata18\\StataMP-64.exe)."
     )
+
+
+def main() -> int:
+    """CLI helper to print discovered Stata binary and edition."""
+    try:
+        path, edition = find_stata_path()
+        # Print so CLI users and tests see the output on stdout.
+        print(f"Stata executable: {path}\nEdition: {edition}")
+        return 0
+    except Exception as exc:  # pragma: no cover - exercised via tests with env
+        print(f"Discovery failed: {exc}")
+        return 1
+
+
+if __name__ == "__main__":  # pragma: no cover - manual utility
+    raise SystemExit(main())

mcp_stata/graph_detector.py

@@ -0,0 +1,385 @@
+"""
+Graph creation detection for streaming Stata output.
+
+This module provides functionality to detect when graphs are created
+during Stata command execution and automatically cache them.
+"""
+
+import asyncio
+import re
+import threading
+import time
+from typing import List, Set, Callable, Dict, Any
+import logging
+
+
+# SFI is always available
+SFI_AVAILABLE = True
+
+logger = logging.getLogger(__name__)
+
+
+class GraphCreationDetector:
+    """Detects graph creation using SFI-only detection with pystata integration."""
+    
+    def __init__(self, stata_client=None):
+        self._lock = threading.Lock()
+        self._detected_graphs: Set[str] = set()
+        self._removed_graphs: Set[str] = set()
+        self._unnamed_graph_counter = 0  # Track unnamed graphs for identification
+        self._stata_client = stata_client
+        self._last_graph_state: Dict[str, Any] = {}  # Track graph state changes
+
+    def _describe_graph_signature(self, graph_name: str) -> str:
+        """Return a stable signature for a graph.
+
+        We intentionally avoid using timestamps as the signature, since that makes
+        every poll look like a modification.
+        """
+        if not self._stata_client or not hasattr(self._stata_client, "stata"):
+            return ""
+        try:
+            # Capture output so we can hash it deterministically.
+            resp = self._stata_client.run_command_structured(f"graph describe {graph_name}", echo=False)
+            if resp.success and resp.stdout:
+                return resp.stdout
+            if resp.error and resp.error.snippet:
+                return resp.error.snippet
+        except Exception:
+            return ""
+        return ""
+    
+    def _detect_graphs_via_pystata(self) -> List[str]:
+        """Detect newly created graphs using direct pystata state access."""
+        if not self._stata_client:
+            return []
+        
+        try:
+            # Get current graph state using pystata's sfi interface
+            current_graphs = self._get_current_graphs_from_pystata()
+            current_state = self._get_graph_state_from_pystata()
+            
+            # Compare with last known state to detect new graphs
+            new_graphs = []
+            
+            # Check for new graph names
+            for graph_name in current_graphs:
+                if graph_name not in self._last_graph_state and graph_name not in self._removed_graphs:
+                    new_graphs.append(graph_name)
+            
+            # Check for state changes in existing graphs (modifications)
+            for graph_name, state in current_state.items():
+                if graph_name in self._last_graph_state:
+                    last_state = self._last_graph_state[graph_name]
+                    # Compare stable signature only.
+                    if state.get("signature") != last_state.get("signature"):
+                        if graph_name not in self._removed_graphs:
+                            new_graphs.append(graph_name)
+            
+            # Update cached state
+            self._last_graph_state = current_state.copy()
+            
+            return new_graphs
+            
+        except (ImportError, RuntimeError, ValueError, AttributeError) as e:
+            # These are expected exceptions when SFI is not available or Stata state is inaccessible
+            logger.debug(f"Failed to detect graphs via pystata (expected): {e}")
+            return []
+        except Exception as e:
+            # Unexpected errors should be logged as errors
+            logger.error(f"Unexpected error in pystata graph detection: {e}")
+            return []
+    
+    def _get_current_graphs_from_pystata(self) -> List[str]:
+        """Get current list of graphs using pystata's sfi interface."""
+        try:
+            # Use pystata to get graph list directly
+            if self._stata_client and hasattr(self._stata_client, 'list_graphs'):
+                return self._stata_client.list_graphs()
+            else:
+                # Fallback to sfi Macro interface - only if stata is available
+                if self._stata_client and hasattr(self._stata_client, 'stata'):
+                    try:
+                        from sfi import Macro
+                        self._stata_client.stata.run("quietly graph dir, memory")
+                        self._stata_client.stata.run("global mcp_graph_list `r(list)'")
+                        graph_list_str = Macro.getGlobal("mcp_graph_list")
+                        return graph_list_str.split() if graph_list_str else []
+                    except ImportError:
+                        logger.warning("sfi.Macro not available for fallback graph detection")
+                        return []
+                else:
+                    return []
+        except Exception as e:
+            logger.warning(f"Failed to get current graphs from pystata: {e}")
+            return []
+    
+    def _get_graph_state_from_pystata(self) -> Dict[str, Any]:
+        """Get detailed graph state information using pystata's sfi interface."""
+        graph_state = {}
+        
+        try:
+            current_graphs = self._get_current_graphs_from_pystata()
+            
+            for graph_name in current_graphs:
+                try:
+                    signature = self._describe_graph_signature(graph_name)
+                    state_info = {
+                        "name": graph_name,
+                        "exists": True,
+                        "valid": bool(signature),
+                        "signature": signature,
+                    }
+
+                    # Only update timestamps when the signature changes.
+                    prev = self._last_graph_state.get(graph_name)
+                    if prev is None or prev.get("signature") != signature:
+                        state_info["timestamp"] = time.time()
+                    else:
+                        state_info["timestamp"] = prev.get("timestamp", time.time())
+                    
+                    graph_state[graph_name] = state_info
+                    
+                except Exception as e:
+                    logger.warning(f"Failed to get state for graph {graph_name}: {e}")
+                    graph_state[graph_name] = {"name": graph_name, "timestamp": time.time(), "exists": False, "signature": ""}
+            
+        except Exception as e:
+            logger.warning(f"Failed to get graph state from pystata: {e}")
+        
+        return graph_state
+    
+        
+        
+    def detect_graph_modifications(self, text: str = None) -> dict:
+        """Detect graph modification/removal using SFI state comparison."""
+        modifications = {"dropped": [], "renamed": [], "cleared": False}
+        
+        if not self._stata_client:
+            return modifications
+        
+        try:
+            # Get current graph state via SFI
+            current_graphs = set(self._get_current_graphs_from_pystata())
+            
+            # Compare with last known state to detect modifications
+            if self._last_graph_state:
+                last_graphs = set(self._last_graph_state.keys())
+                
+                # Detect dropped graphs (in last state but not current)
+                dropped_graphs = last_graphs - current_graphs
+                modifications["dropped"].extend(dropped_graphs)
+                
+                # Detect clear all (no graphs remain when there were some before)
+                if last_graphs and not current_graphs:
+                    modifications["cleared"] = True
+            
+            # Update last known state for next comparison (stable signatures)
+            new_state: Dict[str, Any] = {}
+            for graph in current_graphs:
+                sig = self._describe_graph_signature(graph)
+                new_state[graph] = {
+                    "name": graph,
+                    "exists": True,
+                    "valid": bool(sig),
+                    "signature": sig,
+                    "timestamp": time.time(),
+                }
+            self._last_graph_state = new_state
+            
+        except Exception as e:
+            logger.debug(f"SFI modification detection failed: {e}")
+        
+        return modifications
+    
+        
+    def should_cache_graph(self, graph_name: str) -> bool:
+        """Determine if a graph should be cached."""
+        with self._lock:
+            # Don't cache if already detected or removed
+            if graph_name in self._detected_graphs or graph_name in self._removed_graphs:
+                return False
+            
+            # Mark as detected
+            self._detected_graphs.add(graph_name)
+            return True
+    
+    def mark_graph_removed(self, graph_name: str) -> None:
+        """Mark a graph as removed."""
+        with self._lock:
+            self._removed_graphs.add(graph_name)
+            self._detected_graphs.discard(graph_name)
+    
+    def mark_all_cleared(self) -> None:
+        """Mark all graphs as cleared."""
+        with self._lock:
+            self._detected_graphs.clear()
+            self._removed_graphs.clear()
+    
+    def clear_detection_state(self) -> None:
+        """Clear all detection state."""
+        with self._lock:
+            self._detected_graphs.clear()
+            self._removed_graphs.clear()
+            self._unnamed_graph_counter = 0
+    
+    def process_modifications(self, modifications: dict) -> None:
+        """Process detected modifications."""
+        with self._lock:
+            # Handle dropped graphs
+            for graph_name in modifications.get("dropped", []):
+                self.mark_graph_removed(graph_name)
+            
+            # Handle renamed graphs
+            for old_name, new_name in modifications.get("renamed", []):
+                self.mark_graph_removed(old_name)
+                self._detected_graphs.discard(new_name)  # Allow re-detection with new name
+            
+            # Handle clear all
+            if modifications.get("cleared", False):
+                self.mark_all_cleared()
+
+
+class StreamingGraphCache:
+    """Integrates graph detection with caching during streaming."""
+    
+    def __init__(self, stata_client, auto_cache: bool = False):
+        self.stata_client = stata_client
+        self.auto_cache = auto_cache
+        self.detector = GraphCreationDetector(stata_client)
+        self._lock = threading.Lock()
+        self._cache_callbacks: List[Callable[[str, bool], None]] = []
+        self._graphs_to_cache: List[str] = []
+        self._cached_graphs: Set[str] = set()
+        self._removed_graphs = set()  # Track removed graphs directly
+        self._initial_graphs: Set[str] = set()  # Captured before execution starts
+    
+    def add_cache_callback(self, callback: Callable[[str, bool], None]) -> None:
+        """Add callback for graph cache events."""
+        with self._lock:
+            self._cache_callbacks.append(callback)
+
+    
+    async def cache_detected_graphs_with_pystata(self) -> List[str]:
+        """Enhanced caching method that uses pystata for real-time graph detection."""
+        if not self.auto_cache:
+            return []
+        
+        cached_names = []
+        
+        # First, try to get any newly detected graphs via pystata state
+        if self.stata_client:
+            try:
+                # Get current state and check for new graphs
+
+                pystata_detected = self.detector._detect_graphs_via_pystata()
+                
+                # Add any newly detected graphs to cache queue
+                for graph_name in pystata_detected:
+                    if graph_name not in self._cached_graphs and graph_name not in self._removed_graphs:
+                        self._graphs_to_cache.append(graph_name)
+                        
+            except Exception as e:
+                logger.warning(f"Failed to get pystata graph updates: {e}")
+        
+        # Process the cache queue
+        with self._lock:
+            graphs_to_process = self._graphs_to_cache.copy()
+            self._graphs_to_cache.clear()
+        
+        # Get current graph list for verification
+        try:
+            current_graphs = self.stata_client.list_graphs()
+        except Exception as e:
+            logger.warning(f"Failed to get current graph list: {e}")
+            return cached_names
+        
+        for graph_name in graphs_to_process:
+            if graph_name in current_graphs and graph_name not in self._cached_graphs:
+                try:
+                    success = await asyncio.to_thread(self.stata_client.cache_graph_on_creation, graph_name)
+                    if success:
+                        cached_names.append(graph_name)
+                        with self._lock:
+                            self._cached_graphs.add(graph_name)
+                    
+                    # Notify callbacks
+                    for callback in self._cache_callbacks:
+                        try:
+                            callback(graph_name, success)
+                        except Exception as e:
+                            logger.warning(f"Cache callback failed for {graph_name}: {e}")
+                
+                except Exception as e:
+                    logger.warning(f"Failed to cache graph {graph_name}: {e}")
+                    # Still notify callbacks of failure
+                    for callback in self._cache_callbacks:
+                        try:
+                            callback(graph_name, False)
+                        except Exception:
+                            pass
+        
+        return cached_names
+    
+    async def cache_detected_graphs(self) -> List[str]:
+        """Cache all detected graphs."""
+        if not self.auto_cache:
+            return []
+        
+        cached_names = []
+        
+        with self._lock:
+            graphs_to_process = self._graphs_to_cache.copy()
+            self._graphs_to_cache.clear()
+        
+        # Get current graph list for verification
+        try:
+            current_graphs = self.stata_client.list_graphs()
+        except Exception as e:
+            logger.warning(f"Failed to get current graph list: {e}")
+            return cached_names
+        
+        for graph_name in graphs_to_process:
+            if graph_name in current_graphs and graph_name not in self._cached_graphs:
+                try:
+                    success = await asyncio.to_thread(self.stata_client.cache_graph_on_creation, graph_name)
+                    if success:
+                        cached_names.append(graph_name)
+                        with self._lock:
+                            self._cached_graphs.add(graph_name)
+                    
+                    # Notify callbacks
+                    for callback in self._cache_callbacks:
+                        try:
+                            callback(graph_name, success)
+                        except Exception as e:
+                            logger.warning(f"Cache callback failed for {graph_name}: {e}")
+                
+                except Exception as e:
+                    logger.warning(f"Failed to cache graph {graph_name}: {e}")
+                    # Still notify callbacks of failure
+                    for callback in self._cache_callbacks:
+                        try:
+                            callback(graph_name, False)
+                        except Exception:
+                            pass
+        
+        return cached_names
+    
+    def get_cache_stats(self) -> dict:
+        """Get caching statistics."""
+        with self._lock:
+            return {
+                "auto_cache_enabled": self.auto_cache,
+                "pending_cache_count": len(self._graphs_to_cache),
+                "cached_graphs_count": len(self._cached_graphs),
+                "detected_graphs_count": len(self.detector._detected_graphs),
+                "removed_graphs_count": len(self.detector._removed_graphs),
+            }
+    
+    def reset(self) -> None:
+        """Reset the cache state."""
+        with self._lock:
+            self._graphs_to_cache.clear()
+            self._cached_graphs.clear()
+        self.detector.clear_detection_state()

mcp_stata/models.py

@@ -7,6 +7,7 @@ class ErrorEnvelope(BaseModel):
     rc: Optional[int] = None
     line: Optional[int] = None
     command: Optional[str] = None
+    log_path: Optional[str] = None
     stdout: Optional[str] = None
     stderr: Optional[str] = None
     snippet: Optional[str] = None
@@ -18,6 +19,7 @@ class CommandResponse(BaseModel):
     rc: int
     stdout: str
     stderr: Optional[str] = None
+    log_path: Optional[str] = None
     success: bool
     error: Optional[ErrorEnvelope] = None
 
@@ -49,7 +51,8 @@ class GraphListResponse(BaseModel):
 
 class GraphExport(BaseModel):
     name: str
-    image_base64: str
+    file_path: Optional[str] = None
+    image_base64: Optional[str] = None
 
 
 class GraphExportResponse(BaseModel):