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/config.py +20 -0
- mcp_stata/discovery.py +134 -59
- mcp_stata/graph_detector.py +29 -26
- mcp_stata/models.py +3 -0
- mcp_stata/server.py +647 -19
- mcp_stata/stata_client.py +1881 -989
- mcp_stata/streaming_io.py +3 -1
- mcp_stata/test_stata.py +54 -0
- mcp_stata/ui_http.py +178 -19
- {mcp_stata-1.7.3.dist-info → mcp_stata-1.13.0.dist-info}/METADATA +15 -3
- mcp_stata-1.13.0.dist-info/RECORD +16 -0
- mcp_stata-1.7.3.dist-info/RECORD +0 -14
- {mcp_stata-1.7.3.dist-info → mcp_stata-1.13.0.dist-info}/WHEEL +0 -0
- {mcp_stata-1.7.3.dist-info → mcp_stata-1.13.0.dist-info}/entry_points.txt +0 -0
- {mcp_stata-1.7.3.dist-info → mcp_stata-1.13.0.dist-info}/licenses/LICENSE +0 -0
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
|
-
|
|
20
|
-
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
|
|
24
|
-
|
|
25
|
-
|
|
26
|
-
|
|
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
|
-
|
|
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
|
-
|
|
422
|
-
|
|
423
|
-
|
|
424
|
-
|
|
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()
|