mcp-stata 1.7.6__py3-none-any.whl → 1.16.6__py3-none-any.whl

mcp_stata/config.py

@@ -0,0 +1,20 @@
+
+"""
+Central configuration for mcp-stata server and UI channel.
+"""
+from typing import Final
+
+# Server Limits
+MAX_LIMIT: Final[int] = 500  # Default row limit for JSON endpoints
+MAX_VARS: Final[int] = 32_767  # Max variables in Stata
+MAX_CHARS: Final[int] = 500  # Max chars per string cell to return
+MAX_REQUEST_BYTES: Final[int] = 1_000_000  # Max size of HTTP request body
+MAX_ARROW_LIMIT: Final[int] = 1_000_000  # Default row limit for Arrow IPC streams
+
+# Timeouts (seconds)
+TOKEN_TTL_S: Final[int] = 20 * 60  # Bearer token validity
+VIEW_TTL_S: Final[int] = 30 * 60  # Filtered view handle validity
+
+# Network
+DEFAULT_HOST: Final[str] = "127.0.0.1"
+DEFAULT_PORT: Final[int] = 0  # 0 = random ephemeral port

mcp_stata/discovery.py

@@ -1,10 +1,11 @@
 """
-Improved discovery.py with better error handling for intermittent failures.
+Optimized discovery.py with fast auto-discovery and targeted retry logic.
 Key improvements:
-1. Retry logic for file existence checks
-2. Better diagnostic logging
-3. Fuzzy path matching for common typos
-4. Case-insensitive path resolution on Windows
+1. Fast path checking during discovery (no retries)
+2. Retry logic only for validation of user-provided paths
+3. Better diagnostic logging
+4. Fuzzy path matching for common typos
+5. Case-insensitive path resolution on Windows
 """
 
 import os
@@ -15,15 +16,17 @@ import logging
 import shutil
 import ntpath
 import time
+import re
 from typing import Tuple, List, Optional
 
 logger = logging.getLogger("mcp_stata.discovery")
 
 
-def _exists_with_retry(path: str, max_attempts: int = 3, delay: float = 0.1) -> bool:
+def _exists_with_retry(path: str, max_attempts: int = 1, delay: float = 0.01) -> bool:
     """
     Check if file exists with retry logic to handle transient failures.
     This helps with antivirus scans, file locks, and other temporary issues.
+    Only use this for validating user-provided paths, not during discovery.
     """
     for attempt in range(max_attempts):
         if os.path.exists(path):
@@ -36,6 +39,11 @@ def _exists_with_retry(path: str, max_attempts: int = 3, delay: float = 0.1) ->
     return False
 
 
+def _exists_fast(path: str) -> bool:
+    """Fast existence check without retries for auto-discovery."""
+    return os.path.exists(path)
+
+
 def _find_similar_stata_dirs(target_path: str) -> List[str]:
     """
     Find similar Stata directories to help diagnose path typos.
@@ -72,6 +80,7 @@ def _validate_path_with_diagnostics(path: str, system: str) -> Tuple[bool, str]:
     """
     Validate path exists and provide detailed diagnostics if not.
     Returns (exists, diagnostic_message)
+    Uses retry logic for validation since this is for user-provided paths.
     """
     if _exists_with_retry(path):
         return True, ""
@@ -132,8 +141,14 @@ def _normalize_env_path(raw: str, system: str) -> str:
     return os.path.normpath(expanded)
 
 
-def _is_executable(path: str, system: str) -> bool:
-    if not _exists_with_retry(path):  # Use retry logic
+def _is_executable(path: str, system: str, use_retry: bool = True) -> bool:
+    """
+    Check if path is executable.
+    use_retry: Use retry logic for user-provided paths, fast check for discovery.
+    """
+    exists_check = _exists_with_retry if use_retry else _exists_fast
+    
+    if not exists_check(path):
         return False
     if system == "Windows":
         # On Windows, check if it's a file and has .exe extension
@@ -165,6 +180,37 @@ def _dedupe_str_preserve(items: List[str]) -> List[str]:
     return out
 
 
+def _extract_version_number(path: str) -> int:
+    """
+    Extract the highest Stata version number found in path components that
+    mention 'stata'. Returns 0 if no version is found.
+    """
+    version = 0
+    normalized = path.lower().replace("\\", os.sep)
+    for part in normalized.split(os.sep):
+        if "stata" not in part:
+            continue
+        for match in re.findall(r"(\d{1,3})", part):
+            try:
+                version = max(version, int(match))
+            except ValueError:
+                continue
+    return version
+
+
+def _sort_candidates(candidates: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+    """Sort candidates by version desc, edition (mp>se>be), then path for stability."""
+    edition_rank = {"mp": 3, "se": 2, "be": 1}
+
+    def sort_key(item: Tuple[str, str]):
+        path, edition = item
+        version = _extract_version_number(path)
+        rank = edition_rank.get((edition or "").lower(), 0)
+        return (-version, -rank, path)
+
+    return sorted(candidates, key=sort_key)
+
+
 def _resolve_windows_host_path(path: str, system: str) -> str:
     """
     On non-Windows hosts running Windows-discovery code, a Windows-style path
