mcp-stata 1.18.0__cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

mcp_stata/server.py

@@ -0,0 +1,1130 @@
+import anyio
+import asyncio
+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 (
+    DataResponse,
+    GraphListResponse,
+    VariablesResponse,
+    GraphExportResponse,
+)
+import logging
+import sys
+import json
+import os
+import re
+import traceback
+import uuid
+from functools import wraps
+from typing import Optional, Dict
+
+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():
+    global _LOGGING_CONFIGURED
+    if _LOGGING_CONFIGURED:
+        return
+    _LOGGING_CONFIGURED = True
+    log_level = os.getenv("MCP_STATA_LOGLEVEL", "DEBUG").upper()
+    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))
+
+    try:
+        _mcp_stata_version = version("mcp-stata")
+    except PackageNotFoundError:
+        _mcp_stata_version = "unknown"
+
+    logger.info("=== mcp-stata server starting ===")
+    logger.info("mcp-stata version: %s", _mcp_stata_version)
+    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()
+ui_channel = UIChannelManager(client)
+
+
+@dataclass
+class BackgroundTask:
+    task_id: str
+    kind: str
+    task: asyncio.Task
+    created_at: datetime
+    log_path: Optional[str] = None
+    result: Optional[str] = None
+    error: Optional[str] = None
+    done: bool = False
+
+
+_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:
+    _background_tasks[task_info.task_id] = task_info
+    if len(_background_tasks) <= max_tasks:
+        return
+    completed = [task for task in _background_tasks.values() if task.done]
+    completed.sort(key=lambda item: item.created_at)
+    for task in completed[: max(0, len(_background_tasks) - max_tasks)]:
+        _background_tasks.pop(task.task_id, None)
+
+
+def _format_command_result(result, raw: bool, as_json: bool) -> str:
+    if raw:
+        if result.success:
+            return result.log_path or ""
+        if result.error:
+            msg = result.error.message
+            if result.error.rc is not None:
+                msg = f"{msg}\nrc={result.error.rc}"
+            return msg
+        return result.log_path or ""
+    if as_json:
+        return result.model_dump_json()
+    return result.model_dump_json()
+
+
+async def _wait_for_log_path(task_info: BackgroundTask) -> None:
+    while task_info.log_path is None and not task_info.done:
+        await anyio.sleep(0.01)
+
+
+async def _notify_task_done(session: object | None, task_info: BackgroundTask, request_id: object | None) -> None:
+    if session is None:
+        return
+    payload = {
+        "event": "task_done",
+        "task_id": task_info.task_id,
+        "status": "done" if task_info.done else "unknown",
+        "log_path": task_info.log_path,
+        "error": task_info.error,
+    }
+    try:
+        await session.send_log_message(level="info", data=json.dumps(payload), related_request_id=request_id)
+    except Exception:
+        return
+
+
+def _debug_notification(kind: str, payload: object, request_id: object | None = None) -> None:
+    try:
+        serialized = payload if isinstance(payload, str) else json.dumps(payload, ensure_ascii=False)
+    except Exception:
+        serialized = str(payload)
+    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:
+    if ctx is None:
+        return
+    session = ctx.request_context.session
+    if session is None:
+        return
+    task_id = None
+    meta = ctx.request_context.meta
+    if meta is not None:
+        task_id = getattr(meta, "task_id", None) or getattr(meta, "taskId", None)
+    payload = {
+        "event": "tool_error",
+        "tool": tool_name,
+        "error": str(exc),
+        "traceback": traceback.format_exc(),
+    }
+    if task_id is not None:
+        payload["task_id"] = task_id
+    try:
+        await session.send_log_message(
+            level="error",
+            data=json.dumps(payload),
+            related_request_id=ctx.request_id,
+        )
+    except Exception:
+        logger.exception("Failed to emit tool_error notification for %s", tool_name)
+
+
+def _log_tool_call(tool_name: str, ctx: Context | None = None) -> None:
+    request_id = None
+    if ctx is not 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:
+        return
+    meta = ctx.request_context.meta
+    if meta is None:
+        meta = types.RequestParams.Meta()
+        ctx.request_context.meta = meta
+    try:
+        setattr(meta, "task_id", task_id)
+    except Exception:
+        logger.debug("Unable to attach task_id to request meta", exc_info=True)
+
+
+def _extract_ctx(args: tuple[object, ...], kwargs: dict[str, object]) -> Context | None:
+    ctx = kwargs.get("ctx")
+    if isinstance(ctx, Context):
+        return ctx
+    for arg in args:
+        if isinstance(arg, Context):
+            return arg
+    return None
+
+
+_mcp_tool = mcp.tool
+_mcp_resource = mcp.resource
+
+
+def tool(*tool_args, **tool_kwargs):
+    decorator = _mcp_tool(*tool_args, **tool_kwargs)
+
+    def outer(func):
+        if asyncio.iscoroutinefunction(func):
+            @wraps(func)
+            async def async_inner(*args, **kwargs):
+                ctx = _extract_ctx(args, kwargs)
+                _log_tool_call(func.__name__, ctx)
+                try:
+                    return await func(*args, **kwargs)
+                except Exception as exc:
+                    await _notify_tool_error(ctx, func.__name__, exc)
+                    raise
+
+            return decorator(async_inner)
+
+        @wraps(func)
+        def sync_inner(*args, **kwargs):
+            ctx = _extract_ctx(args, kwargs)
+            _log_tool_call(func.__name__, ctx)
+            try:
+                return func(*args, **kwargs)
+            except Exception:
+                logger.exception("Tool %s failed", func.__name__)
+                raise
+
+        return decorator(sync_inner)
+
+    return outer
+
+
+mcp.tool = tool
+
+
+def resource(*resource_args, **resource_kwargs):
+    decorator = _mcp_resource(*resource_args, **resource_kwargs)
+
+    def outer(func):
+        if asyncio.iscoroutinefunction(func):
+            @wraps(func)
+            async def async_inner(*args, **kwargs):
+                _log_tool_call(func.__name__, _extract_ctx(args, kwargs))
+                return await func(*args, **kwargs)
+
+            return decorator(async_inner)
+
+        @wraps(func)
+        def sync_inner(*args, **kwargs):
+            _log_tool_call(func.__name__, _extract_ctx(args, kwargs))
+            return func(*args, **kwargs)
+
+        return decorator(sync_inner)
+
+    return outer
+
+
+mcp.resource = resource
+
+
+@mcp.tool()
+async def run_do_file_background(
+    path: str,
+    ctx: Context | None = None,
+    echo: bool = True,
+    as_json: bool = True,
+    trace: bool = False,
+    raw: bool = False,
+    max_output_lines: int = None,
+    cwd: str | None = None,
+) -> str:
+    """Run a Stata do-file in the background and return a task id.
+
+    Notifications:
+      - logMessage: {"event":"log_path","path":"..."}
+      - logMessage: {"event":"task_done","task_id":"...","status":"done","log_path":"...","error":null}
+    """
+    session = ctx.request_context.session if ctx is not None else None
+    request_id = ctx.request_id if ctx is not None else None
+    task_id = uuid.uuid4().hex
+    _attach_task_id(ctx, task_id)
+    task_info = BackgroundTask(
+        task_id=task_id,
+        kind="do_file",
+        task=None,
+        created_at=datetime.utcnow(),
+    )
+
+    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)
+            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
+
+    progress_token = None
+    if ctx is not None and ctx.request_context.meta is not None:
+        progress_token = ctx.request_context.meta.progressToken
+
+    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,
+            total=total,
+            message=message,
+            related_request_id=ctx.request_id,
+        )
+
+    async def _run() -> None:
+        try:
+            result = await client.run_do_file_streaming(
+                path,
+                notify_log=notify_log,
+                notify_progress=notify_progress if progress_token is not None else None,
+                echo=echo,
+                trace=trace,
+                max_output_lines=max_output_lines,
+                cwd=cwd,
+                emit_graph_ready=True,
+                graph_ready_task_id=task_id,
+                graph_ready_format="svg",
+            )
+            # 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.done = True
+            task_info.error = str(exc)
+            await _notify_task_done(session, task_info, request_id)
+
+    task_info.task = asyncio.create_task(_run())
+    _register_task(task_info)
+    await _wait_for_log_path(task_info)
+    return json.dumps({"task_id": task_id, "status": "started", "log_path": task_info.log_path})
+
+
+@mcp.tool()
+def get_task_status(task_id: str, allow_polling: bool = False) -> str:
+    """Return task status for background executions.
+
+    Polling is disabled by default; set allow_polling=True for legacy callers.
+    """
+    notice = "Prefer task_done logMessage notifications over polling get_task_status."
+    if not allow_polling:
+        logger.warning(
+            "get_task_status called without allow_polling; clients must use task_done logMessage notifications"
+        )
+        return json.dumps({
+            "task_id": task_id,
+            "status": "polling_not_allowed",
+            "error": "Polling is disabled; use task_done logMessage notifications.",
+            "notice": notice,
+        })
+    logger.warning("get_task_status called; clients should use task_done logMessage notifications instead of polling")
+    task_info = _background_tasks.get(task_id)
+    if task_info is None:
+        return json.dumps({"task_id": task_id, "status": "not_found", "notice": notice})
+    return json.dumps({
+        "task_id": task_id,
+        "status": "done" if task_info.done else "running",
+        "kind": task_info.kind,
+        "created_at": task_info.created_at.isoformat(),
+        "log_path": task_info.log_path,
+        "error": task_info.error,
+        "notice": notice,
+    })
+
+
+@mcp.tool()
+def get_task_result(task_id: str, allow_polling: bool = False) -> str:
+    """Return task result for background executions.
+
+    Polling is disabled by default; set allow_polling=True for legacy callers.
+    """
+    notice = "Prefer task_done logMessage notifications over polling get_task_result."
+    if not allow_polling:
+        logger.warning(
+            "get_task_result called without allow_polling; clients must use task_done logMessage notifications"
+        )
+        return json.dumps({
+            "task_id": task_id,
+            "status": "polling_not_allowed",
+            "error": "Polling is disabled; use task_done logMessage notifications.",
+            "notice": notice,
+        })
+    logger.warning("get_task_result called; clients should use task_done logMessage notifications instead of polling")
+    task_info = _background_tasks.get(task_id)
+    if task_info is None:
+        return json.dumps({"task_id": task_id, "status": "not_found", "notice": notice})
+    if not task_info.done:
+        return json.dumps({
+            "task_id": task_id,
+            "status": "running",
+            "log_path": task_info.log_path,
+            "notice": notice,
+        })
+    return json.dumps({
+        "task_id": task_id,
+        "status": "done",
+        "log_path": task_info.log_path,
+        "error": task_info.error,
+        "notice": notice,
+        "result": task_info.result,
+    })
+
+
+@mcp.tool()
+def cancel_task(task_id: str) -> str:
+    """Request cancellation of a background task."""
+    task_info = _background_tasks.get(task_id)
+    if task_info is None:
+        return json.dumps({"task_id": task_id, "status": "not_found"})
+    if task_info.task and not task_info.task.done():
+        task_info.task.cancel()
+        return json.dumps({"task_id": task_id, "status": "cancelling"})
+    return json.dumps({"task_id": task_id, "status": "done", "log_path": task_info.log_path})
+
+
+@mcp.tool()
+async def run_command_background(
+    code: str,
+    ctx: Context | None = None,
+    echo: bool = True,
+    as_json: bool = True,
+    trace: bool = False,
+    raw: bool = False,
+    max_output_lines: int = None,
+    cwd: str | None = None,
+) -> str:
+    """Run a Stata command in the background and return a task id.
+
+    Notifications:
+      - logMessage: {"event":"log_path","path":"..."}
+      - logMessage: {"event":"task_done","task_id":"...","status":"done","log_path":"...","error":null}
+    """
+    session = ctx.request_context.session if ctx is not None else None
+    request_id = ctx.request_id if ctx is not None else None
+    task_id = uuid.uuid4().hex
+    _attach_task_id(ctx, task_id)
+    task_info = BackgroundTask(
+        task_id=task_id,
+        kind="command",
+        task=None,
+        created_at=datetime.utcnow(),
+    )
+
+    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
+
+    progress_token = None
+    if ctx is not None and ctx.request_context.meta is not None:
+        progress_token = ctx.request_context.meta.progressToken
+
+    async def notify_progress(progress: float, total: float | None, message: str | None) -> None:
+        if session is None or progress_token is None:
+            return
+        await session.send_progress_notification(
+            progress_token=progress_token,
+            progress=progress,
+            total=total,
+            message=message,
+            related_request_id=ctx.request_id,
+        )
+
+    async def _run() -> None:
+        try:
+            result = await client.run_command_streaming(
+                code,
+                notify_log=notify_log,
+                notify_progress=notify_progress if progress_token is not None else None,
+                echo=echo,
+                trace=trace,
+                max_output_lines=max_output_lines,
+                cwd=cwd,
+                emit_graph_ready=True,
+                graph_ready_task_id=task_id,
+                graph_ready_format="svg",
+            )
+            # 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.done = True
+            task_info.error = str(exc)
+            await _notify_task_done(session, task_info, request_id)
+
+    task_info.task = asyncio.create_task(_run())
+    _register_task(task_info)
+    await _wait_for_log_path(task_info)
+    return json.dumps({"task_id": task_id, "status": "started", "log_path": task_info.log_path})
+
+@mcp.tool()
+async def run_command(
+    code: str,
+    ctx: Context | None = None,
+    echo: bool = True,
+    as_json: bool = True,
+    trace: bool = False,
+    raw: bool = False,
+    max_output_lines: int = None,
+    cwd: str | None = None,
+) -> str:
+    """
+    Executes Stata code.
+
+    This is the primary tool for interacting with Stata.
+
+    Stata output is written to a temporary log file on disk.
+    The server emits a single `notifications/logMessage` event containing the log file path
+    (JSON payload: {"event":"log_path","path":"..."}) so the client can tail it locally.
+    If the client supplies a progress callback/token, progress updates may also be emitted
+    via `notifications/progress`.
+
+    Args:
+        code: The Stata command(s) to execute (e.g., "sysuse auto", "regress price mpg", "summarize").
+        ctx: FastMCP-injected request context (used to send MCP notifications). Optional for direct Python calls.
+        echo: If True, the command itself is included in the output. Default is True.
+        as_json: If True, returns a JSON envelope with rc/stdout/stderr/error.
+        trace: If True, enables `set trace on` for deeper error diagnostics (automatically disabled after).
+        raw: If True, return raw output/error message rather than a JSON envelope.
+        max_output_lines: If set, truncates stdout to this many lines for token efficiency.
+                         Useful for verbose commands (regress, codebook, etc.).
+        Note: This tool always uses log-file streaming semantics; there is no non-streaming mode.
+    """
+    session = ctx.request_context.session if ctx is not None else None
+
+    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:
+        progress_token = ctx.request_context.meta.progressToken
+
+    async def notify_progress(progress: float, total: float | None, message: str | None) -> None:
+        if session is None or progress_token is None:
+            return
+        await session.send_progress_notification(
+            progress_token=progress_token,
+            progress=progress,
+            total=total,
+            message=message,
+            related_request_id=ctx.request_id,
+        )
+
+    async def _noop_log(_text: str) -> None:
+        return
+
+    result = await client.run_command_streaming(
+        code,
+        notify_log=notify_log if session is not None else _noop_log,
+        notify_progress=notify_progress if progress_token is not None else None,
+        echo=echo,
+        trace=trace,
+        max_output_lines=max_output_lines,
+        cwd=cwd,
+        emit_graph_ready=True,
+        graph_ready_task_id=ctx.request_id if ctx else None,
+        graph_ready_format="svg",
+    )
+
+    # Conservative invalidation: arbitrary Stata commands may change data.
+    ui_channel.notify_potential_dataset_change()
+    if raw:
+        if result.success:
+            return result.log_path or ""
+        if result.error:
+            msg = result.error.message
+            if result.error.rc is not None:
+                msg = f"{msg}\nrc={result.error.rc}"
+            return msg
+        return result.log_path or ""
+    if as_json:
+        return result.model_dump_json()
+
+
+@mcp.tool()
+def read_log(path: str, offset: int = 0, max_bytes: int = 65536) -> str:
+    """Read a slice of a log file.
+
+    Intended for clients that want to display a terminal-like view without pushing MBs of
+    output through MCP log notifications.
+
+    Args:
+        path: Absolute path to the log file previously provided by the server.
+        offset: Byte offset to start reading from.
+        max_bytes: Maximum bytes to read.
+
+    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:
+        return json.dumps({"path": path, "offset": offset, "next_offset": offset, "data": ""})
+    except Exception as e:
+        return json.dumps({"path": path, "offset": offset, "next_offset": offset, "data": f"ERROR: {e}"})
+
+
+@mcp.tool()
+def find_in_log(
+    path: str,
+    query: str,
+    start_offset: int = 0,
+    max_bytes: int = 5_000_000,
+    before: int = 2,
+    after: int = 2,
+    case_sensitive: bool = False,
+    regex: bool = False,
+    max_matches: int = 50,
+) -> str:
+    """Find text within a log file and return context windows.
+
+    Args:
+        path: Absolute path to the log file previously provided by the server.
+        query: Text or regex pattern to search for.
+        start_offset: Byte offset to start searching from.
+        max_bytes: Maximum bytes to read from the log.
+        before: Number of context lines to include before each match.
+        after: Number of context lines to include after each match.
+        case_sensitive: If True, match case-sensitively.
+        regex: If True, treat query as a regular expression.
+        max_matches: Maximum number of matches to return.
+
+    Returns a JSON string with matches and offsets:
+        {"path":..., "query":..., "start_offset":..., "next_offset":..., "truncated":..., "matches":[...]}.
+    """
+    try:
+        if start_offset < 0:
+            start_offset = 0
+        if max_bytes <= 0:
+            return json.dumps({
+                "path": path,
+                "query": query,
+                "start_offset": start_offset,
+                "next_offset": start_offset,
+                "truncated": False,
+                "matches": [],
+            })
+        with open(path, "rb") as f:
+            f.seek(start_offset)
+            data = f.read(max_bytes)
+            next_offset = f.tell()
+
+        text = data.decode("utf-8", errors="replace")
+        lines = text.splitlines()
+
+        if regex:
+            flags = 0 if case_sensitive else re.IGNORECASE
+            pattern = re.compile(query, flags=flags)
+            def is_match(line: str) -> bool:
+                return pattern.search(line) is not None
+        else:
+            needle = query if case_sensitive else query.lower()
+            def is_match(line: str) -> bool:
+                haystack = line if case_sensitive else line.lower()
+                return needle in haystack
+
+        matches = []
+        for idx, line in enumerate(lines):
+            if not is_match(line):
+                continue
+            start_idx = max(0, idx - max(0, before))
+            end_idx = min(len(lines), idx + max(0, after) + 1)
+            context = lines[start_idx:end_idx]
+            matches.append({
+                "line_index": idx,
+                "context_start": start_idx,
+                "context_end": end_idx,
+                "context": context,
+            })
+            if len(matches) >= max_matches:
+                break
+
+        truncated = len(matches) >= max_matches
+        return json.dumps({
+            "path": path,
+            "query": query,
+            "start_offset": start_offset,
+            "next_offset": next_offset,
+            "truncated": truncated,
+            "matches": matches,
+        })
+    except FileNotFoundError:
+        return json.dumps({
+            "path": path,
+            "query": query,
+            "start_offset": start_offset,
+            "next_offset": start_offset,
+            "truncated": False,
+            "matches": [],
+        })
+    except Exception as e:
+        return json.dumps({
+            "path": path,
+            "query": query,
+            "start_offset": start_offset,
+            "next_offset": start_offset,
+            "truncated": False,
+            "matches": [],
+            "error": f"ERROR: {e}",
+        })
+
+
+@mcp.tool()
+def get_data(start: int = 0, count: int = 50) -> str:
+    """
+    Returns a slice of the active dataset as a JSON-formatted list of dictionaries.
+
+    Use this to inspect the actual data values in memory. Useful for checking data quality or content.
+    
+    Args:
+        start: The zero-based index of the first observation to retrieve.
+        count: The number of observations to retrieve. Defaults to 50.
+    """
+    data = client.get_data(start, count)
+    resp = DataResponse(start=start, count=count, data=data)
+    return resp.model_dump_json()
+
+
+@mcp.tool()
+def get_ui_channel() -> str:
+    """Return localhost HTTP endpoint + bearer token for the extension UI data plane."""
+    info = ui_channel.get_channel()
+    payload = {
+        "baseUrl": info.base_url,
+        "token": info.token,
+        "expiresAt": info.expires_at,
+        "capabilities": ui_channel.capabilities(),
+    }
+    return json.dumps(payload)
+
+@mcp.tool()
+def describe() -> str:
+    """
+    Returns variable descriptions, storage types, and labels (equivalent to Stata's `describe` command).
+
+    Use this to understand the structure of the dataset, variable names, and their formats before running analysis.
+    """
+    result = client.run_command_structured("describe", echo=True)
+    if result.success:
+        return result.stdout
+    if result.error:
+        return result.error.message
+    return ""
+
+@mcp.tool()
+def list_graphs() -> str:
+    """
+    Lists the names of all graphs currently stored in Stata's memory.
+
+    Use this to see which graphs are available for export via `export_graph`. The
+    response marks the active graph so the agent knows which one will export by
+    default.
+    """
+    graphs = client.list_graphs_structured()
+    return graphs.model_dump_json()
+
+@mcp.tool()
+def export_graph(graph_name: str = None, format: str = "pdf") -> str:
+    """
+    Exports a stored Stata graph to a file and returns its path.
+    
+    Args:
+        graph_name: The name of the graph to export (as seen in `list_graphs`). 
+                   If None, exports the currently active graph.
+        format: Output format, defaults to "pdf". Supported: "pdf", "png". Use
+                "png" to view the plot directly so the agent can visually check
+                titles, labels, legends, colors, and other user requirements.
+    """
+    try:
+        return client.export_graph(graph_name, format=format)
+    except Exception as e:
+        raise RuntimeError(f"Failed to export graph: {e}")
+
+@mcp.tool()
+def get_help(topic: str, plain_text: bool = False) -> str:
+    """
+    Returns the official Stata help text for a given command or topic.
+
+    Args:
+        topic: The command name or help topic (e.g., "regress", "graph", "options").
+               Returns Markdown by default, or plain text when plain_text=True.
+    """
+    return client.get_help(topic, plain_text=plain_text)
+
+@mcp.tool()
+def get_stored_results() -> str:
+    """
+    Returns the current stored results (r-class and e-class scalars/macros) as a JSON-formatted string.
+
+    Use this after running a command (like `summarize` or `regress`) to programmatically retrieve
+    specific values (e.g., means, coefficients, sample sizes) for validation or further calculation.
+    """
+    import json
+    return json.dumps(client.get_stored_results())
+
+@mcp.tool()
+def load_data(source: str, clear: bool = True, as_json: bool = True, raw: bool = False, max_output_lines: int = None) -> str:
+    """
+    Loads data using sysuse/webuse/use heuristics based on the source string.
+    Automatically appends , clear unless clear=False.
+
+    Args:
+        source: Dataset source (e.g., "auto", "auto.dta", "/path/to/file.dta").
+        clear: If True, clears data in memory before loading.
+        as_json: If True, returns JSON envelope.
+        raw: If True, returns raw output only.
+        max_output_lines: If set, truncates stdout to this many lines for token efficiency.
+    """
+    result = client.load_data(source, clear=clear, max_output_lines=max_output_lines)
+    ui_channel.notify_potential_dataset_change()
+    if raw:
+        return result.stdout if result.success else (result.error.message if result.error else result.stdout)
+    return result.model_dump_json()
+
+@mcp.tool()
+def codebook(variable: str, as_json: bool = True, trace: bool = False, raw: bool = False, max_output_lines: int = None) -> str:
+    """
+    Returns codebook/summary for a specific variable.
+
+    Args:
+        variable: The variable name to analyze.
+        as_json: If True, returns JSON envelope.
+        trace: If True, enables trace mode.
+        raw: If True, returns raw output only.
+        max_output_lines: If set, truncates stdout to this many lines for token efficiency.
+    """
+    result = client.codebook(variable, trace=trace, max_output_lines=max_output_lines)
+    if raw:
+        return result.stdout if result.success else (result.error.message if result.error else result.stdout)
+    return result.model_dump_json()
+
+@mcp.tool()
+async def run_do_file(
+    path: str,
+    ctx: Context | None = None,
+    echo: bool = True,
+    as_json: bool = True,
+    trace: bool = False,
+    raw: bool = False,
+    max_output_lines: int = None,
+    cwd: str | None = None,
+) -> str:
+    """
+    Executes a .do file.
+
+    Stata output is written to a temporary log file on disk.
+    The server emits a single `notifications/logMessage` event containing the log file path
+    (JSON payload: {"event":"log_path","path":"..."}) so the client can tail it locally.
+    If the client supplies a progress callback/token, progress updates are emitted via
+    `notifications/progress`.
+
+    Args:
+        path: Path to the .do file to execute.
+        ctx: FastMCP-injected request context (used to send MCP notifications). Optional for direct Python calls.
+        echo: If True, includes command in output.
+        as_json: If True, returns JSON envelope.
+        trace: If True, enables trace mode.
+        raw: If True, returns raw output only.
+        max_output_lines: If set, truncates stdout to this many lines for token efficiency.
+        Note: This tool always uses log-file streaming semantics; there is no non-streaming mode.
+    """
+    session = ctx.request_context.session if ctx is not None else None
+
+    async def notify_log(text: str) -> None:
+        if session is None:
+            return
+        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:
+        progress_token = ctx.request_context.meta.progressToken
+
+    async def notify_progress(progress: float, total: float | None, message: str | None) -> None:
+        if session is None or progress_token is None:
+            return
+        await session.send_progress_notification(
+            progress_token=progress_token,
+            progress=progress,
+            total=total,
+            message=message,
+            related_request_id=ctx.request_id,
+        )
+
+    async def _noop_log(_text: str) -> None:
+        return
+
+    result = await client.run_do_file_streaming(
+        path,
+        notify_log=notify_log if session is not None else _noop_log,
+        notify_progress=notify_progress if progress_token is not None else None,
+        echo=echo,
+        trace=trace,
+        max_output_lines=max_output_lines,
+        cwd=cwd,
+        emit_graph_ready=True,
+        graph_ready_task_id=ctx.request_id if ctx else None,
+        graph_ready_format="svg",
+    )
+
+    ui_channel.notify_potential_dataset_change()
+
+    if raw:
+        if result.success:
+            return result.log_path or ""
+        if result.error:
+            return result.error.message
+        return result.log_path or ""
+    return result.model_dump_json()
+
+@mcp.resource("stata://data/summary")
+def get_summary() -> str:
+    """
+    Returns the output of the `summarize` command for the dataset in memory.
+    Provides descriptive statistics (obs, mean, std. dev, min, max) for all variables.
+    """
+    result = client.run_command_structured("summarize", echo=True)
+    if result.success:
+        return result.stdout
+    if result.error:
+        return result.error.message
+    return ""
+
+@mcp.resource("stata://data/metadata")
+def get_metadata() -> str:
+    """
+    Returns the output of the `describe` command.
+    Provides metadata about the dataset, including variable names, storage types, display formats, and labels.
+    """
+    result = client.run_command_structured("describe", echo=True)
+    if result.success:
+        return result.stdout
+    if result.error:
+        return result.error.message
+    return ""
+
+@mcp.resource("stata://graphs/list")
+def list_graphs_resource() -> str:
+    """Resource wrapper for the graph list (uses tool list_graphs)."""
+    return list_graphs()
+
+@mcp.tool()
+def get_variable_list() -> str:
+    """Returns JSON list of all variables."""
+    variables = client.list_variables_structured()
+    return variables.model_dump_json()
+
+@mcp.resource("stata://variables/list")
+def get_variable_list_resource() -> str:
+    """Resource wrapper for the variable list."""
+    return get_variable_list()
+
+@mcp.resource("stata://results/stored")
+def get_stored_results_resource() -> str:
+    """Returns stored r() and e() results."""
+    import json
+    return json.dumps(client.get_stored_results())
+
+@mcp.tool()
+def export_graphs_all() -> str:
+    """
+    Exports all graphs in memory to file paths.
+
+    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()
+    return exports.model_dump_json(exclude_none=False)
+
+def main():
+    if "--version" in sys.argv:
+        try:
+            from importlib.metadata import version
+            print(version("mcp-stata"))
+        except Exception:
+            print("unknown")
+        return
+
+    setup_logging()
+    
+    # Initialize Stata here on the main thread to ensure any issues are logged early.
+    # On Windows, this is critical for COM registration. On other platforms, it helps
+    # catch license or installation errors before the first tool call.
+    try:
+        client.init()
+    except BaseException as e:
+        # Use sys.stderr.write and flush to ensure visibility before exit
+        msg = f"\n{'='*60}\n[mcp_stata] FATAL: STATA INITIALIZATION FAILED\n{'='*60}\nError: {repr(e)}\n"
+        sys.stderr.write(msg)
+        if isinstance(e, SystemExit):
+            sys.stderr.write(f"Stata triggered a SystemExit (code: {e.code}). This is usually a license error.\n")
+        sys.stderr.write(f"{'='*60}\n\n")
+        sys.stderr.flush()
+        
+        # We exit here because the user wants a clear failure when Stata cannot be loaded.
+        sys.exit(1)
+
+    mcp.run()
+
+if __name__ == "__main__":
+    main()