mcp-stata 1.22.1__cp311-abi3-macosx_11_0_x86_64.whl

mcp_stata/utils.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+import os
+import tempfile
+import pathlib
+import uuid
+import logging
+import threading
+import shutil
+import atexit
+import signal
+import sys
+from typing import Optional, List
+
+logger = logging.getLogger("mcp_stata")
+
+_temp_dir_cache: Optional[str] = None
+_temp_dir_lock = threading.Lock()
+_files_to_cleanup: set[pathlib.Path] = set()
+_dirs_to_cleanup: set[pathlib.Path] = set()
+
+def register_temp_file(path: str | pathlib.Path) -> None:
+    """
+    Register a file to be deleted on process exit.
+    Using this instead of NamedTemporaryFile(delete=True) because on Windows,
+    delete=True prevents Stata from opening the file simultaneously.
+    """
+    with _temp_dir_lock:
+        p = pathlib.Path(path).absolute()
+        _files_to_cleanup.add(p)
+
+def register_temp_dir(path: str | pathlib.Path) -> None:
+    """Register a directory to be recursively deleted on process exit."""
+    with _temp_dir_lock:
+        p = pathlib.Path(path).absolute()
+        _dirs_to_cleanup.add(p)
+
+def is_windows() -> bool:
+    """Returns True if the current operating system is Windows."""
+    return os.name == "nt"
+
+def _cleanup_temp_resources():
+    """Cleanup registered temporary files and directories."""
+    with _temp_dir_lock:
+        # Sort and copy to avoid modification during iteration
+        files = sorted(list(_files_to_cleanup), reverse=True)
+        for p in files:
+            try:
+                # missing_ok=True is Python 3.8+
+                p.unlink(missing_ok=True)
+                _files_to_cleanup.discard(p)
+            except Exception:
+                pass
+        
+        dirs = sorted(list(_dirs_to_cleanup), reverse=True)
+        for p in dirs:
+            try:
+                if p.exists() and p.is_dir():
+                    shutil.rmtree(p, ignore_errors=True)
+                _dirs_to_cleanup.discard(p)
+            except Exception:
+                pass
+
+atexit.register(_cleanup_temp_resources)
+
+def _signal_handler(signum, frame):
+    """Handle signals by cleaning up and exiting."""
+    _cleanup_temp_resources()
+    sys.exit(0)
+
+# Register signal handlers for graceful cleanup on termination
+try:
+    # Avoid hijacking signals if we are running in a test environment or not in main thread
+    is_pytest = "pytest" in sys.modules or "PYTEST_CURRENT_TEST" in os.environ
+    if threading.current_thread() is threading.main_thread() and not is_pytest:
+        signal.signal(signal.SIGTERM, _signal_handler)
+        signal.signal(signal.SIGINT, _signal_handler)
+except (ValueError, RuntimeError):
+    # Not in main thread or other signal handling restriction
+    pass
+
+def get_writable_temp_dir() -> str:
+    """
+    Finds a writable temporary directory by trying multiple fallback locations.
+    Priority:
+    1. MCP_STATA_TEMP environment variable
+    2. System Temp (tempfile.gettempdir())
+    3. User Home subdirectory (~/.mcp-stata/temp)
+    4. Current Working Directory subdirectory (.tmp)
+    
+    Results are cached after the first successful identification.
+    """
+    global _temp_dir_cache
+    
+    with _temp_dir_lock:
+        if _temp_dir_cache is not None:
+            return _temp_dir_cache
+        
+        candidates = []
+        
+        # 1. Environment variable
+        env_temp = os.getenv("MCP_STATA_TEMP")
+        if env_temp:
+            candidates.append((pathlib.Path(env_temp), "MCP_STATA_TEMP environment variable"))
+            
+        # 2. System Temp
+        candidates.append((pathlib.Path(tempfile.gettempdir()), "System temp directory"))
+        
+        # 3. User Home
+        try:
+            home_temp = pathlib.Path.home() / ".mcp-stata" / "temp"
+            candidates.append((home_temp, "User home directory"))
+        except Exception:
+            pass
+            
+        # 4. Current working directory subdirectory (.tmp)
+        candidates.append((pathlib.Path.cwd() / ".tmp", "Working directory (.tmp)"))
+        
+        tested_paths = []
+        for path, description in candidates:
+            try:
+                # Ensure directory exists
+                path.mkdir(parents=True, exist_ok=True)
+                
+                # Test writability using standard tempfile logic
+                try:
+                    fd, temp_path = tempfile.mkstemp(
+                        prefix=".mcp_write_test_", 
+                        suffix=".tmp", 
+                        dir=str(path)
+                    )
+                    os.close(fd)
+                    os.unlink(temp_path)
+                    
+                    # Success
+                    validated_path = str(path.absolute())
+                    
+                    # Log if we fell back from the first preferred (non-env) candidate
+                    # (System temp is second, index 1 if env_temp is set, else index 0)
+                    first_preferred_idx = 1 if env_temp else 0
+                    if candidates.index((path, description)) > first_preferred_idx:
+                        logger.warning(f"Falling back to temporary directory: {validated_path} ({description})")
+                    else:
+                        logger.debug(f"Using temporary directory: {validated_path} ({description})")
+                        
+                    _temp_dir_cache = validated_path
+                    # Globally set tempfile.tempdir so other parts of the app and libraries
+                    # use our validated writable path by default.
+                    tempfile.tempdir = validated_path
+                    return validated_path
+                except (OSError, PermissionError) as e:
+                    tested_paths.append(f"{path} ({description}): {e}")
+                    continue
+            except (OSError, PermissionError) as e:
+                tested_paths.append(f"{path} ({description}): {e}")
+                continue
+                
+        error_msg = "Failed to find any writable temporary directory. Errors:\n" + "\n".join(tested_paths)
+        logger.error(error_msg)
+        raise RuntimeError(error_msg)

