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

mcp_stata/discovery.py

@@ -375,18 +375,19 @@ def find_stata_candidates() -> List[Tuple[str, str]]:
     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",
             "/Applications/Stata/StataMP.app",
             "/Applications/Stata/StataSE.app",
             "/Applications/Stata/Stata.app",
-            "/Applications/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_fast(binary_dir):
@@ -395,6 +396,7 @@ def find_stata_candidates() -> List[Tuple[str, str]]:
                     full_path = os.path.join(binary_dir, binary)
                     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.
@@ -443,6 +445,7 @@ def find_stata_candidates() -> List[Tuple[str, str]]:
                 full_path = os.path.join(stata_dir, exe)
                 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("~")
@@ -486,7 +489,8 @@ def find_stata_candidates() -> List[Tuple[str, str]]:
 
     # Final validation of candidates (still using fast checks)
     validated: List[Tuple[str, str]] = []
-    for path, edition in _sort_candidates(candidates):
+    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

mcp_stata/graph_detector.py

@@ -6,6 +6,7 @@ during Stata command execution and automatically cache them.
 """
 
 import asyncio
+import contextlib
 import inspect
 import re
 import threading
@@ -34,24 +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:
-            # Use lightweight execution to avoid heavy FS I/O for high-frequency polling
-            resp = self._stata_client.exec_lightweight(f"graph describe {graph_name}")
-                
-            if resp.success and resp.stdout:
-                return resp.stdout
-            if resp.error and resp.error.snippet:
-                # If using lightweight, error might be None or just string in stderr, 
-                # but run_command_structured returns proper error envelope.
-                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."""
@@ -103,21 +107,26 @@ class GraphCreationDetector:
             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
-                        hold_name = f"_mcp_detector_hold_{int(time.time() * 1000 % 1000000)}"
-                        self._stata_client.stata.run(f"capture _return hold {hold_name}", echo=False)
+                    # 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:
-                            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 []
+                            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:
@@ -256,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] = []

mcp_stata/models.py

@@ -55,7 +55,6 @@ class GraphListResponse(BaseModel):
 class GraphExport(BaseModel):
     name: str
     file_path: Optional[str] = None
-    image_base64: Optional[str] = None
 
 
 class GraphExportResponse(BaseModel):

mcp_stata/server.py

@@ -4,6 +4,7 @@ from dataclasses import dataclass
 from datetime import datetime
 from importlib.metadata import PackageNotFoundError, version
 from mcp.server.fastmcp import Context, FastMCP
