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

mcp_stata/server.py

@@ -1,6 +1,10 @@
 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 (
@@ -10,26 +14,575 @@ from .models import (
     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
 
 
-LOG_LEVEL = os.getenv("MCP_STATA_LOGLEVEL", "INFO").upper()
-logging.basicConfig(level=LOG_LEVEL, format="%(asctime)s %(levelname)s %(name)s - %(message)s")
-try:
-    _mcp_stata_version = version("mcp-stata")
-except PackageNotFoundError:
-    _mcp_stata_version = "unknown"
-logging.info("mcp-stata version: %s", _mcp_stata_version)
-logging.info("STATA_PATH env at startup: %s", os.getenv("STATA_PATH", "<not set>"))
+# 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,
@@ -68,7 +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:
@@ -96,6 +659,9 @@ async def run_command(
         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.
@@ -128,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:
@@ -142,6 +716,111 @@ def read_log(path: str, offset: int = 0, max_bytes: int = 65536) -> str:
         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:
     """
@@ -306,7 +985,16 @@ async def run_do_file(
     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:
@@ -334,6 +1022,9 @@ async def run_do_file(
         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()
@@ -395,35 +1086,45 @@ 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.
+    Exports all graphs in memory to file paths.
 
-    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.
-
-    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():
-    # On Windows, Stata automation relies on COM, which is sensitive to threading models.
-    # The FastMCP server executes tool calls in a thread pool. If Stata is initialized
-    # lazily inside a worker thread, it may fail or hang due to COM/UI limitations.
-    # We explicitly initialize Stata here on the main thread to ensure the COM server
-    # is properly registered and accessible.
-    if os.name == "nt":
+    if "--version" in sys.argv:
         try:
-            client.init()
-        except Exception as e:
-            # Log error but let the server start; specific tools will fail gracefully later
-            logging.error(f"Stata initialization failed: {e}")
+            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()
+    main()