diff --git a/mcpb/manifest.json b/mcpb/manifest.json index cea87a4..9a40836 100644 --- a/mcpb/manifest.json +++ b/mcpb/manifest.json @@ -51,11 +51,15 @@ }, { "name": "execute_code", - "description": "Execute Python code against the OpenSTAAD API in a sandboxed environment." + "description": "Execute Python code against the OpenSTAAD API in a sandboxed environment. Supports a progress callback and long-running/background execution." }, { "name": "get_status", - "description": "Check the connection to STAAD.Pro. Returns connection state, STAAD version, model path, and analysis status." + "description": "Check the connection to STAAD.Pro. Returns connection state, STAAD version, model path, analysis status, and whether the executor is busy." + }, + { + "name": "get_job_result", + "description": "Wait for a background `execute_code` job (started with `mode='poll'`) and return its progress or final result." } ], "tools_generated": false, diff --git a/mcpb/openstaad-mcp.spec b/mcpb/openstaad-mcp.spec index 8cb16ac..eb13982 100644 --- a/mcpb/openstaad-mcp.spec +++ b/mcpb/openstaad-mcp.spec @@ -11,7 +11,7 @@ STAAD skills content. import os from pathlib import Path -from PyInstaller.utils.hooks import copy_metadata +from PyInstaller.utils.hooks import collect_all, copy_metadata block_cipher = None @@ -30,11 +30,16 @@ if skills_dir.exists(): # Include distribution metadata so frozen builds can resolve it. package_metadata = copy_metadata("fastmcp") +# fastmcp bundles `docket`, whose memory client backend dynamically imports +# `burner_redis` (a native extension) via importlib at lifespan startup. +# PyInstaller's static analysis misses this, so collect it explicitly. +burner_datas, burner_binaries, burner_hiddenimports = collect_all("burner_redis") + a = Analysis( [str(ROOT / "src" / "openstaad_mcp" / "main.py")], pathex=[str(ROOT / "src")], - binaries=[], - datas=skills_data + package_metadata, + binaries=burner_binaries, + datas=skills_data + package_metadata + burner_datas, hiddenimports=[ "openstaad_mcp", "openstaad_mcp.server", @@ -45,7 +50,8 @@ a = Analysis( "openstaadpy.os_analytical", "uvicorn", "fastmcp", - ], + ] + + burner_hiddenimports, hookspath=[], hooksconfig={}, runtime_hooks=[], diff --git a/pyproject.toml b/pyproject.toml index aba216f..da9e3c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,7 +22,7 @@ classifiers = [ "License :: OSI Approved :: MIT License", ] dependencies = [ - "fastmcp>=3.2.3,<4", + "fastmcp[tasks]>=3.2.3,<4", "packaging>=24", "starlette>=1.0.0,<2.0.0", "pydantic>=2,<3", diff --git a/src/openstaad_mcp/execution.py b/src/openstaad_mcp/execution.py new file mode 100644 index 0000000..d95d07b --- /dev/null +++ b/src/openstaad_mcp/execution.py @@ -0,0 +1,472 @@ +""" +--------------------------------------------------------------------------------------------- +Copyright (c) Bentley Systems, Incorporated. All rights reserved. +See LICENSE.md in the project root for license terms and full copyright notice. +--------------------------------------------------------------------------------------------- + +Code-execution orchestration layer. + +Keeps the MCP tools in :mod:`openstaad_mcp.server` thin: they simply delegate +to :class:`ExecutionService`, which owns instance resolution, file I/O wiring, +sandbox dispatch and the two long-running-work strategies below. The public +``execute_code`` ``mode`` parameter exposes three values that map onto these two +strategies: ``"native"`` and ``"poll"`` select one explicitly, while ``"auto"`` +(the default) picks one per request from the signals described in :meth:`ExecutionService.execute`. + +- **MCP tasks (SEP-1686)** — used when FastMCP has actually accepted *this specific* + call as a native background task (``ctx.is_background_task``). FastMCP runs the + ``task="optional"`` tool in a Docket worker and the client polls via ``tasks/get`` + / ``tasks/result``. Nothing extra is needed here; the service just runs the code + and returns the result. +- **Server-paced polling fallback** — used whenever ``mode="poll"`` is requested + (or ``mode="auto"`` resolves to ``"poll"``) and this call is *not* running as a + native background task. A client's general ``tasks`` capability declaration does + not guarantee it task-augmented this particular request, nor that it surfaces the + resulting notifications to the user — so capability alone must never suppress this + fallback. The service starts a background job, returns a ``job_id`` immediately, + and ``get_job_result`` waits (with adaptive pacing) for completion. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import functools +import itertools +import logging +import time +import uuid +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal + +from openstaad_mcp.connection import InstanceRegistry, StaadInstance, connect_and_run +from openstaad_mcp.file_io.helpers import ( + detect_input_output_collision, + get_allowed_dirs, + get_input_data, + write_output_file, +) +from openstaad_mcp.file_io.path_validator import FileIOError + +if TYPE_CHECKING: + from collections.abc import Callable + from concurrent.futures import Future + + from fastmcp.server.context import Context + + from openstaad_mcp.sandbox.executor import Executor + +logger = logging.getLogger(__name__) + +DEFAULT_TIMEOUT_SECONDS = 120.0 +ExecutionMode = Literal["native", "poll", "auto"] +# native — progress via MCP protocol notifications (client must render them) +# poll — AI polls get_job_result and writes progress to the user +# auto — detects progress token / background-task; falls back to "poll" + +# Poll pacing (see poll_hint): sleep grows linearly with elapsed time, clamped to +# [MIN, MAX] seconds; past ABANDON seconds get_job_result returns immediately. +_POLL_ABANDON_SECONDS = 1200.0 # 20 min — stop auto-pacing, let the user decide +_POLL_RAMP_SECONDS = 600.0 # reach the ceiling after 10 min +_POLL_MIN_INTERVAL = 10 +_POLL_MAX_INTERVAL = 55 + + +# ── Background job store (non-task-client fallback) ─────────────── + + +@dataclass +class _Job: + """In-flight or completed background execution.""" + + future: asyncio.Future[dict[str, Any]] + created: float + timeout: float = DEFAULT_TIMEOUT_SECONDS + progress_message: str = "" + task: asyncio.Task[None] | None = None + + +class JobStore: + """Minimal in-memory store for background code executions. + + Jobs are evicted after *ttl_seconds* (default 2 h). + """ + + def __init__(self, ttl_seconds: float = 7200.0) -> None: + self._jobs: dict[str, _Job] = {} + self._delivered: dict[str, float] = {} + self._ttl = ttl_seconds + + def create(self, future: asyncio.Future[dict[str, Any]], timeout: float = DEFAULT_TIMEOUT_SECONDS) -> str: + self._evict() + job_id = uuid.uuid4().hex[:12] + now = time.monotonic() + self._jobs[job_id] = _Job(future=future, created=now, timeout=timeout) + return job_id + + def get(self, job_id: str) -> _Job | None: + self._evict() + return self._jobs.get(job_id) + + def pop(self, job_id: str) -> _Job | None: + self._evict() + job = self._jobs.pop(job_id, None) + if job is not None: + self._delivered[job_id] = time.monotonic() + return job + + def was_delivered(self, job_id: str) -> bool: + return job_id in self._delivered + + def _evict(self) -> None: + now = time.monotonic() + expired = [k for k, v in self._jobs.items() if now - v.created > self._ttl] + for k in expired: + job = self._jobs.pop(k) + self._delivered[k] = now + job.future.cancel() + # Bound the delivered-ID history so it doesn't grow forever on a long-lived server. + stale = [k for k, delivered_at in self._delivered.items() if now - delivered_at > self._ttl] + for k in stale: + del self._delivered[k] + + +def poll_hint(job: _Job) -> int: + """Return how many seconds ``get_job_result`` should sleep before responding. + + Returns 0 for jobs running longer than ``_POLL_ABANDON_SECONDS`` (caller returns + immediately and tells the user to ask again manually). Otherwise the sleep grows + linearly with elapsed time, clamped to ``[_POLL_MIN_INTERVAL, _POLL_MAX_INTERVAL]`` + seconds to avoid burning tokens on overly frequent polls while still giving timely + updates. + """ + elapsed = time.monotonic() - job.created + if elapsed > _POLL_ABANDON_SECONDS: + return 0 # return immediately, let the user decide + span = _POLL_MAX_INTERVAL - _POLL_MIN_INTERVAL + interval = int(_POLL_MIN_INTERVAL + (elapsed / _POLL_RAMP_SECONDS) * span) + return max(_POLL_MIN_INTERVAL, min(_POLL_MAX_INTERVAL, interval)) + + +# ── Client capability detection ─────────────────────────────────── + + +def _has_progress_token(ctx: Context | None) -> bool: + """Return True when the client included a progressToken in this specific request. + + A progress token is a strong signal: the client is actively requesting + ``notifications/progress`` for this call and is expected to surface them. + Absence of a token means ``ctx.report_progress()`` is a silent no-op and + the user will see nothing from native MCP progress notifications. + """ + if ctx is None: + return False + try: + return getattr(ctx.meta, "progressToken", None) is not None + except Exception: + return False + + +# ── Result helpers ──────────────────────────────────────────────── + + +def _log_notify_failure(fut: Future[Any]) -> None: + """Consume a fire-and-forget notification future so its exception isn't swallowed silently.""" + with contextlib.suppress(BaseException): + exc = fut.exception() + if exc is not None: + logger.debug("progress notification failed: %s", exc) + + +def _error_result(message: str, duration: float = 0.0) -> dict[str, Any]: + return { + "success": False, + "result": None, + "stdout": "", + "stderr": "", + "error": message, + "duration_seconds": duration, + } + + +# ── Execution service ───────────────────────────────────────────── + + +class ExecutionService: + """Owns instance resolution, file I/O wiring and code-execution dispatch.""" + + def __init__( + self, + registry: InstanceRegistry, + executor: Executor, + args_allowed_dirs: list[Path], + jobs: JobStore | None = None, + ) -> None: + self._registry = registry + self._executor = executor + self._args_allowed_dirs = args_allowed_dirs + self._jobs = jobs if jobs is not None else JobStore() + + @property + def executor_busy(self) -> bool: + return self._executor.is_busy + + def resolve_target(self, instance: str | None) -> StaadInstance: + """Return the target StaadInstance or raise ValueError.""" + instances = self._registry.get_active_instances() + if not instances: + raise ValueError("No STAAD.Pro instances found") + if instance is None: + if len(instances) > 1: + aliases = [i.alias for i in instances] + raise ValueError(f"Multiple instances running — specify one: {aliases}") + return instances[0] + pid = self._registry.resolve(instance) + if pid is None: + alive = [i.alias for i in instances] + raise ValueError(f"{instance!r} is unknown. Available: {alive}") + matches = [i for i in instances if i.pid == pid] + if not matches: + alive = [i.alias for i in instances] + raise ValueError(f"{instance!r} is no longer running. Available: {alive}") + return matches[0] + + async def execute( + self, + *, + ctx: Context | None, + code: str, + instance: str | None = None, + input_data_path: str | None = None, + output_data_path: str | None = None, + overwrite: bool = False, + mode: ExecutionMode = "auto", + timeout: float | None = None, + ) -> dict[str, Any]: + """Resolve target + inputs, then run *code* synchronously or as a background job.""" + try: + target = self.resolve_target(instance) + except ValueError as e: + return _error_result(str(e)) + + # ── File I/O resolution (request-context bound; must run before any job) ── + allowed_dirs = await get_allowed_dirs(ctx, self._args_allowed_dirs, input_data_path, output_data_path) + try: + await detect_input_output_collision(input_data_path, output_data_path, allowed_dirs) + input_data, _ = await get_input_data(input_data_path, allowed_dirs) + except FileIOError as e: + return _error_result(f"{e.code}: {e.message}") + + effective_timeout = timeout if timeout is not None else DEFAULT_TIMEOUT_SECONDS + + def make_blocking(progress_fn: Callable[[str], None]) -> Callable[[], dict[str, Any]]: + return functools.partial( + self._run_blocking, + target=target, + code=code, + input_data=input_data, + output_data_path=output_data_path, + overwrite=overwrite, + allowed_dirs=allowed_dirs, + timeout=effective_timeout, + progress_fn=progress_fn, + ) + + loop = asyncio.get_running_loop() + + # Auto-detect: use "poll" unless there is strong evidence the client will + # surface progress natively. Two reliable signals: + # 1. ctx.is_background_task — FastMCP actually assigned a task_id to *this* + # request; the client is expected to poll tasks/get for status. + # 2. _has_progress_token — the client included a progressToken, explicitly + # requesting notifications/progress for this call. + # A client's general ``tasks`` capability alone is NOT reliable: it may not + # have task-augmented this specific call (tool is ``task="optional"``), or it + # may declare the capability but silently swallow the resulting notifications. + resolved = mode + if resolved == "auto": + is_bg = ctx is not None and getattr(ctx, "is_background_task", False) + resolved = "native" if is_bg or _has_progress_token(ctx) else "poll" + + use_fallback = resolved == "poll" and ctx is not None and not getattr(ctx, "is_background_task", False) + + if not use_fallback: + return await self._run_sync(loop, ctx, make_blocking, effective_timeout) + return self._start_job(loop, make_blocking, effective_timeout) + + def _run_blocking( + self, + *, + target: StaadInstance, + code: str, + input_data: Any, + output_data_path: str | None, + overwrite: bool, + allowed_dirs: list[Path], + timeout: float, + progress_fn: Callable[[str], None], + ) -> dict[str, Any]: + """Blocking work: connect on the COM thread, run the sandbox, write output.""" + + def _run(staad: Any) -> dict[str, Any]: + return self._executor.execute( + code, staad, input_data=input_data, progress_fn=progress_fn, lock_timeout=timeout + ).to_dict() + + result = connect_and_run(_run, target.file_path, timeout) + + if output_data_path is not None and result.get("success"): + try: + result["result"] = write_output_file( + output_data_path, result["result"], allowed_dirs, overwrite=overwrite + ) + except FileIOError as e: + return { + "success": False, + "result": None, + "stdout": result.get("stdout", ""), + "stderr": result.get("stderr", ""), + "error": f"{e.code}: {e.message}", + "duration_seconds": result.get("duration_seconds", 0.0), + } + + if target.warning: + result["warning"] = target.warning + return result + + async def _run_sync( + self, + loop: asyncio.AbstractEventLoop, + ctx: Context | None, + make_blocking: Callable[[Callable[[str], None]], Callable[[], dict[str, Any]]], + timeout: float, + ) -> dict[str, Any]: + start = time.monotonic() + blocking = make_blocking(self._sync_progress_fn(loop, ctx)) + try: + return await loop.run_in_executor(None, blocking) + except TimeoutError: + return _error_result( + f"Code execution timed out after {timeout:.0f}s. Retry with mode='poll' to run in the background.", + round(time.monotonic() - start, 1), + ) + except Exception as e: + return _error_result(str(e)) + + def _start_job( + self, + loop: asyncio.AbstractEventLoop, + make_blocking: Callable[[Callable[[str], None]], Callable[[], dict[str, Any]]], + timeout: float, + ) -> dict[str, Any]: + future: asyncio.Future[dict[str, Any]] = loop.create_future() + job_id = self._jobs.create(future, timeout=timeout) + job = self._jobs.get(job_id) + + def _progress_fn(message: str) -> None: + if job is not None: + job.progress_message = message + + blocking = make_blocking(_progress_fn) + + async def _bg_task() -> None: + try: + result = await loop.run_in_executor(None, blocking) + except TimeoutError: + result = _error_result( + f"Code execution timed out after {timeout:.0f}s. " + "The operation took longer than expected — consider splitting it into smaller chunks.", + timeout, + ) + except Exception as e: + result = _error_result(str(e)) + if not future.done(): + future.set_result(result) + + if job is not None: + job.task = asyncio.create_task(_bg_task()) + + return { + "job_id": job_id, + "next_action": ( + f"IMMEDIATELY call get_job_result('{job_id}') — do NOT write any text to the user first. " + "Write the 'message' field to the user after each get_job_result response, then call it again immediately. " + "Stop after 5 consecutive 'running' responses: tell the user the job is still in progress " + f"(job_id='{job_id}') and wait for them to ask for an update before polling again." + ), + } + + @staticmethod + def _sync_progress_fn(loop: asyncio.AbstractEventLoop, ctx: Context | None) -> Callable[[str], None]: + """Build the ``progress()`` callback injected into the sandbox for synchronous/task execution. + + Reports on two independent MCP channels so a client renders whichever it supports: + + - ``ctx.log()`` — a logging notification (``notifications/message``). + - ``ctx.report_progress()`` — the dedicated progress-token notification + (``notifications/progress``) in the foreground, or — critically, when this call is + running as a native background task (``ctx.is_background_task``) — an update to the + task's own progress, which becomes visible to the client via ``tasks/get`` and + ``notifications/tasks/status``. This is the one channel that reaches a client polling + task status directly, so it must not be limited to the logging notification alone. + + Both are fire-and-forget: neither is awaited, matching the sandbox's threaded execution + model where ``progress()`` is called synchronously from a worker thread. + """ + step = itertools.count(1) + + def _progress_fn(message: str) -> None: + if ctx is None: + return + log_fut = asyncio.run_coroutine_threadsafe(ctx.log(message, level="info"), loop) + progress_fut = asyncio.run_coroutine_threadsafe(ctx.report_progress(next(step), message=message), loop) + log_fut.add_done_callback(_log_notify_failure) + progress_fut.add_done_callback(_log_notify_failure) + + return _progress_fn + + async def get_job_result(self, job_id: str) -> dict[str, Any]: + """Wait (server-paced) for a background job and return its status or result.""" + job = self._jobs.get(job_id) + if job is None: + if self._jobs.was_delivered(job_id): + return { + "status": "delivered", + "message": f"Result for {job_id!r} was already delivered. Do not poll again.", + } + return { + "status": "unknown", + "message": f"Job not found or expired: {job_id!r}. Results are kept for 2 hours.", + } + + # Sleep to enforce pacing — wake early if the job completes. + if not job.future.done(): + wait = poll_hint(job) + if wait > 0: + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(asyncio.shield(job.future), timeout=wait) + + elapsed = round(time.monotonic() - job.created, 1) + + if not job.future.done(): + progress = job.progress_message + if poll_hint(job) == 0: + msg = ( + f"⏳ Still running ({elapsed:.0f}s). Tell the user the operation is still in progress " + f"(job_id={job_id!r}) and wait for them to ask for an update." + ) + else: + msg = f"⏳ Running ({elapsed:.0f}s): {progress}" if progress else f"⏳ Running ({elapsed:.0f}s)..." + return {"status": "running", "message": msg, "elapsed_seconds": elapsed} + + # Done — pop and return the full result. + self._jobs.pop(job_id) + result = job.future.result() + status = "completed" if result.get("success") else "failed" + result["status"] = status + result["elapsed_seconds"] = elapsed + result["message"] = ( + f"✅ Completed in {elapsed:.0f}s" if status == "completed" else f"❌ Failed after {elapsed:.0f}s" + ) + return result diff --git a/src/openstaad_mcp/main.py b/src/openstaad_mcp/main.py index ee1f18b..7db54ac 100644 --- a/src/openstaad_mcp/main.py +++ b/src/openstaad_mcp/main.py @@ -96,8 +96,30 @@ def setup_logging(log_level: str) -> None: handlers=[ logging.StreamHandler(sys.stderr), ], + force=True, ) logging.info(f"Logging initialized at {log_level} level") + _quiet_fastmcp_to_client_logger() + + +def _quiet_fastmcp_to_client_logger() -> None: + """Silence FastMCP's internal "Sending to client" trace logger. + + FastMCP already delivers ``ctx.log()`` messages to clients out-of-band via + ``session.send_log_message`` — the ``to_client`` logger call in + ``fastmcp.server.context`` is a redundant internal trace, not the actual + delivery mechanism. Left at its default level, it duplicates every + progress message onto FastMCP's own Rich console handler, which resolves + ``sys.stderr`` dynamically on each write (``rich.console.Console(stderr=True)``). + Because ``execute_code`` temporarily reassigns the process-global + ``sys.stderr`` while sandboxed user code runs on a worker thread, a + same-time log call from the event-loop thread (e.g. triggered by + ``progress()``) can have its output captured into that request's sandboxed + ``stderr`` buffer instead of the real console — leaking internal SDK + chatter into the tool result and inflating it for long-running loops. + Raising this logger's level to WARNING removes the trace entirely. + """ + logging.getLogger("fastmcp.server.context.to_client").setLevel(logging.WARNING) def main(argv: list[str] | None = None) -> None: @@ -109,6 +131,7 @@ def main(argv: list[str] | None = None) -> None: if args.transport == "stdio": # Run FastMCP server in the main thread, the COM thread will be started by the lifespan. mcp = create_mcp_server(allowed_dirs) + _quiet_fastmcp_to_client_logger() # re-apply: create_mcp_server may reconfigure fastmcp's logger try: mcp.run(transport="stdio", show_banner=False) except KeyboardInterrupt: @@ -129,6 +152,7 @@ def main(argv: list[str] | None = None) -> None: ) } mcp = create_mcp_server(allowed_dirs, fastmcp_kwargs=fastmcp_kwargs) + _quiet_fastmcp_to_client_logger() # re-apply: create_mcp_server may reconfigure fastmcp's logger try: mcp.run( transport="http", diff --git a/src/openstaad_mcp/sandbox/executor.py b/src/openstaad_mcp/sandbox/executor.py index e862cbc..25a701e 100644 --- a/src/openstaad_mcp/sandbox/executor.py +++ b/src/openstaad_mcp/sandbox/executor.py @@ -21,6 +21,7 @@ import sys import threading import time +from collections.abc import Callable from dataclasses import dataclass from typing import Any @@ -78,12 +79,23 @@ def __init__( # sys.stdout / sys.stderr reassignment. self._exec_lock = threading.Lock() + @property + def is_busy(self) -> bool: + """Return True if the executor lock is currently held (execution in progress).""" + acquired = self._exec_lock.acquire(blocking=False) + if acquired: + self._exec_lock.release() + return False + return True + def execute( self, code: str, staad_object: Any, *, input_data: Any = None, + progress_fn: Callable[[str], None] | None = None, + lock_timeout: float = 5.0, ) -> ExecutionResult: """Validate and execute *code* in the sandbox. @@ -96,6 +108,14 @@ def execute( input_data: Optional pre-parsed, deep-frozen data injected as ``input_data`` in the sandbox globals. ``None`` when no input file is provided. + progress_fn: + Optional callable injected as ``progress`` in the sandbox. + Called with a human-readable message string for real-time updates. + lock_timeout: + How long (seconds) to wait for the execution lock before returning + an "Executor busy" error. Callers should pass a value at least as + large as the expected execution duration so that a second call waits + for the first to finish rather than failing instantly. Returns ------- @@ -117,16 +137,18 @@ def execute( sandbox_globals.update(self.injected_modules) sandbox_globals["staad"] = COMProxy(staad_object) sandbox_globals["input_data"] = input_data + sandbox_globals["progress"] = progress_fn if progress_fn is not None else lambda _: None - # ── 4. Execute with stdout/stderr capture ─────────────────── + # ── 4. Execute with stdout/stderr capture ──────────────────── captured_out, captured_err = LimitedStringIO(), LimitedStringIO() exec_error: BaseException | None = None duration = 0.0 - if not self._exec_lock.acquire(timeout=5.0): + if not self._exec_lock.acquire(timeout=lock_timeout): return ExecutionResult( success=False, - error="Executor busy — a previous operation may have timed out. Restart the server.", + error="Executor busy — a previous operation is still running. " + "Check get_status() and retry when executor_busy is false.", ) try: old_stdout, old_stderr = sys.stdout, sys.stderr diff --git a/src/openstaad_mcp/server.py b/src/openstaad_mcp/server.py index 9a370ae..b214004 100644 --- a/src/openstaad_mcp/server.py +++ b/src/openstaad_mcp/server.py @@ -7,10 +7,12 @@ MCP server definition — tools, lifespan, and ASGI app factory. Exposes MCP tools: -- ``discover_api`` — lists available skills and usage guidance -- ``read_skills`` — returns requested skill content -- ``execute_code`` — runs validated Python against the COM bridge -- ``get_status`` — reports connection health +- ``discover_api`` — lists available skills and usage guidance +- ``read_skills`` — returns requested skill content +- ``list_instances`` — lists running STAAD.Pro instances +- ``get_status`` — reports connection health +- ``execute_code`` — runs validated Python against the COM bridge (native MCP task, poll fallback, or auto-detect) +- ``get_job_result`` — returns the status/result of a background execution job """ from __future__ import annotations @@ -23,16 +25,11 @@ from fastmcp import FastMCP from fastmcp.server.context import Context from fastmcp.server.lifespan import lifespan +from fastmcp.utilities.tasks import TaskConfig from mcp.types import ToolAnnotations -from openstaad_mcp.connection import InstanceRegistry, StaadInstance, connect_and_run -from openstaad_mcp.file_io.helpers import ( - detect_input_output_collision, - get_allowed_dirs, - get_input_data, - write_output_file, -) -from openstaad_mcp.file_io.path_validator import FileIOError +from openstaad_mcp.connection import InstanceRegistry, connect_and_run +from openstaad_mcp.execution import ExecutionMode, ExecutionService from openstaad_mcp.sandbox.executor import Executor from openstaad_mcp.skills import SkillsManager from openstaad_mcp.version import check_version_warning @@ -50,27 +47,9 @@ def _register_tools( skills_mgr: SkillsManager, args_allowed_dirs: list[Path], ) -> None: - """Register MCP tools on *mcp*, closing over the *InstanceRegistry*.""" - - def _resolve_target(instance: str | None) -> StaadInstance: - """Return the target StaadInstance or raise ValueError.""" - instances = registry.get_active_instances() - if not instances: - raise ValueError("No STAAD.Pro instances found") - if instance is None: - if len(instances) > 1: - aliases = [i.alias for i in instances] - raise ValueError(f"Multiple instances running — specify one: {aliases}") - return instances[0] - pid = registry.resolve(instance) - if pid is None: - alive = [i.alias for i in instances] - raise ValueError(f"{instance!r} is unknown. Available: {alive}") - matches = [i for i in instances if i.pid == pid] - if not matches: - alive = [i.alias for i in instances] - raise ValueError(f"{instance!r} is no longer running. Available: {alive}") - return matches[0] + """Register MCP tools on *mcp*, delegating execution to an :class:`ExecutionService`.""" + + service = ExecutionService(registry, exc, args_allowed_dirs) @mcp.tool( annotations=ToolAnnotations( @@ -152,12 +131,13 @@ def get_status(instance: str | None = None) -> dict[str, Any]: Pass ``instance`` (alias from ``list_instances``) to target a specific instance. Omit it when only one instance is running. - Returns connection state, STAAD version, and model path. + Returns connection state, STAAD version, model path, and whether the + executor is currently busy (``executor_busy``). """ try: - target = _resolve_target(instance) + target = service.resolve_target(instance) except ValueError as e: - return {"connected": False, "error": str(e)} + return {"connected": False, "executor_busy": service.executor_busy, "error": str(e)} def _read_status(staad: Any) -> dict[str, Any]: version = staad.GetApplicationVersion() @@ -175,6 +155,7 @@ def _read_status(staad: Any) -> dict[str, Any]: "model_path": model_path, "alias": target.alias, "analyzing": analyzing, + "executor_busy": service.executor_busy, } warning = check_version_warning(version) if warning: @@ -184,9 +165,9 @@ def _read_status(staad: Any) -> dict[str, Any]: try: return connect_and_run(_read_status, target.file_path, timeout=10.0) except TimeoutError: - return {"connected": False, "error": "Connection timed out"} + return {"connected": False, "executor_busy": service.executor_busy, "error": "Connection timed out"} except Exception as e: - return {"connected": False, "error": str(e)} + return {"connected": False, "executor_busy": service.executor_busy, "error": str(e)} @mcp.tool( annotations=ToolAnnotations( @@ -195,7 +176,8 @@ def _read_status(staad: Any) -> dict[str, Any]: destructiveHint=True, idempotentHint=False, # Different result for repeated calls openWorldHint=False, # Only internal data - ) + ), + task=TaskConfig(mode="optional"), ) async def execute_code( ctx: Context, @@ -204,11 +186,16 @@ async def execute_code( input_data_path: str | None = None, output_data_path: str | None = None, overwrite: bool = False, + mode: ExecutionMode = "auto", + timeout: float | None = None, ) -> dict[str, Any]: """Execute Python code in a sandbox against the OpenSTAAD API (don't forget to call discover_api and read_skills for API guidance). The sandbox provides pre-connected ``staad`` (the OpenSTAAD root object) and ``input_data`` (if input_data_path is provided) variables (plus ``json`` - and ``math`` modules). `import` statements, `dir()`, `getattr()`, ... are **BLOCKED**. + and ``math`` modules) and a ``progress(message)`` callback. `import` statements, `dir()`, `getattr()`, ... are **BLOCKED**. + + Call ``progress(f"Processing {i}/{total}")`` inside long loops or before a long single operation + (analysis, design) so the user sees real-time feedback. The last expression value or an explicit ``result = ...`` assignment is returned as the result. If ``output_data_path`` is provided, the sandbox will write the result to the specified file. @@ -217,6 +204,20 @@ async def execute_code( On Claude Desktop, users can configure allowed directories in the extension settings and Claude can use the filesystem ``copy_file_to_claude`` tool to move files to Claude's filesystem. + Long-running work: by default (``mode="auto"``) the server detects whether the client will + surface progress natively (via a progress token or native MCP background task) and falls back + to ``"poll"`` automatically when it cannot. Pass ``mode="poll"`` explicitly to always use the + AI-polling path: the server returns a ``job_id`` immediately and you poll ``get_job_result(job_id)`` + (writing its ``message`` to the user each time). Pass ``mode="native"`` to always trust the + client's MCP progress notifications. + + IMPORTANT — when a ``job_id`` is returned: your very next tool call MUST be + ``get_job_result``. Do NOT write any text to the user before the first poll — every + second of delay is progress the user cannot see. Write the ``message`` field to the + user *after* each ``get_job_result`` response, then call it again immediately. + Stop after 5 consecutive ``"running"`` responses: tell the user the job is still in + progress (include the ``job_id``) and wait for them to ask for an update. + Parameters ---------- code: str @@ -241,80 +242,47 @@ async def execute_code( } overwrite: bool, optional Allow overwriting an existing output file. + mode: {"auto", "native", "poll"}, optional + ``"auto"`` (default) detects whether the client will surface progress natively and + falls back to ``"poll"`` when it cannot. ``"native"`` trusts the client's MCP progress + notifications (``notifications/progress``). ``"poll"`` always returns a ``job_id`` + immediately so you can poll ``get_job_result`` and relay progress in your text responses. + timeout: float, optional + Max seconds to wait for the call to complete (default: 120). The COM operation + cannot be safely interrupted, so on timeout the call returns an error but the work + continues running in the background (the executor stays busy, reported via + ``get_status``'s ``executor_busy``, until it finishes on its own). """ - try: - target = _resolve_target(instance) - except ValueError as e: - return { - "success": False, - "result": None, - "stdout": "", - "stderr": "", - "error": str(e), - "duration_seconds": 0.0, - } - - # ── Resolve allowed dirs for path validation ── - allowed_dirs = await get_allowed_dirs(ctx, args_allowed_dirs, input_data_path, output_data_path) - - # ── Input file handling (server-side, outside sandbox) ─────── - try: - await detect_input_output_collision(input_data_path, output_data_path, allowed_dirs) - input_data, _ = await get_input_data(input_data_path, allowed_dirs) - except FileIOError as e: - return { - "success": False, - "result": None, - "stdout": "", - "stderr": "", - "error": f"{e.code}: {e.message}", - "duration_seconds": 0.0, - } - - # ── Execute code in sandbox ────────────────────────────────── - def _run(staad: Any) -> dict[str, Any]: - return exc.execute(code, staad, input_data=input_data).to_dict() - - try: - result = connect_and_run(_run, target.file_path) - except TimeoutError: - return { - "success": False, - "result": None, - "stdout": "", - "stderr": "", - "error": "Code execution timed out", - "duration_seconds": 0.0, - } - except Exception as e: - return { - "success": False, - "result": None, - "stdout": "", - "stderr": "", - "error": str(e), - "duration_seconds": 0.0, - } - - # ── Output file handling (server-side, outside sandbox) ────── - if output_data_path is not None and result.get("success"): - try: - result["result"] = write_output_file( - output_data_path, result["result"], allowed_dirs, overwrite=overwrite - ) - except FileIOError as e: - return { - "success": False, - "result": None, - "stdout": result.get("stdout", ""), - "stderr": result.get("stderr", ""), - "error": f"{e.code}: {e.message}", - "duration_seconds": result.get("duration_seconds", 0.0), - } + return await service.execute( + ctx=ctx, + code=code, + instance=instance, + input_data_path=input_data_path, + output_data_path=output_data_path, + overwrite=overwrite, + mode=mode, + timeout=timeout, + ) - if target.warning: - result["warning"] = target.warning - return result + @mcp.tool( + annotations=ToolAnnotations( + title="Get background job result — ALWAYS show message to user", + readOnlyHint=True, + destructiveHint=False, + idempotentHint=False, # Job state can change between calls + openWorldHint=False, + ) + ) + async def get_job_result(job_id: str) -> dict[str, Any]: + """Wait (server-paced) for a background ``execute_code`` job and return its status or result. + + The server sleeps internally before responding — call this again immediately after each + response regardless of status. Write the ``message`` field to the user after each call; + that is the only way they see progress. Stop when ``status`` is ``"completed"``, + ``"failed"`` (full result payload included; job is then removed from the store), or + ``"delivered"`` (the terminal result was already returned on an earlier poll). + """ + return await service.get_job_result(job_id) def create_mcp_server(allowed_dirs: list[Path], fastmcp_kwargs: dict | None = None) -> FastMCP: @@ -336,6 +304,8 @@ async def mcp_lifespan(server: Any) -> AsyncIterator[None]: "instructions. Use `list_instances` to see running STAAD instances, " "`execute_code` to run code against a live STAAD.Pro model, and " "`get_status` to check connection. " + "When `execute_code` returns a `job_id`, follow the `next_action` field exactly: " + "call `get_job_result` immediately — no text output before the first poll. " "When a `warning` field appears in any tool response, report it to the user." ), lifespan=mcp_lifespan, diff --git a/src/openstaad_mcp/staad_skills/staad-core/SKILL.md b/src/openstaad_mcp/staad_skills/staad-core/SKILL.md index a377abb..8ba4cf7 100644 --- a/src/openstaad_mcp/staad_skills/staad-core/SKILL.md +++ b/src/openstaad_mcp/staad_skills/staad-core/SKILL.md @@ -1,6 +1,6 @@ --- name: staad-core -description: "ALWAYS load first for any STAAD.Pro automation. Covers: Python sandbox (staad pre-injected — import blocked), sub-module access (Geometry, Property, Support, Load, Command, Output, Design), units and axis check via execute_code, unit conversion (English=inches/KIP, Metric=meters/kN), GetBaseUnit, IsZUp, SetSilentMode required before UpdateStructure/AnalyzeModel/AnalyzeEx/SaveModel/file operations, UpdateStructure semantics, application control (ShowApplication, GetApplicationVersion, Quit). Do not auto-save." +description: "ALWAYS load first for any STAAD.Pro automation. Covers: Python sandbox (staad/input_data/progress pre-injected — import blocked), sub-module access (Geometry, Property, Support, Load, Command, Output, Design), execution modes (auto default / native = MCP protocol progress / poll = AI polls get_job_result + writes message to user), progress reporting via progress(), timeout guidelines (default 120s, override with timeout=), large model query patterns, units and axis check via execute_code, unit conversion (English=inches/KIP, Metric=meters/kN), GetBaseUnit, IsZUp, SetSilentMode required before UpdateStructure/AnalyzeModel/AnalyzeEx/SaveModel/file operations, UpdateStructure semantics, application control (ShowApplication, GetApplicationVersion, Quit). Do not auto-save." --- # STAAD.Pro Core — Sandbox & Model Setup @@ -9,20 +9,174 @@ description: "ALWAYS load first for any STAAD.Pro automation. Covers: Python san ### Sandbox -- Pre-injected names (do NOT import): `staad`, `input_data`, `json`, `math` +- Pre-injected names (do NOT import): `staad`, `input_data`, `json`, `math`, `progress` - `import` statements, `dir()`, `getattr()`, ... are **BLOCKED** — use skills for discovery, only use pre-injected names in code - `staad` is already connected and ready — do NOT call any initialization function - `input_data` is injected if `input_data_path` is provided in `execute_code` params — use it to feed large datasets into the sandbox without hardcoding - Sub-modules: `geo = staad.Geometry`, `prop = staad.Property`, `sup = staad.Support`, `load = staad.Load`, `cmd = staad.Command`, `out = staad.Output`, `design = staad.Design` +- `progress(message)` — send a real-time status update (e.g. `progress(f"Node {i}/{total}")`) - If `output_data_path` is provided, write the `result` variable to that file path instead of returning it in the context (use for large/tabular data). The `execute_code` return value will contain a summary of the `result` content instead (e.g. number of rows, columns and a sample of rows). - Both `input_data_path` and `output_data_path` must be on the user LOCAL filesystem and inside MCP roots or configured `allowed_dirs`. On Claude Desktop, users can configure allowed directories in the extension settings and Claude can use the filesystem `copy_file_to_claude` tool to move files to Claude's filesystem. +### Progress + +Call `progress(message)` every 10–50 iterations of a loop (with a counter), and once before any long single op — otherwise the client looks frozen. Don't call it every iteration (flooding). + +```python +node_ids = staad.Geometry.GetNodeList() +total = len(node_ids) +for i, nid in enumerate(node_ids): + if i % 20 == 0: + progress(f"Node {i}/{total}") + ... # work +# or, before a single long op: progress("Running analysis...") +``` + +### Before executing + +If a prior `execute_code` may still be running, call `get_status` first — it returns `executor_busy`. If `true`, wait and retry; a new execution is rejected while busy. + +### Execution modes + +`execute_code(mode=...)` chooses how progress reaches the user, via two channels: + +- **native** — real-time MCP notifications rendered *by the client* (needs client support). +- **poll** — returns a `job_id`; *you* relay progress by polling `get_job_result` (works on every client). + +Pick by expected duration (only you know it — the server can't): + +| Work | Mode | Why | +| --- | --- | --- | +| Quick query / small mutation — `GetNodeCount`, `AddNode`, assign to a known list | `native` | nothing to show; one call, no poll overhead | +| Medium/heavy loop — 100s–10 000s of elements | `auto` | server uses native progress if the client supports it, else falls back to poll — best channel per client | +| Analysis / nonlinear / P-Delta / buckling / design | `poll` | guarantees background execution so a long run can't hit a transport timeout | + +`auto` (the param default) resolves to native or poll by inspecting the client; use `native`/`poll` explicitly for the ends of the range above. + +**Poll loop — when a call returns a `job_id`, follow exactly:** + +1. Response has `job_id` + `next_action`. **Write nothing yet** — call `get_job_result(job_id)` immediately. +2. Write the returned `message` to the user. +3. `status == "running"` → call `get_job_result` again immediately (server paces at 10–55 s internally). +4. `status == "completed"` / `"failed"` → present the result; the job is then gone (a later poll returns `"delivered"` — don't re-poll). +5. After **5** consecutive `"running"` (or a job >20 min), stop: tell the user it's still running (include `job_id`) and wait for them to ask before polling again. + +### Timeout + +Default is **120 s** per call (no auto-detection). Pass `timeout=` for longer work; pair long runs with `mode="poll"`: + +``` +execute_code(code="...", timeout=3600, mode="poll") # any analysis +``` + +Estimate loops at ~3 ms per COM call — `timeout ≈ elements × calls_per_element × 0.003` — and round up generously (too short kills the op and wastes progress): + +- 1 call/element (single property): `elements × 0.003` +- 2 calls/element (coords + property): `elements × 0.006` +- per load case: `elements × load_cases × 0.003` +- any analysis: `timeout=3600` + +e.g. 100 000 plates × 2 calls → `timeout=600` (10 min). + +### Result Structure for Bulk Queries + +Build a **dict keyed by element ID** rather than parallel lists or nested loops. This avoids O(n²) lookups, produces self-describing JSON, and keeps results compact. + +**Preferred pattern:** + +```python +node_ids = staad.Geometry.GetNodeList() +nodes = {} +total = len(node_ids) +for i, nid in enumerate(node_ids): + if i % 20 == 0: + progress(f"Reading node {i}/{total}") + x, y, z = staad.Geometry.GetNodeCoordinates(nid) + nodes[nid] = {"x": x, "y": y, "z": z} +result = nodes +``` + +Rules: + +- Use `result = {...}` assignment so the sandbox returns it as the tool result. +- Top-level key: entity type in plural (`"nodes"`, `"beams"`, `"loads"`). +- Each value: a flat dict of properties; avoid deeply nested structures. +- For cross-entity data nest one level: `{beam_id: {"loads": [...]}}`. +- Never use parallel lists (`ids = [...]`, `xs = [...]`) — they break if order differs. + +### Large Model Queries + +For models with >500 nodes/beams, query counts first and adapt strategy before bulk-fetching. + +**Step 1 — count before fetch:** + +```python +node_count = staad.Geometry.GetNodeCount() +beam_count = staad.Geometry.GetMemberCount() +progress(f"Model: {node_count} nodes, {beam_count} beams") +``` + +**Step 2 — choose strategy:** + +| Count | Strategy | +| --- | --- | +| < 500 | Fetch all, return full dict | +| 500–2 000 | Fetch all with `progress()`, summarise if result near 200 KB | +| > 2 000 | Targeted ranges or summary statistics | + +**Summary statistics (avoids large result):** + +```python +node_ids = staad.Geometry.GetNodeList() +xs, ys, zs = [], [], [] +for nid in node_ids: + x, y, z = staad.Geometry.GetNodeCoordinates(nid) + xs.append(x); ys.append(y); zs.append(z) +result = { + "node_count": len(node_ids), + "x_range": [min(xs), max(xs)], + "y_range": [min(ys), max(ys)], + "z_range": [min(zs), max(zs)], +} +``` + +**Sampling for initial exploration:** + +```python +step = max(1, node_count // 50) # ~50 samples +sample_ids = node_ids[::step] +nodes = {} +for nid in sample_ids: + x, y, z = staad.Geometry.GetNodeCoordinates(nid) + nodes[nid] = {"x": x, "y": y, "z": z} +result = {"sampled": True, "sample_size": len(sample_ids), "nodes": nodes} +``` + +Rules: + +- Never fetch all coordinates for >2 000 nodes in a single call — the result may be too large. +- Prefer summary statistics for initial exploration; fetch raw data only when the user needs specific elements. + +### Typical Workflows + +Load skills in this order for common tasks: + +| Task | Skills to load | +| --- | --- | +| Query an existing model | `staad-core` → `staad-results` | +| Build a model from scratch | `staad-core` → `staad-geometry` → `staad-properties` → `staad-supports` → `staad-loading` → `staad-analysis` → `staad-results` | +| Run steel design | `staad-core` → `staad-steel-design` | +| Add loads to existing geometry | `staad-core` → `staad-loading` → `staad-analysis` → `staad-results` | +| Export a screenshot | `staad-core` → `staad-view` | +| Robust scripting / error handling | Add `staad-errors` to any of the above | + ### Discovery Before writing any script: 1. Call `discover_api` → lists available skills and usage guidance 2. Call `read_skills` with skill names → detailed instructions for that domain +3. If you already know a function name but not which skill covers it, read `./assets/FUNCTION_SKILL_MAP.md` for a quick function → skill lookup Never guess or invent function names — only use names from the skill documentation. @@ -32,6 +186,17 @@ Never guess or invent function names — only use names from the skill documenta - Call `get_status(instance)` to verify a specific instance is reachable - Pass `instance` (alias like `staadPro1`) to `execute_code` when multiple instances are running +### Tool Reference + +| Tool | Purpose | +| --- | --- | +| `discover_api` | List available skills | +| `read_skills` | Load skill instructions | +| `list_instances` | List running STAAD.Pro instances | +| `get_status` | Check connection to an instance (includes `executor_busy`) | +| `execute_code` | Run code — default `mode="auto"`/120 s; pass `mode="poll"` and `timeout=` for long work | +| `get_job_result` | Long-poll a running poll job or collect its result when done | + ### Version Compatibility The MCP is built against a single bundled **openstaadpy** wrapper, identical for @@ -151,7 +316,7 @@ staad.CloseSTAADFile() ## Gotchas -- `import`, `dir()`, `getattr()`, ... are blocked — only `staad`, `input_data`, `json`, `math` are available +- `import`, `dir()`, `getattr()`, ... are blocked — only `staad`, `input_data`, `json`, `math`, `progress` are available - If `input_data_path` is provided, `input_data` is injected as an immutable variable — use it to feed large datasets into the sandbox without hardcoding - If `output_data_path` is provided, write the `result` variable to that file path instead of returning it in the context (use for large/tabular data). The `execute_code` return value will contain a summary of the `result` content instead (e.g. number of rows, columns and a sample of rows). - Both `input_data_path` and `output_data_path` must be on the user LOCAL filesystem and inside MCP roots or configured `allowed_dirs`. On Claude Desktop, users can configure allowed directories in the extension settings and Claude can use the filesystem `copy_file_to_claude` tool to move files to Claude's filesystem. diff --git a/src/openstaad_mcp/staad_skills/staad-geometry/SKILL.md b/src/openstaad_mcp/staad_skills/staad-geometry/SKILL.md index e51dfc5..9714a72 100644 --- a/src/openstaad_mcp/staad_skills/staad-geometry/SKILL.md +++ b/src/openstaad_mcp/staad_skills/staad-geometry/SKILL.md @@ -133,6 +133,7 @@ pm_count = geo.GetPhysicalMemberCount() - [add-beam.py](./scripts/add-beam.py) — add a single beam between two new nodes - [add-plate.py](./scripts/add-plate.py) — create quad and triangular plates (shows the `0` triangle convention) - [select-members.py](./scripts/select-members.py) — select single and multiple beams +- [plate-max-area.py](./scripts/plate-max-area.py) — find the largest plate in a 100k+ plate model (shows `progress()` + coordinate pre-fetch for heavy loops) ## Gotchas - **Triangle plates use `0` as a sentinel:** all plate functions always take exactly 4 node arguments; `0` tells STAAD the slot is empty (i.e., this is a 3-node element). See [add-plate.py](./scripts/add-plate.py) for a full example. diff --git a/src/openstaad_mcp/staad_skills/staad-geometry/scripts/plate-max-area.py b/src/openstaad_mcp/staad_skills/staad-geometry/scripts/plate-max-area.py new file mode 100644 index 0000000..8ac2532 --- /dev/null +++ b/src/openstaad_mcp/staad_skills/staad-geometry/scripts/plate-max-area.py @@ -0,0 +1,66 @@ +# plate-max-area.py +# Finds the plate with the maximum area in a large model (100k+ plates). +# Optimized: fetches all node coordinates once into a dict, then iterates plates. +# Uses progress() for real-time feedback on large models. +# GetPlateIncidence(pid) returns (n1, n2, n3, n4) — n4 is 0 for triangles. + +geo = staad.Geometry + +# ── 1. Area calculation helpers (defined before use) ───────────── + +def tri_area(p0, p1, p2): + """Area of a 3D triangle via cross product of edge vectors.""" + v1x, v1y, v1z = p1[0] - p0[0], p1[1] - p0[1], p1[2] - p0[2] + v2x, v2y, v2z = p2[0] - p0[0], p2[1] - p0[1], p2[2] - p0[2] + cx = v1y*v2z - v1z*v2y + cy = v1z*v2x - v1x*v2z + cz = v1x*v2y - v1y*v2x + return 0.5 * math.sqrt(cx*cx + cy*cy + cz*cz) + +def plate_area(pts): + """Area of a plate (3 or 4 coordinate tuples).""" + if len(pts) == 3: + return tri_area(pts[0], pts[1], pts[2]) + if len(pts) == 4: + return (tri_area(pts[0], pts[1], pts[2]) + + tri_area(pts[0], pts[2], pts[3])) + return 0.0 + +# ── 2. Pre-fetch all node coordinates into a dict (single pass) ── +node_ids = geo.GetNodeList() +total_nodes = len(node_ids) +coords = {} +for i, nid in enumerate(node_ids): + if i % 500 == 0: + progress(f"Reading nodes {i}/{total_nodes}") + coords[nid] = geo.GetNodeCoordinates(nid) + +progress(f"Loaded {total_nodes} nodes. Scanning plates...") + +# ── 3. Iterate plates and find maximum area ────────────────────── +plates = geo.GetPlateList() +total_plates = len(plates) + +max_area = 0.0 +max_plate = None + +for i, pid in enumerate(plates): + if i % 500 == 0: + progress(f"Plate {i}/{total_plates}") + + incidence = geo.GetPlateIncidence(pid) + pts = [coords[nid] for nid in incidence if nid != 0] + + area = plate_area(pts) + if area > max_area: + max_area = area + max_plate = pid + +# Include incidence nodes of the winner for context +winner_nodes = [nid for nid in geo.GetPlateIncidence(max_plate) if nid != 0] +result = { + "plate": max_plate, + "area": round(max_area, 4), + "nodes": winner_nodes, + "total_plates": total_plates, +} diff --git a/tests/test_connection.py b/tests/test_connection.py index f20b205..d016c00 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -139,7 +139,7 @@ def test_auto_select_single_instance(self): with ( _mock_get_active_instances(single), - patch("openstaad_mcp.server.connect_and_run", return_value=expected) as mock_run, + patch("openstaad_mcp.execution.connect_and_run", return_value=expected) as mock_run, ): mcp = create_mcp_server(allowed_dirs=[]) asyncio.run(mcp.call_tool("execute_code", {"code": "result = 42"})) diff --git a/tests/test_execution.py b/tests/test_execution.py new file mode 100644 index 0000000..096c901 --- /dev/null +++ b/tests/test_execution.py @@ -0,0 +1,343 @@ +""" +--------------------------------------------------------------------------------------------- +Copyright (c) Bentley Systems, Incorporated. All rights reserved. +See LICENSE.md in the project root for license terms and full copyright notice. +--------------------------------------------------------------------------------------------- + +Tests for the code-execution orchestration layer (openstaad_mcp.execution). +""" + +from __future__ import annotations + +import asyncio +import time +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +from openstaad_mcp.connection import StaadInstance +from openstaad_mcp.execution import ( + DEFAULT_TIMEOUT_SECONDS, + ExecutionService, + JobStore, + _error_result, + _Job, + poll_hint, +) + + +class _FakeRegistry: + def __init__(self, instances: list[StaadInstance], resolve_map: dict[str, int] | None = None) -> None: + self._instances = instances + self._resolve_map = resolve_map or {} + + def get_active_instances(self) -> list[StaadInstance]: + return self._instances + + def resolve(self, instance: str) -> int | None: + return self._resolve_map.get(instance) + + +def _instance(alias: str = "staadPro1", pid: int = 1, path: str = "C:\\A.std") -> StaadInstance: + return StaadInstance(alias=alias, pid=pid, file_path=path, version="22.12") + + +def _fake_ctx(*, is_background_task: bool = False, tasks: Any = None) -> SimpleNamespace: + return SimpleNamespace( + is_background_task=is_background_task, + session=SimpleNamespace(client_params=SimpleNamespace(capabilities=SimpleNamespace(tasks=tasks))), + ) + + +def _service(instances: list[StaadInstance], resolve_map: dict[str, int] | None = None) -> ExecutionService: + executor = MagicMock() + executor.is_busy = False + return ExecutionService(_FakeRegistry(instances, resolve_map), executor, []) + + +# --------------------------------------------------------------------------- +# poll_hint +# --------------------------------------------------------------------------- + + +def _job_with_age(age_seconds: float) -> _Job: + loop = asyncio.new_event_loop() + try: + fut: asyncio.Future[dict[str, Any]] = loop.create_future() + finally: + loop.close() + return _Job(future=fut, created=time.monotonic() - age_seconds) + + +class TestPollHint: + def test_fresh_job_returns_floor(self): + assert poll_hint(_job_with_age(0)) == 10 + + def test_old_job_clamped_to_ceiling(self): + assert poll_hint(_job_with_age(700)) == 55 + + def test_very_old_job_returns_zero(self): + assert poll_hint(_job_with_age(1300)) == 0 + + +# --------------------------------------------------------------------------- +# JobStore +# --------------------------------------------------------------------------- + + +class TestJobStore: + def test_create_get_pop_delivered(self): + loop = asyncio.new_event_loop() + try: + store = JobStore() + fut: asyncio.Future[dict[str, Any]] = loop.create_future() + job_id = store.create(fut) + assert len(job_id) == 12 + assert store.get(job_id) is not None + popped = store.pop(job_id) + assert popped is not None + assert store.was_delivered(job_id) is True + assert store.get(job_id) is None + finally: + loop.close() + + def test_pop_missing_returns_none(self): + store = JobStore() + assert store.pop("does-not-exist") is None + + def test_eviction_after_ttl(self): + loop = asyncio.new_event_loop() + try: + store = JobStore(ttl_seconds=1.0) + fut: asyncio.Future[dict[str, Any]] = loop.create_future() + job_id = store.create(fut) + store._jobs[job_id].created -= 5.0 + assert store.get(job_id) is None + assert store.was_delivered(job_id) is True + assert fut.cancelled() is True + finally: + loop.close() + + +# --------------------------------------------------------------------------- +# _error_result +# --------------------------------------------------------------------------- + + +def test_error_result_shape(): + result = _error_result("nope", duration=1.5) + assert result == { + "success": False, + "result": None, + "stdout": "", + "stderr": "", + "error": "nope", + "duration_seconds": 1.5, + } + + +# --------------------------------------------------------------------------- +# ExecutionService.resolve_target +# --------------------------------------------------------------------------- + + +class TestResolveTarget: + def test_no_instances(self): + with pytest.raises(ValueError, match=r"No STAAD\.Pro instances found"): + _service([]).resolve_target(None) + + def test_auto_select_single(self): + inst = _instance() + assert _service([inst]).resolve_target(None) is inst + + def test_multiple_requires_explicit(self): + two = [_instance("staadPro1", 1), _instance("staadPro2", 2)] + with pytest.raises(ValueError, match="Multiple instances running"): + _service(two).resolve_target(None) + + def test_unknown_alias(self): + with pytest.raises(ValueError, match="is unknown"): + _service([_instance()]).resolve_target("nope") + + def test_alias_no_longer_running(self): + svc = _service([_instance(pid=1)], resolve_map={"ghost": 999}) + with pytest.raises(ValueError, match="no longer running"): + svc.resolve_target("ghost") + + def test_alias_resolves_to_match(self): + inst = _instance(pid=42) + svc = _service([inst], resolve_map={"staadPro1": 42}) + assert svc.resolve_target("staadPro1") is inst + + +# --------------------------------------------------------------------------- +# ExecutionService.execute — sync +# --------------------------------------------------------------------------- + + +class TestExecuteSync: + async def test_success(self): + svc = _service([_instance()]) + expected = {"success": True, "result": 42, "stdout": "", "stderr": "", "error": None, "duration_seconds": 0.1} + with patch("openstaad_mcp.execution.connect_and_run", return_value=expected) as mock_run: + result = await svc.execute(ctx=None, code="result = 42") + assert result["result"] == 42 + assert mock_run.call_args[0][1] == "C:\\A.std" + assert mock_run.call_args[0][2] == DEFAULT_TIMEOUT_SECONDS + + async def test_resolve_error(self): + result = await _service([]).execute(ctx=None, code="result = 1") + assert result["success"] is False + assert "No STAAD.Pro instances found" in result["error"] + + async def test_timeout(self): + svc = _service([_instance()]) + with patch("openstaad_mcp.execution.connect_and_run", side_effect=TimeoutError): + result = await svc.execute(ctx=None, code="result = 1", timeout=0.5) + assert result["success"] is False + assert "timed out" in result["error"] + + async def test_generic_error(self): + svc = _service([_instance()]) + with patch("openstaad_mcp.execution.connect_and_run", side_effect=RuntimeError("kaboom")): + result = await svc.execute(ctx=None, code="result = 1") + assert result["success"] is False + assert "kaboom" in result["error"] + + async def test_poll_mode_without_ctx_runs_native(self): + """Without a ctx (no MCP session), poll mode cannot start a background job + and falls through to native synchronous execution.""" + svc = _service([_instance()]) + expected = {"success": True, "result": 1, "stdout": "", "stderr": "", "error": None, "duration_seconds": 0.0} + with patch("openstaad_mcp.execution.connect_and_run", return_value=expected): + result = await svc.execute(ctx=None, code="result = 1", mode="poll") + assert "job_id" not in result + assert result["result"] == 1 + + async def test_poll_mode_not_background_task_uses_job(self): + """poll mode with a real ctx that is not a background task returns a job_id + so the AI can relay progress via get_job_result.""" + svc = _service([_instance()]) + ctx = _fake_ctx(is_background_task=False, tasks=object()) + expected = {"success": True, "result": 1, "stdout": "", "stderr": "", "error": None, "duration_seconds": 0.0} + with patch("openstaad_mcp.execution.connect_and_run", return_value=expected): + result = await svc.execute(ctx=ctx, code="result = 1", mode="poll") + assert "job_id" in result + + async def test_poll_mode_background_task_runs_native(self): + """When this call is genuinely running as a native MCP background task + (a task_id was actually assigned), the outer tasks protocol already + handles async delivery — starting our own job would just leak in the + job store since nothing would ever poll it.""" + svc = _service([_instance()]) + ctx = _fake_ctx(is_background_task=True, tasks=object()) + expected = {"success": True, "result": 1, "stdout": "", "stderr": "", "error": None, "duration_seconds": 0.0} + with patch("openstaad_mcp.execution.connect_and_run", return_value=expected): + result = await svc.execute(ctx=ctx, code="result = 1", mode="poll") + assert "job_id" not in result + + async def test_output_written(self): + svc = _service([_instance()]) + raw = {"success": True, "result": [[1, 2]], "stdout": "", "stderr": "", "error": None, "duration_seconds": 0.0} + + async def _list_roots() -> list[Any]: + return [] + + ctx = SimpleNamespace(list_roots=_list_roots) + with ( + patch("openstaad_mcp.execution.connect_and_run", return_value=raw), + patch("openstaad_mcp.execution.get_allowed_dirs", return_value=[]), + patch("openstaad_mcp.execution.write_output_file", return_value={"path": "out.csv"}) as mock_write, + ): + result = await svc.execute(ctx=ctx, code="result = 1", output_data_path="out.csv", mode="native") + assert mock_write.called + assert result["result"] == {"path": "out.csv"} + + +# --------------------------------------------------------------------------- +# ExecutionService.execute — poll fallback + get_job_result +# --------------------------------------------------------------------------- + + +class TestExecutePollFallback: + async def test_job_lifecycle(self): + svc = _service([_instance()]) + ctx = _fake_ctx(is_background_task=False, tasks=None) + expected = {"success": True, "result": 7, "stdout": "", "stderr": "", "error": None, "duration_seconds": 0.0} + with patch("openstaad_mcp.execution.connect_and_run", return_value=expected): + started = await svc.execute(ctx=ctx, code="result = 7", mode="poll") + job_id = started["job_id"] + assert "job_id" in started + + for _ in range(50): + result = await svc.get_job_result(job_id) + if result.get("status") != "running": + break + await asyncio.sleep(0.01) + + assert result["status"] == "completed" + assert result["result"] == 7 + + # Second lookup: already delivered. + again = await svc.get_job_result(job_id) + assert again["status"] == "delivered" + assert "already delivered" in again["message"] + + async def test_unknown_job(self): + result = await _service([_instance()]).get_job_result("deadbeefdead") + assert result["status"] == "unknown" + + async def test_failed_job(self): + svc = _service([_instance()]) + ctx = _fake_ctx(is_background_task=False, tasks=None) + failed = { + "success": False, + "result": None, + "stdout": "", + "stderr": "", + "error": "boom", + "duration_seconds": 0.0, + } + with patch("openstaad_mcp.execution.connect_and_run", return_value=failed): + started = await svc.execute(ctx=ctx, code="raise", mode="poll") + for _ in range(50): + result = await svc.get_job_result(started["job_id"]) + if result.get("status") != "running": + break + await asyncio.sleep(0.01) + assert result["status"] == "failed" + + +# --------------------------------------------------------------------------- +# ExecutionService._sync_progress_fn +# --------------------------------------------------------------------------- + + +class TestSyncProgressFn: + async def test_reports_on_both_log_and_progress_channels(self): + log_calls: list[tuple[str, str]] = [] + progress_calls: list[tuple[float, str | None]] = [] + + async def _log(message: str, level: str) -> None: + log_calls.append((message, level)) + + async def _report_progress(progress: float, message: str | None = None) -> None: + progress_calls.append((progress, message)) + + ctx = SimpleNamespace(log=_log, report_progress=_report_progress) + loop = asyncio.get_running_loop() + progress_fn = ExecutionService._sync_progress_fn(loop, ctx) + + progress_fn("step 1") + progress_fn("step 2") + await asyncio.sleep(0.01) + + assert log_calls == [("step 1", "info"), ("step 2", "info")] + assert progress_calls == [(1, "step 1"), (2, "step 2")] + + async def test_noop_without_ctx(self): + loop = asyncio.get_running_loop() + progress_fn = ExecutionService._sync_progress_fn(loop, None) + progress_fn("ignored") # must not raise