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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of mcp-stata might be problematic. Click here for more details.

mcp_stata/server.py CHANGED
@@ -1,4 +1,7 @@
1
1
  import anyio
2
+ import asyncio
3
+ from dataclasses import dataclass
4
+ from datetime import datetime
2
5
  from importlib.metadata import PackageNotFoundError, version
3
6
  from mcp.server.fastmcp import Context, FastMCP
4
7
  import mcp.types as types
@@ -10,26 +13,513 @@ from .models import (
10
13
  GraphExportResponse,
11
14
  )
12
15
  import logging
16
+ import sys
13
17
  import json
14
18
  import os
19
+ import re
20
+ import traceback
21
+ import uuid
22
+ from functools import wraps
23
+ from typing import Optional, Dict
15
24
 
16
25
  from .ui_http import UIChannelManager
17
26
 
18
27
 
19
- LOG_LEVEL = os.getenv("MCP_STATA_LOGLEVEL", "INFO").upper()
20
- logging.basicConfig(level=LOG_LEVEL, format="%(asctime)s %(levelname)s %(name)s - %(message)s")
21
- try:
22
- _mcp_stata_version = version("mcp-stata")
23
- except PackageNotFoundError:
24
- _mcp_stata_version = "unknown"
25
- logging.info("mcp-stata version: %s", _mcp_stata_version)
26
- logging.info("STATA_PATH env at startup: %s", os.getenv("STATA_PATH", "<not set>"))
28
+ # Configure logging
29
+ logger = logging.getLogger("mcp_stata")
30
+
31
+ def setup_logging():
32
+ # Configure logging to stderr with immediate flush for MCP transport
33
+ log_level = os.getenv("MCP_STATA_LOGLEVEL", "DEBUG").upper()
34
+ configure_root = os.getenv("MCP_STATA_CONFIGURE_LOGGING", "1").lower() not in {"0", "false", "no"}
35
+
36
+ # Create a handler that flushes immediately
37
+ handler = logging.StreamHandler(sys.stderr)
38
+ handler.setLevel(getattr(logging, log_level, logging.DEBUG))
39
+ handler.setFormatter(logging.Formatter("[%(name)s] %(levelname)s: %(message)s"))
40
+
41
+ # Configure root logger only if requested; avoid clobbering existing handlers.
42
+ if configure_root:
43
+ root_logger = logging.getLogger()
44
+ if not root_logger.handlers:
45
+ root_logger.addHandler(handler)
46
+ root_logger.setLevel(getattr(logging, log_level, logging.DEBUG))
47
+
48
+ # Also configure the mcp_stata logger explicitly without duplicating handlers.
49
+ if logger.level == logging.NOTSET:
50
+ logger.setLevel(getattr(logging, log_level, logging.DEBUG))
51
+ if not logger.handlers:
52
+ logger.addHandler(handler)
53
+ logger.propagate = False
54
+
55
+ try:
56
+ _mcp_stata_version = version("mcp-stata")
57
+ except PackageNotFoundError:
58
+ _mcp_stata_version = "unknown"
59
+
60
+ logger.info("=== mcp-stata server starting ===")
61
+ logger.info("mcp-stata version: %s", _mcp_stata_version)
62
+ logger.info("STATA_PATH env at startup: %s", os.getenv("STATA_PATH", "<not set>"))
63
+ logger.info("LOG_LEVEL: %s", log_level)
27
64
 
28
65
  # Initialize FastMCP
29
66
  mcp = FastMCP("mcp_stata")
30
67
  client = StataClient()
31
68
  ui_channel = UIChannelManager(client)
32
69
 