mcp_stata/worker.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+import os
+import sys
+import threading
+import logging
+import json
+import traceback
+from typing import Any, Dict, Optional
+from multiprocessing.connection import Connection
+import asyncio
+
+# Ensure the parent directory is in sys.path so we can import modules
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from mcp_stata.stata_client import StataClient
+
+logger = logging.getLogger("mcp_stata.worker")
+
+class StataWorker:
+    def __init__(self, conn: Connection):
+        self.conn = conn
+        self.client: Optional[StataClient] = None
+        self.loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(self.loop)
+
+    def run(self):
+        """Main loop for the worker process."""
+        try:
+            # Initialize Stata in this process
+            self.client = StataClient()
+            # StataClient.init() will be called on first command if not already done,
+            # but we can do it here explicitly.
+            self.client.init()
+            
+            logger.info("StataWorker initialized and ready.")
+            self.conn.send({"event": "ready", "pid": os.getpid()})
+
+            while True:
+                if self.conn.poll(0.1):
+                    msg = self.conn.recv()
+                    if msg.get("type") == "stop":
+                        break
+                    
+                    # Handle command
+                    self.loop.run_until_complete(self.handle_message(msg))
+        except Exception as e:
+            logger.error(f"Worker process failed: {e}")
+            self.conn.send({"event": "error", "message": str(e), "traceback": traceback.format_exc()})
+        finally:
+            logger.info("Worker process exiting.")
+            self.conn.close()
+
+    async def handle_message(self, msg: Dict[str, Any]):
+        msg_type = msg.get("type")
+        msg_id = msg.get("id")
+        args = msg.get("args", {})
+
+        async def notify_log(text: str):
+            self.conn.send({"event": "log", "id": msg_id, "text": text})
+
+        async def notify_progress(progress: float, total: Optional[float], message: Optional[str]):
+            self.conn.send({"event": "progress", "id": msg_id, "progress": progress, "total": total, "message": message})
+
+        try:
+            if msg_type == "run_command":
+                result = await self.client.run_command_streaming(
+                    args["code"],
+                    notify_log=notify_log,
+                    notify_progress=notify_progress,
+                    **args.get("options", {})
+                )
+                self.conn.send({"event": "result", "id": msg_id, "result": result.model_dump()})
+            
+            elif msg_type == "run_do_file":
+                result = await self.client.run_do_file_streaming(
+                    args["path"],
+                    notify_log=notify_log,
+                    notify_progress=notify_progress,
+                    **args.get("options", {})
+                )
+                self.conn.send({"event": "result", "id": msg_id, "result": result.model_dump()})
+
+            elif msg_type == "get_data":
+                data = self.client.get_data(args.get("start", 0), args.get("count", 50))
+                self.conn.send({"event": "result", "id": msg_id, "result": data})
+
+            elif msg_type == "list_graphs":
+                graphs = self.client.list_graphs_structured()
+                self.conn.send({"event": "result", "id": msg_id, "result": graphs.model_dump()})
+
+            elif msg_type == "export_graph":
+                path = self.client.export_graph(args.get("graph_name"), format=args.get("format", "pdf"))
+                self.conn.send({"event": "result", "id": msg_id, "result": path})
+
+            elif msg_type == "get_help":
+                help_text = self.client.get_help(args["topic"], plain_text=args.get("plain_text", False))
+                self.conn.send({"event": "result", "id": msg_id, "result": help_text})
+
+            elif msg_type == "run_command_structured":
+                result = self.client.run_command_structured(args["code"], **args.get("options", {}))
+                self.conn.send({"event": "result", "id": msg_id, "result": result.model_dump()})
+
+            elif msg_type == "load_data":
+                result = self.client.load_data(args["source"], **args.get("options", {}))
+                self.conn.send({"event": "result", "id": msg_id, "result": result.model_dump()})
+
+            elif msg_type == "codebook":
+                result = self.client.codebook(args["variable"], **args.get("options", {}))
+                self.conn.send({"event": "result", "id": msg_id, "result": result.model_dump()})
+
+            elif msg_type == "get_dataset_state":
+                state = self.client.get_dataset_state()
+                self.conn.send({"event": "result", "id": msg_id, "result": state})
+
+            elif msg_type == "get_arrow_stream":
+                # StataClient.get_arrow_stream supports offset, limit, vars, etc.
+                arrow_bytes = self.client.get_arrow_stream(**args)
+                self.conn.send({"event": "result", "id": msg_id, "result": arrow_bytes})
+
+            elif msg_type == "list_variables_rich":
+                variables = self.client.list_variables_rich()
+                self.conn.send({"event": "result", "id": msg_id, "result": variables})
+
+            elif msg_type == "compute_view_indices":
+                indices = self.client.compute_view_indices(args["filter_expr"])
+                self.conn.send({"event": "result", "id": msg_id, "result": indices})
+
+            elif msg_type == "validate_filter_expr":
+                self.client.validate_filter_expr(args["filter_expr"])
+                self.conn.send({"event": "result", "id": msg_id, "result": None})
+
+            elif msg_type == "get_page":
+                page = self.client.get_page(**args)
+                self.conn.send({"event": "result", "id": msg_id, "result": page})
+
+            elif msg_type == "list_variables_structured":
+                variables = self.client.list_variables_structured()
+                self.conn.send({"event": "result", "id": msg_id, "result": variables.model_dump()})
+
+            elif msg_type == "export_graphs_all":
+                exports = self.client.export_graphs_all()
+                self.conn.send({"event": "result", "id": msg_id, "result": exports.model_dump()})
+
+            elif msg_type == "get_stored_results":
+                results = self.client.get_stored_results()
+                self.conn.send({"event": "result", "id": msg_id, "result": results})
+
+            else:
+                self.conn.send({"event": "error", "id": msg_id, "message": f"Unknown message type: {msg_type}"})
+
+        except Exception as e:
+            logger.error(f"Error handling message {msg_type}: {e}")
+            self.conn.send({
+                "event": "error", 
+                "id": msg_id, 
+                "message": str(e), 
+                "traceback": traceback.format_exc()
+            })
+
+def main(conn):
+    worker = StataWorker(conn)
+    worker.run()
+
+if __name__ == "__main__":
+    # This entry point is used when the process is started via multiprocessing
+    # But usually we'll pass the connection object from the parent.
+    pass