@@ -174,11 +220,11 @@ def _resolve_windows_host_path(path: str, system: str) -> str:
     """
     if system != "Windows":
         return path
-    if _exists_with_retry(path):  # Use retry logic
+    if _exists_fast(path):
         return path
     if os.sep != "\\" and "\\" in path:
         alt_path = path.replace("\\", os.sep)
-        if _exists_with_retry(alt_path):  # Use retry logic
+        if _exists_fast(alt_path):
             return alt_path
     return path
 
@@ -193,13 +239,18 @@ def _detect_system() -> str:
     return platform.system()
 
 
-def find_stata_path() -> Tuple[str, str]:
+def find_stata_candidates() -> List[Tuple[str, str]]:
     """
-    Attempts to automatically locate the Stata installation path.
-    Returns (path_to_executable, edition_string).
+    Locate all viable Stata installations ordered by preference.
+
+    Returns:
+        List of (path_to_executable, edition_string) sorted by:
+        - Newest version number found in path (desc)
+        - Edition preference: mp > se > be
+        - Path name (stable tie-breaker)
 
     Behavior:
-    - If STATA_PATH is set and valid, use it.
+    - If STATA_PATH is set and valid, use it (may yield multiple binaries in dir).
     - If STATA_PATH is set but invalid, provide detailed diagnostics and fall back.
     - If auto-discovery fails, raise an error with helpful suggestions.
     """
@@ -215,52 +266,62 @@ def find_stata_path() -> Tuple[str, str]:
         ("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)
-    raw_env_path = os.environ.get("STATA_PATH")
-    if raw_env_path:
+    # 1. Check STATA_PATH override with enhanced diagnostics
+    raw_stata_path = os.environ.get("STATA_PATH")
+    if raw_stata_path:
         try:
-            path = _normalize_env_path(raw_env_path, system)
+            path = _normalize_env_path(raw_stata_path, system)
             path = _resolve_windows_host_path(path, system)
-            logger.info("Trying STATA_PATH override (normalized): %s", path)
 
-            # If a directory is provided, try standard binaries for the platform
             if os.path.isdir(path):
-                search_set = []
+                candidates_in_dir = []
                 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
+                    for exe, edition in windows_binaries:
+                        candidate = os.path.join(path, exe)
+                        if _is_executable(candidate, system, use_retry=True):
+                            candidates_in_dir.append((candidate, edition))
+                elif system == "Darwin" or (system != "Windows" and path.endswith(".app")):
+                    # macOS app bundle logic
+                    sub_path = os.path.join(path, "Contents", "MacOS")
+                    if os.path.isdir(sub_path):
+                        for binary, edition in [("stata-mp", "mp"), ("stata-se", "se"), ("stata", "be")]:
+                            candidate = os.path.join(sub_path, binary)
+                            if _is_executable(candidate, system, use_retry=True):
+                                candidates_in_dir.append((candidate, edition))
+                    
+                    # Also try direct if not in a bundle
+                    if not candidates_in_dir:
+                        for binary, edition in linux_binaries:
+                            candidate = os.path.join(path, binary)
+                            if _is_executable(candidate, system, use_retry=True):
+                                candidates_in_dir.append((candidate, edition))
+                else:
+                    for binary, edition in linux_binaries:
+                        candidate = os.path.join(path, binary)
+                        if _is_executable(candidate, system, use_retry=True):
+                            candidates_in_dir.append((candidate, edition))
+
+                if candidates_in_dir:
+                    resolved = []
+                    for candidate, edition in _sort_candidates(candidates_in_dir):
+                        if _is_executable(candidate, system, use_retry=True):
+                            logger.info(
+                                "Found Stata via STATA_PATH directory: %s (%s)",
+                                candidate,
+                                edition,
+                            )
+                            resolved.append((candidate, edition))
+                    if resolved:
+                        return resolved
 
                 # Enhanced error with diagnostics
                 exists, diagnostics = _validate_path_with_diagnostics(path, system)
@@ -281,7 +342,7 @@ def find_stata_path() -> Tuple[str, str]:
             elif "be" in lower_path:
                 edition = "be"
 
-            # Use enhanced validation with diagnostics
+            # Use enhanced validation with diagnostics (with retry for user path)
             exists, diagnostics = _validate_path_with_diagnostics(path, system)
             if not exists:
                 error_msg = (
@@ -293,14 +354,14 @@ def find_stata_path() -> Tuple[str, str]:
                 )
                 raise FileNotFoundError(error_msg)
                 
-            if not _is_executable(path, system):
+            if not _is_executable(path, system, use_retry=True):
                 raise PermissionError(
                     f"STATA_PATH points to '{path}', but it is not executable. "
                     "Ensure this is the Stata binary, not the .app directory."
                 )
 
             logger.info("Using STATA_PATH override: %s (%s)", path, edition)
-            return path, edition
+            return [(path, edition)]
 
         except Exception as exc:
             stata_path_error = exc
@@ -310,11 +371,12 @@ def find_stata_path() -> Tuple[str, str]:
                 exc,
             )
 
-    # 2. Platform-specific search
+    # 2. Platform-specific search (using fast checks, no retries)
     candidates: List[Tuple[str, str]] = []  # List of (path, edition)
 
     if system == "Darwin":  # macOS
-        app_globs = [
+        # Search targets specific to macOS installation patterns
+        patterns = [
             "/Applications/StataNow/StataMP.app",
             "/Applications/StataNow/StataSE.app",
             "/Applications/StataNow/Stata.app",
@@ -322,17 +384,19 @@ def find_stata_path() -> Tuple[str, str]:
             "/Applications/Stata/StataSE.app",
             "/Applications/Stata/Stata.app",
             "/Applications/Stata*/Stata*.app",
+            "/Applications/Stata*.app",
         ]
 
-        for pattern in app_globs:
+        for pattern in patterns:
             for app_dir in glob.glob(pattern):
                 binary_dir = os.path.join(app_dir, "Contents", "MacOS")
-                if not _exists_with_retry(binary_dir):  # Use retry logic
+                if not _exists_fast(binary_dir):
                     continue
                 for binary, edition in [("stata-mp", "mp"), ("stata-se", "se"), ("stata", "be")]:
                     full_path = os.path.join(binary_dir, binary)
-                    if _exists_with_retry(full_path):  # Use retry logic
+                    if _exists_fast(full_path):
                         candidates.append((full_path, edition))
+        candidates = _dedupe_preserve(candidates)
 
     elif system == "Windows":
         # Include ProgramW6432 (real 64-bit Program Files) and hardcode fallbacks.
@@ -379,8 +443,9 @@ def find_stata_path() -> Tuple[str, str]:
                 continue
             for exe, edition in windows_binaries:
                 full_path = os.path.join(stata_dir, exe)
-                if _exists_with_retry(full_path):  # Use retry logic
+                if _exists_fast(full_path):
                     candidates.append((full_path, edition))
+        candidates = _dedupe_preserve(candidates)
 
     elif system == "Linux":
         home_base = os.environ.get("HOME") or os.path.expanduser("~")
@@ -417,20 +482,26 @@ def find_stata_path() -> Tuple[str, str]:
                         continue
                     for binary, edition in linux_binaries:
                         full_path = os.path.join(base_dir, binary)
-                        if _exists_with_retry(full_path):  # Use retry logic
+                        if _exists_fast(full_path):
                             candidates.append((full_path, edition))
 
-    candidates = _dedupe_preserve(candidates)
+        candidates = _dedupe_preserve(candidates)
 
-    for path, edition in candidates:
-        if not _exists_with_retry(path):  # Use retry logic
+    # Final validation of candidates (still using fast checks)
+    validated: List[Tuple[str, str]] = []
+    unique_candidates = _dedupe_preserve(candidates)
+    for path, edition in _sort_candidates(unique_candidates):
+        if not _exists_fast(path):
             logger.warning("Discovered candidate missing on disk: %s", path)
             continue
-        if not _is_executable(path, system):
+        if not _is_executable(path, system, use_retry=False):
             logger.warning("Discovered candidate is not executable: %s", path)
             continue
         logger.info("Auto-discovered Stata at %s (%s)", path, edition)
-        return path, edition
+        validated.append((path, edition))
+
+    if validated:
+        return validated
 
     # Build comprehensive error message
     error_parts = ["Could not automatically locate Stata."]
@@ -455,6 +526,14 @@ def find_stata_path() -> Tuple[str, str]:
         raise FileNotFoundError("\n".join(error_parts))
 
 
+def find_stata_path() -> Tuple[str, str]:
+    """
+    Backward-compatible wrapper returning the top-ranked candidate.
+    """
+    candidates = find_stata_candidates()
+    return candidates[0]
+
+
 def main() -> int:
     """CLI helper to print discovered Stata binary and edition."""
     try:

mcp_stata/graph_detector.py

@@ -6,6 +6,8 @@ during Stata command execution and automatically cache them.
 """
 
 import asyncio