70
+
71
+ @dataclass
72
+ class BackgroundTask:
73
+ task_id: str
74
+ kind: str
75
+ task: asyncio.Task
76
+ created_at: datetime
77
+ log_path: Optional[str] = None
78
+ result: Optional[str] = None
79
+ error: Optional[str] = None
80
+ done: bool = False
81
+
82
+
83
+ _background_tasks: Dict[str, BackgroundTask] = {}
84
+
85
+
86
+ def _register_task(task_info: BackgroundTask, max_tasks: int = 100) -> None:
87
+ _background_tasks[task_info.task_id] = task_info
88
+ if len(_background_tasks) <= max_tasks:
89
+ return
90
+ completed = [task for task in _background_tasks.values() if task.done]
91
+ completed.sort(key=lambda item: item.created_at)
92
+ for task in completed[: max(0, len(_background_tasks) - max_tasks)]:
93
+ _background_tasks.pop(task.task_id, None)
94
+
95
+
96
+ def _format_command_result(result, raw: bool, as_json: bool) -> str:
97
+ if raw:
98
+ if result.success:
99
+ return result.log_path or ""
100
+ if result.error:
101
+ msg = result.error.message
102
+ if result.error.rc is not None:
103
+ msg = f"{msg}\nrc={result.error.rc}"
104
+ return msg
105
+ return result.log_path or ""
106
+ if as_json:
107
+ return result.model_dump_json()
108
+ return result.model_dump_json()
109
+
110
+
111
+ async def _wait_for_log_path(task_info: BackgroundTask) -> None:
112
+ while task_info.log_path is None and not task_info.done:
113
+ await anyio.sleep(0.01)
114
+
115
+
116
+ async def _notify_task_done(session: object | None, task_info: BackgroundTask, request_id: object | None) -> None:
117
+ if session is None:
118
+ return
119
+ payload = {
120
+ "event": "task_done",
121
+ "task_id": task_info.task_id,
122
+ "status": "done" if task_info.done else "unknown",
123
+ "log_path": task_info.log_path,
124
+ "error": task_info.error,
125
+ }
126
+ try:
127
+ _debug_notification("logMessage", payload, request_id)
128
+ await session.send_log_message(level="info", data=json.dumps(payload), related_request_id=request_id)
129
+ except Exception:
130
+ return
131
+
132
+
133
+ def _debug_notification(kind: str, payload: object, request_id: object | None = None) -> None:
134
+ try:
135
+ serialized = payload if isinstance(payload, str) else json.dumps(payload, ensure_ascii=False)
136
+ except Exception:
137
+ serialized = str(payload)
138
+ logger.debug("MCP notify %s request_id=%s payload=%s", kind, request_id, serialized)
139
+
140
+
141
+ async def _notify_tool_error(ctx: Context | None, tool_name: str, exc: Exception) -> None:
142
+ if ctx is None:
143
+ return
144
+ session = ctx.request_context.session
145
+ if session is None:
146
+ return
147
+ task_id = None
148
+ meta = ctx.request_context.meta
149
+ if meta is not None:
150
+ task_id = getattr(meta, "task_id", None) or getattr(meta, "taskId", None)
151
+ payload = {
152
+ "event": "tool_error",
153
+ "tool": tool_name,
154
+ "error": str(exc),
155
+ "traceback": traceback.format_exc(),
156
+ }
157
+ if task_id is not None:
158
+ payload["task_id"] = task_id
159
+ try:
160
+ _debug_notification("logMessage", payload, ctx.request_id)
161
+ await session.send_log_message(
162
+ level="error",
163
+ data=json.dumps(payload),
164
+ related_request_id=ctx.request_id,
165
+ )
166
+ except Exception:
167
+ logger.exception("Failed to emit tool_error notification for %s", tool_name)
168
+
169
+
170
+ def _log_tool_call(tool_name: str, ctx: Context | None = None) -> None:
171
+ request_id = None
172
+ if ctx is not None:
173
+ request_id = getattr(ctx, "request_id", None)
174
+ logger.info("MCP tool call: %s request_id=%s", tool_name, request_id)
175
+
176
+
177
+ def _attach_task_id(ctx: Context | None, task_id: str) -> None:
178
+ if ctx is None:
179
+ return
180
+ meta = ctx.request_context.meta
181
+ if meta is None:
182
+ meta = types.RequestParams.Meta()
183
+ ctx.request_context.meta = meta
184
+ try:
185
+ setattr(meta, "task_id", task_id)
186
+ except Exception:
187
+ logger.debug("Unable to attach task_id to request meta", exc_info=True)
188
+
189
+
190
+ def _extract_ctx(args: tuple[object, ...], kwargs: dict[str, object]) -> Context | None:
191
+ ctx = kwargs.get("ctx")
192
+ if isinstance(ctx, Context):
193
+ return ctx
194
+ for arg in args:
195
+ if isinstance(arg, Context):
196
+ return arg
197
+ return None
198
+
199
+
200
+ _mcp_tool = mcp.tool
201
+ _mcp_resource = mcp.resource
202
+
203
+
204
+ def tool(*tool_args, **tool_kwargs):
205
+ decorator = _mcp_tool(*tool_args, **tool_kwargs)
206
+
207
+ def outer(func):
208
+ if asyncio.iscoroutinefunction(func):
209
+ @wraps(func)
210
+ async def async_inner(*args, **kwargs):
211
+ ctx = _extract_ctx(args, kwargs)
212
+ _log_tool_call(func.__name__, ctx)
213
+ try:
214
+ return await func(*args, **kwargs)
215
+ except Exception as exc:
216
+ await _notify_tool_error(ctx, func.__name__, exc)
217
+ raise
218
+
219
+ return decorator(async_inner)
220
+
221
+ @wraps(func)
222
+ def sync_inner(*args, **kwargs):
223
+ ctx = _extract_ctx(args, kwargs)
224
+ _log_tool_call(func.__name__, ctx)
225
+ try:
226
+ return func(*args, **kwargs)
227
+ except Exception:
228
+ logger.exception("Tool %s failed", func.__name__)
229
+ raise
230
+
231
+ return decorator(sync_inner)
232
+
233
+ return outer
234
+
235
+
236
+ mcp.tool = tool
237
+
238
+
239
+ def resource(*resource_args, **resource_kwargs):
240
+ decorator = _mcp_resource(*resource_args, **resource_kwargs)
241
+
242
+ def outer(func):
243
+ if asyncio.iscoroutinefunction(func):
244
+ @wraps(func)
245
+ async def async_inner(*args, **kwargs):
246
+ _log_tool_call(func.__name__, _extract_ctx(args, kwargs))
247
+ return await func(*args, **kwargs)
248
+
249
+ return decorator(async_inner)
250
+
251
+ @wraps(func)
252
+ def sync_inner(*args, **kwargs):
253
+ _log_tool_call(func.__name__, _extract_ctx(args, kwargs))
254
+ return func(*args, **kwargs)
255
+
256
+ return decorator(sync_inner)
257
+
258
+ return outer
259
+
260
+
261
+ mcp.resource = resource
262
+
263
+
264
+ @mcp.tool()
265
+ async def run_do_file_background(
266
+ path: str,
267
+ ctx: Context | None = None,
268
+ echo: bool = True,
269
+ as_json: bool = True,
270
+ trace: bool = False,
271
+ raw: bool = False,
272
+ max_output_lines: int = None,
273
+ cwd: str | None = None,
274
+ ) -> str:
275
+ """Run a Stata do-file in the background and return a task id.
276
+
277
+ Notifications:
278
+ - logMessage: {"event":"log_path","path":"..."}
279
+ - logMessage: {"event":"task_done","task_id":"...","status":"done","log_path":"...","error":null}
280
+ """
281
+ session = ctx.request_context.session if ctx is not None else None
282
+ request_id = ctx.request_id if ctx is not None else None
283
+ task_id = uuid.uuid4().hex
284
+ _attach_task_id(ctx, task_id)
285
+ task_info = BackgroundTask(
286
+ task_id=task_id,
287
+ kind="do_file",
288
+ task=None,
289
+ created_at=datetime.utcnow(),
290
+ )
291
+
292
+ async def notify_log(text: str) -> None:
293
+ if session is not None:
294
+ _debug_notification("logMessage", text, ctx.request_id)
295
+ await session.send_log_message(level="info", data=text, related_request_id=ctx.request_id)
296
+ try:
297
+ payload = json.loads(text)
298
+ if isinstance(payload, dict) and payload.get("event") == "log_path":
299
+ task_info.log_path = payload.get("path")
300
+ except Exception:
301
+ return
302
+
303
+ progress_token = None
304
+ if ctx is not None and ctx.request_context.meta is not None:
305
+ progress_token = ctx.request_context.meta.progressToken
306
+
307
+ async def notify_progress(progress: float, total: float | None, message: str | None) -> None:
308
+ if session is None or progress_token is None:
309
+ return
310
+ _debug_notification(
311
+ "progress",
312
+ {"progress": progress, "total": total, "message": message},
313
+ ctx.request_id,
314
+ )
315
+ await session.send_progress_notification(
316
+ progress_token=progress_token,
317
+ progress=progress,
318
+ total=total,
319
+ message=message,
320
+ related_request_id=ctx.request_id,
321
+ )
322
+
323
+ async def _run() -> None:
324
+ try:
325
+ result = await client.run_do_file_streaming(
326
+ path,
327
+ notify_log=notify_log,
328
+ notify_progress=notify_progress if progress_token is not None else None,
329
+ echo=echo,
330
+ trace=trace,
331
+ max_output_lines=max_output_lines,
332
+ cwd=cwd,
333
+ emit_graph_ready=True,
334
+ graph_ready_task_id=task_id,
335
+ graph_ready_format="svg",
336
+ )
337
+ ui_channel.notify_potential_dataset_change()
338
+ task_info.result = _format_command_result(result, raw=raw, as_json=as_json)
339
+ if result.error:
340
+ task_info.error = result.error.message
341
+ except Exception as exc: # pragma: no cover - defensive
342
+ task_info.error = str(exc)
343
+ finally:
344
+ task_info.done = True
345
+ await _notify_task_done(session, task_info, request_id)
346
+
347
+ task_info.task = asyncio.create_task(_run())
348
+ _register_task(task_info)
349
+ await _wait_for_log_path(task_info)
350
+ return json.dumps({"task_id": task_id, "status": "started", "log_path": task_info.log_path})
351
+
352
+
353
+ @mcp.tool()
354
+ def get_task_status(task_id: str, allow_polling: bool = False) -> str:
355
+ """Return task status for background executions.
356
+
357
+ Polling is disabled by default; set allow_polling=True for legacy callers.
358
+ """
359
+ notice = "Prefer task_done logMessage notifications over polling get_task_status."
360
+ if not allow_polling:
361
+ logger.warning(
362
+ "get_task_status called without allow_polling; clients must use task_done logMessage notifications"
363
+ )
364
+ return json.dumps({
365
+ "task_id": task_id,
366
+ "status": "polling_not_allowed",
367
+ "error": "Polling is disabled; use task_done logMessage notifications.",
368
+ "notice": notice,
369
+ })
370
+ logger.warning("get_task_status called; clients should use task_done logMessage notifications instead of polling")
371
+ task_info = _background_tasks.get(task_id)
372
+ if task_info is None:
373
+ return json.dumps({"task_id": task_id, "status": "not_found", "notice": notice})
374
+ return json.dumps({
375
+ "task_id": task_id,
376
+ "status": "done" if task_info.done else "running",
377
+ "kind": task_info.kind,
378
+ "created_at": task_info.created_at.isoformat(),
379
+ "log_path": task_info.log_path,
380
+ "error": task_info.error,
381
+ "notice": notice,
382
+ })
383
+
384
+
385
+ @mcp.tool()
386
+ def get_task_result(task_id: str, allow_polling: bool = False) -> str:
387
+ """Return task result for background executions.
388
+
389
+ Polling is disabled by default; set allow_polling=True for legacy callers.
390
+ """
391
+ notice = "Prefer task_done logMessage notifications over polling get_task_result."
392
+ if not allow_polling:
393
+ logger.warning(
394
+ "get_task_result called without allow_polling; clients must use task_done logMessage notifications"
395
+ )
396
+ return json.dumps({
397
+ "task_id": task_id,
398
+ "status": "polling_not_allowed",
399
+ "error": "Polling is disabled; use task_done logMessage notifications.",
400
+ "notice": notice,
401
+ })
402
+ logger.warning("get_task_result called; clients should use task_done logMessage notifications instead of polling")
403
+ task_info = _background_tasks.get(task_id)
404
+ if task_info is None:
405
+ return json.dumps({"task_id": task_id, "status": "not_found", "notice": notice})
406
+ if not task_info.done:
407
+ return json.dumps({
408
+ "task_id": task_id,
409
+ "status": "running",
410
+ "log_path": task_info.log_path,
411
+ "notice": notice,
412
+ })
413
+ return json.dumps({
414
+ "task_id": task_id,
415
+ "status": "done",
416
+ "log_path": task_info.log_path,
417
+ "error": task_info.error,
418
+ "notice": notice,
419
+ "result": task_info.result,
420
+ })
421
+
422
+
423
+ @mcp.tool()
424
+ def cancel_task(task_id: str) -> str:
425
+ """Request cancellation of a background task."""
426
+ task_info = _background_tasks.get(task_id)
427
+ if task_info is None:
428
+ return json.dumps({"task_id": task_id, "status": "not_found"})
429
+ if task_info.task and not task_info.task.done():
430
+ task_info.task.cancel()
431
+ return json.dumps({"task_id": task_id, "status": "cancelling"})
432
+ return json.dumps({"task_id": task_id, "status": "done", "log_path": task_info.log_path})
433
+
434
+
435
+ @mcp.tool()
436
+ async def run_command_background(
437
+ code: str,
438
+ ctx: Context | None = None,
439
+ echo: bool = True,
440
+ as_json: bool = True,
441
+ trace: bool = False,
442
+ raw: bool = False,
443
+ max_output_lines: int = None,
444
+ cwd: str | None = None,
445
+ ) -> str:
446
+ """Run a Stata command in the background and return a task id.
447
+
448
+ Notifications:
449
+ - logMessage: {"event":"log_path","path":"..."}
450
+ - logMessage: {"event":"task_done","task_id":"...","status":"done","log_path":"...","error":null}
451
+ """
452
+ session = ctx.request_context.session if ctx is not None else None
453
+ request_id = ctx.request_id if ctx is not None else None
454
+ task_id = uuid.uuid4().hex
455
+ _attach_task_id(ctx, task_id)
456
+ task_info = BackgroundTask(
457
+ task_id=task_id,
458
+ kind="command",
459
+ task=None,
460
+ created_at=datetime.utcnow(),
461
+ )
462
+
463
+ async def notify_log(text: str) -> None:
464
+ if session is not None:
465
+ _debug_notification("logMessage", text, ctx.request_id)
466
+ await session.send_log_message(level="info", data=text, related_request_id=ctx.request_id)
467
+ try:
468
+ payload = json.loads(text)
469
+ if isinstance(payload, dict) and payload.get("event") == "log_path":
470
+ task_info.log_path = payload.get("path")
471
+ except Exception:
472
+ return
473
+
474
+ progress_token = None
475
+ if ctx is not None and ctx.request_context.meta is not None:
476
+ progress_token = ctx.request_context.meta.progressToken
477
+
478
+ async def notify_progress(progress: float, total: float | None, message: str | None) -> None:
479
+ if session is None or progress_token is None:
480
+ return
481
+ _debug_notification(
482
+ "progress",
483
+ {"progress": progress, "total": total, "message": message},
484
+ ctx.request_id,
485
+ )
486
+ await session.send_progress_notification(
487
+ progress_token=progress_token,
488
+ progress=progress,
489
+ total=total,
490
+ message=message,
491
+ related_request_id=ctx.request_id,
492
+ )
493
+
494
+ async def _run() -> None:
495
+ try:
496
+ result = await client.run_command_streaming(
497
+ code,
498
+ notify_log=notify_log,
499
+ notify_progress=notify_progress if progress_token is not None else None,
500
+ echo=echo,
501
+ trace=trace,
502
+ max_output_lines=max_output_lines,
503
+ cwd=cwd,
504
+ emit_graph_ready=True,
505
+ graph_ready_task_id=task_id,
506
+ graph_ready_format="svg",
507
+ )
508
+ ui_channel.notify_potential_dataset_change()
509
+ task_info.result = _format_command_result(result, raw=raw, as_json=as_json)
510
+ if result.error:
511
+ task_info.error = result.error.message
512
+ except Exception as exc: # pragma: no cover - defensive
513
+ task_info.error = str(exc)
514
+ finally:
515
+ task_info.done = True
516
+ await _notify_task_done(session, task_info, request_id)
517
+
518
+ task_info.task = asyncio.create_task(_run())
519
+ _register_task(task_info)
520
+ await _wait_for_log_path(task_info)
521
+ return json.dumps({"task_id": task_id, "status": "started", "log_path": task_info.log_path})
522
+
33
523
  @mcp.tool()