+from mcp.server.fastmcp.utilities import logging as fastmcp_logging
 import mcp.types as types
 from .stata_client import StataClient
 from .models import (
@@ -27,30 +28,61 @@ from .ui_http import UIChannelManager
 
 # Configure logging
 logger = logging.getLogger("mcp_stata")
+payload_logger = logging.getLogger("mcp_stata.payloads")
+_LOGGING_CONFIGURED = False
 
 def setup_logging():
-    # Configure logging to stderr with immediate flush for MCP transport
+    global _LOGGING_CONFIGURED
+    if _LOGGING_CONFIGURED:
+        return
+    _LOGGING_CONFIGURED = True
     log_level = os.getenv("MCP_STATA_LOGLEVEL", "DEBUG").upper()
-    configure_root = os.getenv("MCP_STATA_CONFIGURE_LOGGING", "1").lower() not in {"0", "false", "no"}
-    
-    # Create a handler that flushes immediately
-    handler = logging.StreamHandler(sys.stderr)
-    handler.setLevel(getattr(logging, log_level, logging.DEBUG))
-    handler.setFormatter(logging.Formatter("[%(name)s] %(levelname)s: %(message)s"))
-    
-    # Configure root logger only if requested; avoid clobbering existing handlers.
-    if configure_root:
-        root_logger = logging.getLogger()
-        if not root_logger.handlers:
-            root_logger.addHandler(handler)
-            root_logger.setLevel(getattr(logging, log_level, logging.DEBUG))
-
-    # Also configure the mcp_stata logger explicitly without duplicating handlers.
+    app_handler = logging.StreamHandler(sys.stderr)
+    app_handler.setLevel(getattr(logging, log_level, logging.DEBUG))
+    app_handler.setFormatter(logging.Formatter("[%(name)s] %(levelname)s: %(message)s"))
+
+    mcp_handler = logging.StreamHandler(sys.stderr)
+    mcp_handler.setLevel(getattr(logging, log_level, logging.DEBUG))
+    mcp_handler.setFormatter(logging.Formatter("[%(name)s] %(levelname)s: %(message)s"))
+
+    payload_handler = logging.StreamHandler(sys.stderr)
+    payload_handler.setLevel(getattr(logging, log_level, logging.DEBUG))
+    payload_handler.setFormatter(logging.Formatter("[%(name)s] %(levelname)s: %(message)s"))
+
+    root_logger = logging.getLogger()
+    root_logger.handlers = []
+    root_logger.setLevel(logging.WARNING)
+
+    for name, item in logging.root.manager.loggerDict.items():
+        if not isinstance(item, logging.Logger):
+            continue
+        item.handlers = []
+        item.propagate = False
+        if item.level == logging.NOTSET:
+            item.setLevel(getattr(logging, log_level, logging.DEBUG))
+
+    logger.handlers = [app_handler]
+    logger.propagate = False
+
+    payload_logger.handlers = [payload_handler]
+    payload_logger.propagate = False
+
+    mcp_logger = logging.getLogger("mcp.server")
+    mcp_logger.handlers = [mcp_handler]
+    mcp_logger.propagate = False
+    mcp_logger.setLevel(getattr(logging, log_level, logging.DEBUG))
+
+    mcp_lowlevel = logging.getLogger("mcp.server.lowlevel.server")
+    mcp_lowlevel.handlers = [mcp_handler]
+    mcp_lowlevel.propagate = False
+    mcp_lowlevel.setLevel(getattr(logging, log_level, logging.DEBUG))
+
+    mcp_root = logging.getLogger("mcp")
+    mcp_root.handlers = [mcp_handler]
+    mcp_root.propagate = False
+    mcp_root.setLevel(getattr(logging, log_level, logging.DEBUG))
     if logger.level == logging.NOTSET:
         logger.setLevel(getattr(logging, log_level, logging.DEBUG))
-    if not logger.handlers:
-        logger.addHandler(handler)
-    logger.propagate = False
 
     try:
         _mcp_stata_version = version("mcp-stata")
@@ -62,6 +94,8 @@ def setup_logging():
     logger.info("STATA_PATH env at startup: %s", os.getenv("STATA_PATH", "<not set>"))
     logger.info("LOG_LEVEL: %s", log_level)
 
+
+
 # Initialize FastMCP
 mcp = FastMCP("mcp_stata")
 client = StataClient()
@@ -81,6 +115,9 @@ class BackgroundTask:
 
 
 _background_tasks: Dict[str, BackgroundTask] = {}
+_request_log_paths: Dict[str, str] = {}
+_read_log_paths: set[str] = set()
+_read_log_offsets: Dict[str, int] = {}
 
 
 def _register_task(task_info: BackgroundTask, max_tasks: int = 100) -> None:
@@ -124,7 +161,6 @@ async def _notify_task_done(session: object | None, task_info: BackgroundTask, r
         "error": task_info.error,
     }
     try:
-        _debug_notification("logMessage", payload, request_id)
         await session.send_log_message(level="info", data=json.dumps(payload), related_request_id=request_id)
     except Exception:
         return
@@ -135,7 +171,7 @@ def _debug_notification(kind: str, payload: object, request_id: object | None =
         serialized = payload if isinstance(payload, str) else json.dumps(payload, ensure_ascii=False)
     except Exception:
         serialized = str(payload)
-    logger.debug("MCP notify %s request_id=%s payload=%s", kind, request_id, serialized)
+    payload_logger.info("MCP notify %s request_id=%s payload=%s", kind, request_id, serialized)
 
 
 async def _notify_tool_error(ctx: Context | None, tool_name: str, exc: Exception) -> None:
@@ -157,7 +193,6 @@ async def _notify_tool_error(ctx: Context | None, tool_name: str, exc: Exception
     if task_id is not None:
         payload["task_id"] = task_id
     try:
-        _debug_notification("logMessage", payload, ctx.request_id)
         await session.send_log_message(
             level="error",
             data=json.dumps(payload),
@@ -173,6 +208,20 @@ def _log_tool_call(tool_name: str, ctx: Context | None = None) -> None:
         request_id = getattr(ctx, "request_id", None)
     logger.info("MCP tool call: %s request_id=%s", tool_name, request_id)
 
+def _should_stream_smcl_chunk(text: str, request_id: object | None) -> bool:
+    if request_id is None:
+        return True
+    try:
+        payload = json.loads(text)
+        if isinstance(payload, dict) and payload.get("event"):
+            return True
+    except Exception:
+        pass
+    log_path = _request_log_paths.get(str(request_id))
+    if log_path and log_path in _read_log_paths:
+        return False
+    return True
+
 
 def _attach_task_id(ctx: Context | None, task_id: str) -> None:
     if ctx is None:
@@ -291,12 +340,21 @@ async def run_do_file_background(
 
     async def notify_log(text: str) -> None:
         if session is not None:
+            if not _should_stream_smcl_chunk(text, ctx.request_id):
+                return
             _debug_notification("logMessage", text, ctx.request_id)
-            await session.send_log_message(level="info", data=text, related_request_id=ctx.request_id)
+            try:
+                await session.send_log_message(level="info", data=text, related_request_id=ctx.request_id)
+            except Exception as e:
+                logger.warning("Failed to send logMessage notification: %s", e)
+                sys.stderr.write(f"[mcp_stata] ERROR: logMessage send failed: {e!r}\n")
+                sys.stderr.flush()
         try:
             payload = json.loads(text)
             if isinstance(payload, dict) and payload.get("event") == "log_path":
                 task_info.log_path = payload.get("path")
+                if ctx.request_id is not None and task_info.log_path:
+                    _request_log_paths[str(ctx.request_id)] = task_info.log_path
         except Exception:
             return
 
@@ -334,14 +392,17 @@ async def run_do_file_background(
                 graph_ready_task_id=task_id,
                 graph_ready_format="svg",
             )
-            ui_channel.notify_potential_dataset_change()
-            task_info.result = _format_command_result(result, raw=raw, as_json=as_json)
+            # Notify task completion as soon as the core operation is finished
+            task_info.done = True
             if result.error:
                 task_info.error = result.error.message
+            await _notify_task_done(session, task_info, request_id)
+
+            ui_channel.notify_potential_dataset_change()
+            task_info.result = _format_command_result(result, raw=raw, as_json=as_json)
         except Exception as exc:  # pragma: no cover - defensive
-            task_info.error = str(exc)
-        finally:
             task_info.done = True
+            task_info.error = str(exc)
             await _notify_task_done(session, task_info, request_id)
 
     task_info.task = asyncio.create_task(_run())
@@ -462,12 +523,16 @@ async def run_command_background(
 
     async def notify_log(text: str) -> None:
         if session is not None:
+            if not _should_stream_smcl_chunk(text, ctx.request_id):
+                return
             _debug_notification("logMessage", text, ctx.request_id)
             await session.send_log_message(level="info", data=text, related_request_id=ctx.request_id)
         try:
             payload = json.loads(text)
             if isinstance(payload, dict) and payload.get("event") == "log_path":
                 task_info.log_path = payload.get("path")
+                if ctx.request_id is not None and task_info.log_path:
+                    _request_log_paths[str(ctx.request_id)] = task_info.log_path
         except Exception:
             return
 
@@ -478,11 +543,6 @@ async def run_command_background(
     async def notify_progress(progress: float, total: float | None, message: str | None) -> None:
         if session is None or progress_token is None:
             return
-        _debug_notification(
-            "progress",
-            {"progress": progress, "total": total, "message": message},
-            ctx.request_id,
-        )
         await session.send_progress_notification(
             progress_token=progress_token,
             progress=progress,
@@ -505,14 +565,17 @@ async def run_command_background(
                 graph_ready_task_id=task_id,
                 graph_ready_format="svg",
             )
-            ui_channel.notify_potential_dataset_change()
-            task_info.result = _format_command_result(result, raw=raw, as_json=as_json)
+            # Notify task completion as soon as the core operation is finished
+            task_info.done = True
             if result.error:
                 task_info.error = result.error.message
+            await _notify_task_done(session, task_info, request_id)
+
+            ui_channel.notify_potential_dataset_change()
+            task_info.result = _format_command_result(result, raw=raw, as_json=as_json)
         except Exception as exc:  # pragma: no cover - defensive
-            task_info.error = str(exc)
-        finally:
             task_info.done = True
+            task_info.error = str(exc)
             await _notify_task_done(session, task_info, request_id)
 
     task_info.task = asyncio.create_task(_run())
@@ -558,8 +621,17 @@ async def run_command(
     async def notify_log(text: str) -> None:
         if session is None:
             return
+        if not _should_stream_smcl_chunk(text, ctx.request_id):
+            return
         _debug_notification("logMessage", text, ctx.request_id)
         await session.send_log_message(level="info", data=text, related_request_id=ctx.request_id)
+        try:
+            payload = json.loads(text)
+            if isinstance(payload, dict) and payload.get("event") == "log_path":
+                if ctx.request_id is not None:
+                    _request_log_paths[str(ctx.request_id)] = payload.get("path")
+        except Exception:
+            return
 
     progress_token = None
     if ctx is not None and ctx.request_context.meta is not None:
@@ -568,11 +640,6 @@ async def run_command(
     async def notify_progress(progress: float, total: float | None, message: str | None) -> None:
         if session is None or progress_token is None:
             return
-        _debug_notification(
-            "progress",
-            {"progress": progress, "total": total, "message": message},
-            ctx.request_id,
-        )
         await session.send_progress_notification(
             progress_token=progress_token,
             progress=progress,
@@ -627,12 +694,20 @@ def read_log(path: str, offset: int = 0, max_bytes: int = 65536) -> str:
     Returns a compact JSON string: {"path":..., "offset":..., "next_offset":..., "data":...}
     """
     try:
+        if path:
+            _read_log_paths.add(path)
         if offset < 0:
             offset = 0
+        if path:
+            last_offset = _read_log_offsets.get(path, 0)
+            if offset < last_offset:
+                offset = last_offset
         with open(path, "rb") as f:
             f.seek(offset)
             data = f.read(max_bytes)
             next_offset = f.tell()
+        if path:
+            _read_log_offsets[path] = next_offset
         text = data.decode("utf-8", errors="replace")
         return json.dumps({"path": path, "offset": offset, "next_offset": next_offset, "data": text})
     except FileNotFoundError:
@@ -910,8 +985,16 @@ async def run_do_file(
     async def notify_log(text: str) -> None:
         if session is None:
             return
-        _debug_notification("logMessage", text, ctx.request_id)
+        if not _should_stream_smcl_chunk(text, ctx.request_id):
+            return
         await session.send_log_message(level="info", data=text, related_request_id=ctx.request_id)
+        try:
+            payload = json.loads(text)
+            if isinstance(payload, dict) and payload.get("event") == "log_path":
+                if ctx.request_id is not None:
+                    _request_log_paths[str(ctx.request_id)] = payload.get("path")
+        except Exception:
+            return
 
     progress_token = None
     if ctx is not None and ctx.request_context.meta is not None:
@@ -920,11 +1003,6 @@ async def run_do_file(
     async def notify_progress(progress: float, total: float | None, message: str | None) -> None:
         if session is None or progress_token is None:
             return
-        _debug_notification(
-            "progress",
-            {"progress": progress, "total": total, "message": message},
-            ctx.request_id,
-        )
         await session.send_progress_notification(
             progress_token=progress_token,
             progress=progress,
@@ -1008,19 +1086,14 @@ def get_stored_results_resource() -> str:
     return json.dumps(client.get_stored_results())
 
 @mcp.tool()
-def export_graphs_all(use_base64: bool = False) -> str:
+def export_graphs_all() -> str:
     """
-    Exports all graphs in memory to file paths (default) or base64-encoded SVGs.
-
-    Args:
-        use_base64: If True, returns base64-encoded images (token-intensive).
-                   If False (default), returns file paths to SVG files (token-efficient).
-                   Use file paths unless you need to embed images directly.
+    Exports all graphs in memory to file paths.
 
-    Returns a JSON envelope listing graph names and either file paths or base64 images.
+    Returns a JSON envelope listing graph names and file paths.
     The agent can open SVG files directly to verify visuals (titles/labels/colors/legends).
     """
-    exports = client.export_graphs_all(use_base64=use_base64)
+    exports = client.export_graphs_all()
     return exports.model_dump_json(exclude_none=False)
 
 def main():