+import contextlib
+import inspect
 import re
 import threading
 import time
@@ -33,21 +35,27 @@ class GraphCreationDetector:
     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.
+        We avoid using Stata calls like 'graph describe' here because they are slow 
+        (each call takes ~35ms) and would be called for every graph on every poll,
+        bottlenecking the streaming output.
+        
+        Instead, we use name-based tracking tied to the Stata command execution 
+        context. The signature is stable within a single command execution but 
+        changes when a new command starts, allowing us to detect modifications 
+        between commands without any Stata overhead.
         """
-        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:
+        if not self._stata_client:
             return ""
-        return ""
+        
+        # Access command_idx from stata_client if available
+        # NOTE: We only use command_idx for the default 'Graph' name to detect
+        # modifications. For named graphs, we only detect creation (name change)
+        # to avoid triggering redundant notifications for all existing graphs 
+        # on every command (since command_idx changes globally).
+        cmd_idx = getattr(self._stata_client, "_command_idx", 0)
+        if graph_name.lower() == "graph":
+            return f"{graph_name}_{cmd_idx}"
+        return graph_name
     
     def _detect_graphs_via_pystata(self) -> List[str]:
         """Detect newly created graphs using direct pystata state access."""
@@ -95,19 +103,30 @@ class GraphCreationDetector:
         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()
+                return self._stata_client.list_graphs(force_refresh=True)
             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 []
+                    # Access the lock from client to prevent concurrency issues with pystata
+                    exec_lock = getattr(self._stata_client, "_exec_lock", None)
+                    ctx = exec_lock if exec_lock else contextlib.nullcontext()
+                    
+                    with ctx:
+                        try:
+                            from sfi import Macro
+                            hold_name = f"_mcp_detector_hold_{int(time.time() * 1000 % 1000000)}"
+                            self._stata_client.stata.run(f"capture _return hold {hold_name}", echo=False)
+                            try:
+                                self._stata_client.stata.run("macro define mcp_graph_list \"\"", echo=False)
+                                self._stata_client.stata.run("quietly graph dir, memory", echo=False)
+                                self._stata_client.stata.run("macro define mcp_graph_list `r(list)'", echo=False)
+                                graph_list_str = Macro.getGlobal("mcp_graph_list")
+                            finally:
+                                self._stata_client.stata.run(f"capture _return restore {hold_name}", echo=False)
+                            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:
@@ -246,7 +265,11 @@ class StreamingGraphCache:
     def __init__(self, stata_client, auto_cache: bool = False):
         self.stata_client = stata_client
         self.auto_cache = auto_cache
-        self.detector = GraphCreationDetector(stata_client)
+        # Use persistent detector from client if available, else create local one
+        if hasattr(stata_client, "_graph_detector"):
+            self.detector = stata_client._graph_detector
+        else:
+            self.detector = GraphCreationDetector(stata_client)
         self._lock = threading.Lock()
         self._cache_callbacks: List[Callable[[str, bool], None]] = []
         self._graphs_to_cache: List[str] = []
@@ -259,6 +282,15 @@ class StreamingGraphCache:
         with self._lock:
             self._cache_callbacks.append(callback)
 
+    async def _notify_cache_callbacks(self, graph_name: str, success: bool) -> None:
+        for callback in self._cache_callbacks:
+            try:
+                result = callback(graph_name, success)
+                if inspect.isawaitable(result):
+                    await result
+            except Exception as e:
+                logger.warning(f"Cache callback failed for {graph_name}: {e}")
+
     
     async def cache_detected_graphs_with_pystata(self) -> List[str]:
         """Enhanced caching method that uses pystata for real-time graph detection."""
@@ -304,20 +336,12 @@ class StreamingGraphCache:
                             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}")
+                    await self._notify_cache_callbacks(graph_name, success)
                 
                 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
+                    await self._notify_cache_callbacks(graph_name, False)
         
         return cached_names
     
@@ -349,20 +373,12 @@ class StreamingGraphCache:
                             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}")
+                    await self._notify_cache_callbacks(graph_name, success)
                 
                 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
+                    await self._notify_cache_callbacks(graph_name, False)
         
         return cached_names
     

mcp_stata/models.py

@@ -13,6 +13,7 @@ class ErrorEnvelope(BaseModel):
     stderr: Optional[str] = None
     snippet: Optional[str] = None
     trace: Optional[bool] = None
+    smcl_output: Optional[str] = None
 
 
 class CommandResponse(BaseModel):
@@ -23,6 +24,7 @@ class CommandResponse(BaseModel):
     log_path: Optional[str] = None
     success: bool
     error: Optional[ErrorEnvelope] = None
+    smcl_output: Optional[str] = None
 
 
 class DataResponse(BaseModel):
@@ -53,7 +55,6 @@ class GraphListResponse(BaseModel):
 class GraphExport(BaseModel):
     name: str
     file_path: Optional[str] = None
-    image_base64: Optional[str] = None
 
 
 class GraphExportResponse(BaseModel):