34
524
  async def run_command(
35
525
  code: str,
@@ -68,6 +558,7 @@ async def run_command(
68
558
  async def notify_log(text: str) -> None:
69
559
  if session is None:
70
560
  return
561
+ _debug_notification("logMessage", text, ctx.request_id)
71
562
  await session.send_log_message(level="info", data=text, related_request_id=ctx.request_id)
72
563
 
73
564
  progress_token = None
@@ -77,6 +568,11 @@ async def run_command(
77
568
  async def notify_progress(progress: float, total: float | None, message: str | None) -> None:
78
569
  if session is None or progress_token is None:
79
570
  return
571
+ _debug_notification(
572
+ "progress",
573
+ {"progress": progress, "total": total, "message": message},
574
+ ctx.request_id,
575
+ )
80
576
  await session.send_progress_notification(
81
577
  progress_token=progress_token,
82
578
  progress=progress,
@@ -96,6 +592,9 @@ async def run_command(
96
592
  trace=trace,
97
593
  max_output_lines=max_output_lines,
98
594
  cwd=cwd,
595
+ emit_graph_ready=True,
596
+ graph_ready_task_id=ctx.request_id if ctx else None,
597
+ graph_ready_format="svg",
99
598
  )
100
599
 
101
600
  # Conservative invalidation: arbitrary Stata commands may change data.
@@ -142,6 +641,111 @@ def read_log(path: str, offset: int = 0, max_bytes: int = 65536) -> str:
142
641
  return json.dumps({"path": path, "offset": offset, "next_offset": offset, "data": f"ERROR: {e}"})
143
642
 
144
643
 
644
+ @mcp.tool()
645
+ def find_in_log(
646
+ path: str,
647
+ query: str,
648
+ start_offset: int = 0,
649
+ max_bytes: int = 5_000_000,
650
+ before: int = 2,
651
+ after: int = 2,
652
+ case_sensitive: bool = False,
653
+ regex: bool = False,
654
+ max_matches: int = 50,
655
+ ) -> str:
656
+ """Find text within a log file and return context windows.
657
+
658
+ Args:
659
+ path: Absolute path to the log file previously provided by the server.
660
+ query: Text or regex pattern to search for.
661
+ start_offset: Byte offset to start searching from.
662
+ max_bytes: Maximum bytes to read from the log.
663
+ before: Number of context lines to include before each match.
664
+ after: Number of context lines to include after each match.
665
+ case_sensitive: If True, match case-sensitively.
666
+ regex: If True, treat query as a regular expression.
667
+ max_matches: Maximum number of matches to return.
668
+
669
+ Returns a JSON string with matches and offsets:
670
+ {"path":..., "query":..., "start_offset":..., "next_offset":..., "truncated":..., "matches":[...]}.
671
+ """
672
+ try:
673
+ if start_offset < 0:
674
+ start_offset = 0
675
+ if max_bytes <= 0:
676
+ return json.dumps({
677
+ "path": path,
678
+ "query": query,
679
+ "start_offset": start_offset,
680
+ "next_offset": start_offset,
681
+ "truncated": False,
682
+ "matches": [],
683
+ })
684
+ with open(path, "rb") as f:
685
+ f.seek(start_offset)
686
+ data = f.read(max_bytes)
687
+ next_offset = f.tell()
688
+
689
+ text = data.decode("utf-8", errors="replace")
690
+ lines = text.splitlines()
691
+
692
+ if regex:
693
+ flags = 0 if case_sensitive else re.IGNORECASE
694
+ pattern = re.compile(query, flags=flags)
695
+ def is_match(line: str) -> bool:
696
+ return pattern.search(line) is not None
697
+ else:
698
+ needle = query if case_sensitive else query.lower()
699
+ def is_match(line: str) -> bool:
700
+ haystack = line if case_sensitive else line.lower()
701
+ return needle in haystack
702
+
703
+ matches = []
704
+ for idx, line in enumerate(lines):
705
+ if not is_match(line):
706
+ continue
707
+ start_idx = max(0, idx - max(0, before))
708
+ end_idx = min(len(lines), idx + max(0, after) + 1)
709
+ context = lines[start_idx:end_idx]
710
+ matches.append({
711
+ "line_index": idx,
712
+ "context_start": start_idx,
713
+ "context_end": end_idx,
714
+ "context": context,
715
+ })
716
+ if len(matches) >= max_matches:
717
+ break
718
+
719
+ truncated = len(matches) >= max_matches
720
+ return json.dumps({
721
+ "path": path,
722
+ "query": query,
723
+ "start_offset": start_offset,
724
+ "next_offset": next_offset,
725
+ "truncated": truncated,
726
+ "matches": matches,
727
+ })
728
+ except FileNotFoundError:
729
+ return json.dumps({
730
+ "path": path,
731
+ "query": query,
732
+ "start_offset": start_offset,
733
+ "next_offset": start_offset,
734
+ "truncated": False,
735
+ "matches": [],
736
+ })
737
+ except Exception as e:
738
+ return json.dumps({
739
+ "path": path,
740
+ "query": query,
741
+ "start_offset": start_offset,
742
+ "next_offset": start_offset,
743
+ "truncated": False,
744
+ "matches": [],
745
+ "error": f"ERROR: {e}",
746
+ })
747
+
748
+
145
749
  @mcp.tool()
146
750
  def get_data(start: int = 0, count: int = 50) -> str:
147
751
  """
@@ -306,6 +910,7 @@ async def run_do_file(
306
910
  async def notify_log(text: str) -> None:
307
911
  if session is None:
308
912
  return
913
+ _debug_notification("logMessage", text, ctx.request_id)
309
914
  await session.send_log_message(level="info", data=text, related_request_id=ctx.request_id)
310
915
 
311
916
  progress_token = None
@@ -315,6 +920,11 @@ async def run_do_file(
315
920
  async def notify_progress(progress: float, total: float | None, message: str | None) -> None:
316
921
  if session is None or progress_token is None:
317
922
  return
923
+ _debug_notification(
924
+ "progress",
925
+ {"progress": progress, "total": total, "message": message},
926
+ ctx.request_id,
927
+ )
318
928
  await session.send_progress_notification(
319
929
  progress_token=progress_token,
320
930
  progress=progress,
@@ -334,6 +944,9 @@ async def run_do_file(
334
944
  trace=trace,
335
945
  max_output_lines=max_output_lines,
336
946
  cwd=cwd,
947
+ emit_graph_ready=True,
948
+ graph_ready_task_id=ctx.request_id if ctx else None,
949
+ graph_ready_format="svg",
337
950
  )
338
951
 
339
952
  ui_channel.notify_potential_dataset_change()
@@ -411,19 +1024,34 @@ def export_graphs_all(use_base64: bool = False) -> str:
411
1024
  return exports.model_dump_json(exclude_none=False)
412
1025
 
413
1026
  def main():
414
- # On Windows, Stata automation relies on COM, which is sensitive to threading models.
415
- # The FastMCP server executes tool calls in a thread pool. If Stata is initialized
416
- # lazily inside a worker thread, it may fail or hang due to COM/UI limitations.
417
- # We explicitly initialize Stata here on the main thread to ensure the COM server
418
- # is properly registered and accessible.
419
- if os.name == "nt":
1027
+ if "--version" in sys.argv:
420
1028
  try:
421
- client.init()
422
- except Exception as e:
423
- # Log error but let the server start; specific tools will fail gracefully later
424
- logging.error(f"Stata initialization failed: {e}")
1029
+ from importlib.metadata import version
1030
+ print(version("mcp-stata"))
1031
+ except Exception:
1032
+ print("unknown")
1033
+ return
1034
+
1035
+ setup_logging()
1036
+
1037
+ # Initialize Stata here on the main thread to ensure any issues are logged early.
1038
+ # On Windows, this is critical for COM registration. On other platforms, it helps
1039
+ # catch license or installation errors before the first tool call.
1040
+ try:
1041
+ client.init()
1042
+ except BaseException as e:
1043
+ # Use sys.stderr.write and flush to ensure visibility before exit
1044
+ msg = f"\n{'='*60}\n[mcp_stata] FATAL: STATA INITIALIZATION FAILED\n{'='*60}\nError: {repr(e)}\n"
1045
+ sys.stderr.write(msg)
1046
+ if isinstance(e, SystemExit):
1047
+ sys.stderr.write(f"Stata triggered a SystemExit (code: {e.code}). This is usually a license error.\n")
1048
+ sys.stderr.write(f"{'='*60}\n\n")
1049
+ sys.stderr.flush()
1050
+
1051
+ # We exit here because the user wants a clear failure when Stata cannot be loaded.
1052
+ sys.exit(1)
425
1053
 
426
1054
  mcp.run()
427
1055
 
428
1056
  if __name__ == "__main__":
429
- main()
1057
+ main()