diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..68ff88b --- /dev/null +++ b/.env.example @@ -0,0 +1,27 @@ +# openstudio-mcp environment variable template. +# Copy to .env and fill in your values — .env is gitignored and must never be committed. + +# --------------------------------------------------------------------------- +# Telemetry (optional — requires pip install 'openstudio-mcp[telemetry]') +# Leave TRACELOOP_BASE_URL unset to disable tracing entirely (zero overhead). +# --------------------------------------------------------------------------- + +# OTLP HTTP endpoint. Examples: +# Local Jaeger: http://localhost:4318 +# Traceloop cloud: https://api.traceloop.com +TRACELOOP_BASE_URL= + +# API key — required only for Traceloop cloud, not for generic OTLP backends. +TRACELOOP_API_KEY= + +# Service name shown on every span. +OTEL_SERVICE_NAME=openstudio-mcp + +# Set to "false" to use synchronous span export (useful in development). +OTEL_EXPORT_BATCH=true + +# IMPORTANT PRIVACY SETTING: when "true" (default), tool arguments and outputs +# — including file paths, model parameters, and simulation results — are +# exported to the OTLP backend. Set to "false" unless you have reviewed the +# data being exported and your backend is self-hosted or trusted. +TRACELOOP_TRACE_CONTENT=false diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ba4cf19..6185f5c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -26,6 +26,15 @@ jobs: mkdir -p runs docker run --rm -v "$PWD:/repo" -v "$PWD/runs:/runs" openstudio-mcp:dev bash -lc 'cd /repo && pytest -vv -m "not integration"' + - name: Smoke-install telemetry extra + # Validates that traceloop-sdk and its deps install cleanly from the + # pinned constraint in [telemetry]. Catches packaging drift that would + # make openstudio-mcp[telemetry] uninstallable without needing a full + # Docker rebuild. + run: | + docker run --rm -v "$PWD:/repo" openstudio-mcp:dev bash -lc \ + 'pip install --quiet -e "/repo[telemetry]" && python -c "from traceloop.sdk import Traceloop; print(\"traceloop-sdk OK\")"' + - name: Save Docker image run: docker save openstudio-mcp:dev | gzip > /tmp/image.tar.gz @@ -81,8 +90,8 @@ jobs: EXTRA_ENV="" ;; 5) - # HVAC supply sim smoke tests + hvac_validation + bar_building + concurrent regression - FILES="tests/test_hvac_supply_sim.py tests/test_hvac_validation.py tests/test_bar_building.py tests/test_concurrent_tools.py tests/test_stdout_logger_silence.py" + # HVAC supply sim smoke tests + hvac_validation + bar_building + concurrent regression + telemetry + FILES="tests/test_hvac_supply_sim.py tests/test_hvac_validation.py tests/test_bar_building.py tests/test_concurrent_tools.py tests/test_stdout_logger_silence.py tests/test_telemetry.py" EXTRA_ENV="" ;; esac diff --git a/.gitignore b/.gitignore index 8396c12..7c70477 100644 --- a/.gitignore +++ b/.gitignore @@ -33,5 +33,11 @@ Thumbs.db # Code review artifacts docs/review/ +# Environment / secrets +.env +.env.* +!.env.example + # Codex CLI .codex/ +.mcp.json diff --git a/CHANGELOG.md b/CHANGELOG.md index 301c256..a8dfb71 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,18 @@ # Changelog +## [Unreleased] + +### Added +- **Optional OpenLLMetry tracing**: `pip install 'openstudio-mcp[telemetry]'` + `TRACELOOP_BASE_URL` env var enables distributed tracing via traceloop-sdk. Zero overhead when unset. Key operations (`run_simulation`, `apply_measure`, `create_measure`, `create_*_building`, `run_qaqc_checks`) emit named spans; every FastMCP tool call is auto-instrumented via `McpInstrumentor`. +- **Docker tracing stack**: `docker/docker-compose.tracing.yml` + `docker/otel-collector-config.yaml` for local Jaeger/OTEL collector. +- **`test_telemetry.py`**: 20 unit tests for telemetry module (no Docker required) — includes startup-wiring and decorator-coverage regression tests. +- **`.env.example`**: template for telemetry environment variables with privacy guidance. + +### Fixed +- **README tracing Docker example**: corrected image tag from `openstudio-mcp:dev` to `openstudio-mcp:tracing` (the dev image does not include traceloop-sdk); added build command and explanatory note. +- **`TRACELOOP_TRACE_CONTENT` docs**: expanded to warn that the default (`true`) exports tool arguments and outputs to the OTLP backend; recommends `false` as the safe starting point. +- **`opentelemetry-sdk` version constraint**: tightened `[dev]` lower bound from `>=1.20` to `>=1.38.0` to match `traceloop-sdk`'s actual minimum; prevents pip resolving an incompatible version when both extras are installed. + ## [0.9.0] - 2026-04-10 ### Added diff --git a/README.md b/README.md index 05e406c..bc8fddc 100644 --- a/README.md +++ b/README.md @@ -532,6 +532,40 @@ In **prod mode**, stdout is reserved exclusively for MCP JSON-RPC messages. Logs --- +## Tracing (OpenLLMetry) + +Distributed tracing via [traceloop-sdk](https://github.com/traceloop/openllmetry) is available as an optional extra. Install it, then set `TRACELOOP_BASE_URL` to enable it: + +```bash +pip install 'openstudio-mcp[telemetry]' +``` + +Or with Docker (requires the tracing image built with `--build-arg TELEMETRY=1`): + +```bash +# Build the tracing-enabled image once: +docker build --build-arg TELEMETRY=1 -t openstudio-mcp:tracing -f docker/Dockerfile . + +# Run with tracing enabled: +docker run --rm -i \ + -e TRACELOOP_BASE_URL=http://host.docker.internal:4318 \ + openstudio-mcp:tracing openstudio-mcp +``` + +The standard `openstudio-mcp:dev` image does **not** include `traceloop-sdk`. Using it with `TRACELOOP_BASE_URL` set will log a warning and disable tracing — it will not work silently. + +| Variable | Default | Description | +|----------|---------|-------------| +| `TRACELOOP_BASE_URL` | *(unset — disabled)* | OTLP endpoint, e.g. `http://localhost:4318` or `https://api.traceloop.com` | +| `TRACELOOP_API_KEY` | *(unset)* | API key for Traceloop cloud (not needed for generic OTLP) | +| `OTEL_SERVICE_NAME` | `openstudio-mcp` | Service name on every span | +| `OTEL_EXPORT_BATCH` | `true` | Set `false` for synchronous export in development | +| `TRACELOOP_TRACE_CONTENT` | `true` | **Set `false` to protect privacy** — when `true`, tool arguments and outputs (including file paths and model data) are exported to the OTLP backend. Recommended: start with `false` and enable only if your backend is self-hosted or you have reviewed the data. | + +Tracing is **off by default** and has zero overhead when `TRACELOOP_BASE_URL` is unset. Key operations (`run_simulation`, `apply_measure`, `create_measure`, the three `create_*_building` variants, and `run_qaqc_checks`) emit named spans. Every FastMCP tool call is auto-instrumented via `McpInstrumentor`. + +--- + ## Architecture - **Transport:** stdio (container spawned by host) diff --git a/docker/Dockerfile b/docker/Dockerfile index 1a4ae9d..3e600f1 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -43,8 +43,13 @@ COPY pyproject.toml /repo/pyproject.toml COPY mcp_server /repo/mcp_server COPY docker /repo/docker +# TELEMETRY=1 installs traceloop-sdk + opentelemetry instrumentation for MCP. +# Default off to keep the image lean. Set to 1 for the tracing variant: +# docker build --build-arg TELEMETRY=1 -t openstudio-mcp:tracing -f docker/Dockerfile . +ARG TELEMETRY=0 RUN pip install --no-cache-dir -U pip \ - && pip install --no-cache-dir -e ".[dev]" + && pip install --no-cache-dir -e ".[dev]" \ + && if [ "$TELEMETRY" = "1" ]; then pip install --no-cache-dir -e ".[telemetry]"; fi # (Optional) If you want the container to include any other repo files too: # COPY . /repo diff --git a/docker/docker-compose.tracing.yml b/docker/docker-compose.tracing.yml new file mode 100644 index 0000000..d75f7e2 --- /dev/null +++ b/docker/docker-compose.tracing.yml @@ -0,0 +1,91 @@ +# OpenLLMetry / Traceloop tracing stack for openstudio-mcp +# +# Quick start: +# # 1. Build the tracing-enabled MCP image (adds traceloop-sdk): +# docker build --build-arg TELEMETRY=1 -t openstudio-mcp:tracing -f docker/Dockerfile . +# +# # 2. Start Jaeger: +# docker compose -f docker/docker-compose.tracing.yml up -d +# open http://localhost:16686 # Jaeger UI — traces appear here +# +# # 3. Configure your MCP client to use the tracing image on the shared network. +# Add these flags to your client's docker run command: +# +# -e TRACELOOP_BASE_URL=http://jaeger:4318 +# --network openstudio-mcp-tracing +# (and use openstudio-mcp:tracing instead of openstudio-mcp:dev) +# +# Example — Claude Code .mcp.json with tracing enabled: +# +# { +# "mcpServers": { +# "openstudio-mcp": { +# "command": "docker", +# "args": [ +# "run", "--rm", "-i", +# "-v", "/abs/path/inputs:/inputs", +# "-v", "/abs/path/runs:/runs", +# "--network", "openstudio-mcp-tracing", +# "-e", "OPENSTUDIO_MCP_MODE=prod", +# "-e", "TRACELOOP_BASE_URL=http://jaeger:4318", +# "-e", "OTEL_SERVICE_NAME=openstudio-mcp", +# "-e", "TRACELOOP_TRACE_CONTENT=true", +# "openstudio-mcp:tracing", "openstudio-mcp" +# ] +# } +# } +# } +# +# Environment variables understood by the MCP server (see mcp_server/telemetry.py): +# TRACELOOP_BASE_URL OTLP HTTP endpoint (required to enable telemetry) +# TRACELOOP_API_KEY API key for Traceloop cloud (omit for local Jaeger) +# OTEL_SERVICE_NAME Service name on spans (default: openstudio-mcp) +# OTEL_EXPORT_BATCH "false" for sync export in dev (default: batch) +# TRACELOOP_TRACE_CONTENT "false" to omit tool args from spans (privacy) + +services: + jaeger: + image: jaegertracing/jaeger:2.5.0 + container_name: openstudio-mcp-jaeger + ports: + - "16686:16686" # Jaeger UI + - "4317:4317" # OTLP gRPC receiver + - "4318:4318" # OTLP HTTP receiver (used by traceloop-sdk) + environment: + - SPAN_STORAGE_TYPE=memory + networks: + - tracing + restart: unless-stopped + healthcheck: + test: ["CMD", "wget", "--spider", "-q", "http://localhost:16686/"] + interval: 10s + timeout: 5s + retries: 5 + start_period: 10s + + # Optional: OpenTelemetry Collector as a middle layer. + # Useful if you want to fan-out to multiple backends (Jaeger + Prometheus + Loki). + # Comment out and point TRACELOOP_BASE_URL directly to jaeger:4318 for simplest setup. + otel-collector: + image: otel/opentelemetry-collector-contrib:0.120.0 + container_name: openstudio-mcp-otelcol + command: ["--config=/etc/otel/config.yaml"] + volumes: + - ./otel-collector-config.yaml:/etc/otel/config.yaml:ro + ports: + - "4319:4318" # OTLP HTTP (external port offset to avoid conflict with jaeger) + - "4320:4317" # OTLP gRPC (external port offset) + - "8888:8888" # Collector metrics (Prometheus scrape endpoint) + networks: + - tracing + depends_on: + jaeger: + condition: service_healthy + restart: unless-stopped + profiles: + - collector # only start with: docker compose --profile collector up + +networks: + tracing: + name: openstudio-mcp-tracing + driver: bridge diff --git a/docker/otel-collector-config.yaml b/docker/otel-collector-config.yaml new file mode 100644 index 0000000..4520179 --- /dev/null +++ b/docker/otel-collector-config.yaml @@ -0,0 +1,59 @@ +# OpenTelemetry Collector config for openstudio-mcp tracing stack. +# Used only when running: docker compose --profile collector up +# For simple setups, skip this and point TRACELOOP_BASE_URL directly to jaeger:4318. + +receivers: + otlp: + protocols: + grpc: + endpoint: 0.0.0.0:4317 + http: + endpoint: 0.0.0.0:4318 + +processors: + batch: + timeout: 1s + send_batch_size: 1024 + memory_limiter: + check_interval: 1s + limit_mib: 256 + spike_limit_mib: 64 + # Add service name tag from resource attributes + resource: + attributes: + - key: service.namespace + value: openstudio-mcp + action: insert + +exporters: + otlp/jaeger: + endpoint: jaeger:4317 + tls: + insecure: true + + # Uncomment to export to Traceloop cloud instead of / in addition to Jaeger: + # otlphttp/traceloop: + # endpoint: https://api.traceloop.com + # headers: + # Authorization: "Bearer ${TRACELOOP_API_KEY}" + + # Prometheus metrics from the collector itself (scrape at :8888/metrics) + prometheus: + endpoint: 0.0.0.0:8888 + + # Debug: print spans to collector stdout (useful for development) + debug: + verbosity: basic + sampling_initial: 5 + sampling_thereafter: 200 + +service: + pipelines: + traces: + receivers: [otlp] + processors: [memory_limiter, resource, batch] + exporters: [otlp/jaeger] + metrics: + receivers: [otlp] + processors: [memory_limiter, batch] + exporters: [prometheus] diff --git a/mcp_server/server.py b/mcp_server/server.py index 290156f..935f790 100644 --- a/mcp_server/server.py +++ b/mcp_server/server.py @@ -1,63 +1,68 @@ from __future__ import annotations -from fastmcp import FastMCP - from mcp_server.config import ENABLE_CODE_MODE from mcp_server.skills import register_all_skills from mcp_server.stdout_suppression import ( redirect_c_stdout_to_stderr, silence_openstudio_stdout_logger, ) - -mcp = FastMCP( - "openstudio-mcp", - instructions=( - "Building energy simulation server (OpenStudio SDK) with 142 tools for " - "creating, modifying, simulating, and analyzing building energy models. " - "Use these tools for all building energy modeling tasks — if no tool " - "exists for a task, ask the user before writing code. " - "NEVER write scripts, code, or files to accomplish tasks that these " - "tools already handle. Specifically: " - "- Measures: ALWAYS use create_measure — never write measure.rb/.py/.xml " - "directly. create_measure handles scaffolding, XML, checksums, and " - "OS App compatibility. Workflow: create_measure → test_measure → apply_measure. " - "- Results/data: use extract_summary_metrics, extract_end_use_breakdown, " - "query_timeseries, extract_envelope_summary, extract_hvac_sizing — " - "never write Python/SQL scripts to parse eplusout.sql. " - "- Visualization: use view_model (3D geometry), view_simulation_data " - "(charts/heatmaps), generate_results_report (HTML report) — never write " - "matplotlib/plotly/HTML scripts. " - "- Models: use create_new_building, create_bar_building, import_floorspacejs " - "— never write raw IDF or OSM files. " - "- Weather: use change_building_location (sets EPW+DDY+CZ in one call) " - "or list_weather_files — never download or write weather files. " - "- HVAC: use add_baseline_system, add_doas_system, add_vrf_system — " - "never write OpenStudio SDK scripts to wire HVAC components. " - "For custom HVAC measures, call search_wiring_patterns to get working " - "Ruby wiring code, and search_api to verify methods exist. " - "If a file path is given, use it directly. If a file operation fails, " - "you may call list_files once to find the right path, then retry — " - "do not call list_files more than once for the same file. " - "Use list_weather_files for EPW discovery — do not use list_files for weather. " - "To find objects by type, use list_model_objects(object_type). " - "List tools default to 10 results — use filters to narrow, or " - "max_results=0 for all. Prefer list tools before detail tools to " - "find the right name. " - "When polling get_run_status, wait at least 1-2 minutes between calls. " - "For multi-step workflows, call list_skills() first." - ), -) - -register_all_skills(mcp) - -if ENABLE_CODE_MODE: - from fastmcp.experimental.transforms.code_mode import CodeMode - mcp.add_transform(CodeMode()) +from mcp_server.telemetry import init_telemetry def main(): silence_openstudio_stdout_logger() redirect_c_stdout_to_stderr() + # init_telemetry() must run before FastMCP is instantiated so that + # McpInstrumentor().instrument() can patch FastMCP.__init__ in time. + init_telemetry() + + from fastmcp import FastMCP + + mcp = FastMCP( + "openstudio-mcp", + instructions=( + "Building energy simulation server (OpenStudio SDK) with 142 tools for " + "creating, modifying, simulating, and analyzing building energy models. " + "Use these tools for all building energy modeling tasks — if no tool " + "exists for a task, ask the user before writing code. " + "NEVER write scripts, code, or files to accomplish tasks that these " + "tools already handle. Specifically: " + "- Measures: ALWAYS use create_measure — never write measure.rb/.py/.xml " + "directly. create_measure handles scaffolding, XML, checksums, and " + "OS App compatibility. Workflow: create_measure → test_measure → apply_measure. " + "- Results/data: use extract_summary_metrics, extract_end_use_breakdown, " + "query_timeseries, extract_envelope_summary, extract_hvac_sizing — " + "never write Python/SQL scripts to parse eplusout.sql. " + "- Visualization: use view_model (3D geometry), view_simulation_data " + "(charts/heatmaps), generate_results_report (HTML report) — never write " + "matplotlib/plotly/HTML scripts. " + "- Models: use create_new_building, create_bar_building, import_floorspacejs " + "— never write raw IDF or OSM files. " + "- Weather: use change_building_location (sets EPW+DDY+CZ in one call) " + "or list_weather_files — never download or write weather files. " + "- HVAC: use add_baseline_system, add_doas_system, add_vrf_system — " + "never write OpenStudio SDK scripts to wire HVAC components. " + "For custom HVAC measures, call search_wiring_patterns to get working " + "Ruby wiring code, and search_api to verify methods exist. " + "If a file path is given, use it directly. If a file operation fails, " + "you may call list_files once to find the right path, then retry — " + "do not call list_files more than once for the same file. " + "Use list_weather_files for EPW discovery — do not use list_files for weather. " + "To find objects by type, use list_model_objects(object_type). " + "List tools default to 10 results — use filters to narrow, or " + "max_results=0 for all. Prefer list tools before detail tools to " + "find the right name. " + "When polling get_run_status, wait at least 1-2 minutes between calls. " + "For multi-step workflows, call list_skills() first." + ), + ) + + register_all_skills(mcp) + + if ENABLE_CODE_MODE: + from fastmcp.experimental.transforms.code_mode import CodeMode + mcp.add_transform(CodeMode()) + mcp.run() diff --git a/mcp_server/skills/common_measures/wrappers.py b/mcp_server/skills/common_measures/wrappers.py index 6b08625..caed813 100644 --- a/mcp_server/skills/common_measures/wrappers.py +++ b/mcp_server/skills/common_measures/wrappers.py @@ -11,6 +11,7 @@ from typing import Any from mcp_server.skills.measures.operations import apply_measure +from mcp_server.telemetry import traced def _ensure_climate_zone() -> None: @@ -154,6 +155,7 @@ def generate_results_report_op(units: str = "IP", run_id: str | None = None) -> # --- 4. run_qaqc_checks: ASHRAE QA/QC --- +@traced() def run_qaqc_checks_op( template: str = "90.1-2013", checks: list[str] | None = None, diff --git a/mcp_server/skills/comstock/operations.py b/mcp_server/skills/comstock/operations.py index 0a07fd0..03935a5 100644 --- a/mcp_server/skills/comstock/operations.py +++ b/mcp_server/skills/comstock/operations.py @@ -18,6 +18,7 @@ from mcp_server.config import RUN_ROOT from mcp_server.model_manager import get_model from mcp_server.skills.measures.operations import apply_measure +from mcp_server.telemetry import traced # Category classification for ComStock measures _BASELINE_PREFIXES = ( @@ -115,6 +116,7 @@ def list_comstock_measures(category: str | None = None) -> dict[str, Any]: } +@traced() def create_typical_building( template: str = "90.1-2019", building_type: str = "SmallOffice", @@ -299,6 +301,7 @@ def _create_empty_model() -> Path: return osm_path +@traced() def create_bar_building( building_type: str = "SmallOffice", total_bldg_floor_area: float = 10000, @@ -407,6 +410,7 @@ def create_bar_building( return result +@traced() def create_new_building( # Bar geometry args building_type: str = "SmallOffice", diff --git a/mcp_server/skills/measure_authoring/operations.py b/mcp_server/skills/measure_authoring/operations.py index 189ab9d..ffccda9 100644 --- a/mcp_server/skills/measure_authoring/operations.py +++ b/mcp_server/skills/measure_authoring/operations.py @@ -14,6 +14,7 @@ import openstudio from mcp_server.config import INPUT_ROOT, RUN_ROOT +from mcp_server.telemetry import traced CUSTOM_MEASURES_DIR = RUN_ROOT / "custom_measures" @@ -740,6 +741,7 @@ def _write_test_file(measure_dir: Path, class_name: str, args: list[dict], # ── Public operations ──────────────────────────────────────────────── +@traced() def create_measure_op( name: str, description: str, diff --git a/mcp_server/skills/measures/operations.py b/mcp_server/skills/measures/operations.py index 52a0c4d..a1238a2 100644 --- a/mcp_server/skills/measures/operations.py +++ b/mcp_server/skills/measures/operations.py @@ -18,6 +18,7 @@ from mcp_server.config import OSCLI_GEM_PATH, OSCLI_GEMFILE, RUN_ROOT from mcp_server.model_manager import get_model, load_model +from mcp_server.telemetry import traced from mcp_server.util import resolve_run_dir @@ -115,6 +116,7 @@ def _parse_runner_messages(out_osw_path: Path) -> dict[str, Any] | None: return None +@traced() def apply_measure( measure_dir: str, arguments: dict[str, Any] | None = None, diff --git a/mcp_server/skills/simulation/operations.py b/mcp_server/skills/simulation/operations.py index 46bdfcd..34673bd 100644 --- a/mcp_server/skills/simulation/operations.py +++ b/mcp_server/skills/simulation/operations.py @@ -14,6 +14,7 @@ import psutil from mcp_server.config import LOG_TAIL_DEFAULT, OSCLI_GEM_PATH, OSCLI_GEMFILE, RUN_ROOT +from mcp_server.telemetry import traced from mcp_server.util import resolve_run_dir # Where the MCP server stores runs inside the container @@ -603,6 +604,7 @@ def validate_model_op() -> dict[str, Any]: } +@traced() def run_simulation(osm_path: str, epw_path: str | None = None, name: str | None = None) -> dict[str, Any]: """Create a minimal OSW from an OSM file and run the simulation. diff --git a/mcp_server/telemetry.py b/mcp_server/telemetry.py new file mode 100644 index 0000000..3691bc2 --- /dev/null +++ b/mcp_server/telemetry.py @@ -0,0 +1,250 @@ +"""OpenLLMetry (Traceloop) tracing for the openstudio-mcp server. + +Optional: a no-op unless traceloop-sdk is installed (included in [telemetry] extra). +Zero overhead when absent: no import errors, all calls become pass-throughs. + +Install: + pip install 'openstudio-mcp[telemetry]' + +Environment variables: + TRACELOOP_BASE_URL OTLP / Traceloop-compatible endpoint, e.g.: + http://localhost:4318 (local OTEL collector) + https://api.traceloop.com (Traceloop cloud, needs API key) + Unset -> telemetry disabled (no-op). + TRACELOOP_API_KEY API key for Traceloop cloud (not required for generic OTLP). + OTEL_SERVICE_NAME Service name emitted on every span. Default: "openstudio-mcp". + OTEL_EXPORT_BATCH "false" -> sync exporting (dev). Default: batch mode. + TRACELOOP_TRACE_CONTENT "false" -> omit tool args from spans (privacy). + +Usage: + from mcp_server.telemetry import init_telemetry, trace_operation, traced + + # In main() before mcp.run(): + init_telemetry() + + # Decorate a key operation: + @traced() + def run_simulation(osm_path: str, ...) -> dict: ... + + # Or use a context manager for finer control: + with trace_operation("prepare_model", {"path": osm_path}) as span: + result = do_work() +""" +from __future__ import annotations + +import importlib.util +import json +import logging +import sys +from contextlib import contextmanager +from typing import Any, Callable, TypeVar + +logger = logging.getLogger(__name__) + +_TELEMETRY_INITIALIZED = False +# True only after Traceloop.init() succeeds with a valid endpoint. +# traced() checks this at call time to avoid traceloop stdout warnings. +_TELEMETRY_ENABLED = False +try: + _SDK_AVAILABLE = importlib.util.find_spec("traceloop.sdk") is not None +except (ModuleNotFoundError, ValueError): + _SDK_AVAILABLE = False + +# Max chars for any single span attribute value. +_MAX_ATTR_LEN = 512 + +F = TypeVar("F", bound=Callable[..., Any]) + + +class _NoopSpan: + """Minimal no-op span used when opentelemetry is not installed.""" + + def set_attribute(self, *args: Any, **kwargs: Any) -> None: + pass + + def set_status(self, *args: Any, **kwargs: Any) -> None: + pass + + def record_exception(self, *args: Any, **kwargs: Any) -> None: + pass + + +def init_telemetry() -> bool: + """Initialize OpenLLMetry tracing. Idempotent — safe to call multiple times. + + Returns True if telemetry was enabled, False otherwise (SDK absent, no endpoint). + When the SDK is installed and an endpoint is configured, calls + McpInstrumentor().instrument() to auto-trace every FastMCP tool call and + Traceloop.init() to configure the OTLP exporter. + """ + global _TELEMETRY_INITIALIZED, _TELEMETRY_ENABLED + + if _TELEMETRY_INITIALIZED: + return _TELEMETRY_ENABLED + + import os + + if not _SDK_AVAILABLE: + endpoint = os.environ.get("TRACELOOP_BASE_URL", "").strip() + if endpoint: + logger.warning( + "TRACELOOP_BASE_URL is set but traceloop-sdk is not installed. " + "Install telemetry extras: pip install 'openstudio-mcp[telemetry]'" + ) + _TELEMETRY_INITIALIZED = True + return False + + endpoint = os.environ.get("TRACELOOP_BASE_URL", "").strip() + if not endpoint: + logger.debug("TRACELOOP_BASE_URL not set -- telemetry disabled") + _TELEMETRY_INITIALIZED = True + return False + + try: + from opentelemetry.instrumentation.mcp import McpInstrumentor + from traceloop.sdk import Traceloop + + service_name = os.environ.get("OTEL_SERVICE_NAME", "openstudio-mcp") + disable_batch = os.environ.get("OTEL_EXPORT_BATCH", "true").lower() == "false" + + # Initialize Traceloop FIRST so its TracerProvider is live before we + # patch FastMCP. McpInstrumentor wraps FastMCP tool calls; if the + # provider isn't established yet those spans have nowhere to go. + # Traceloop.init() uses print() for status messages — redirect sys.stdout + # to stderr to avoid corrupting the MCP JSON-RPC stdio pipe. + _orig_stdout = sys.stdout + sys.stdout = sys.stderr + try: + Traceloop.init( + app_name=service_name, + api_endpoint=endpoint, + disable_batch=disable_batch, + ) + finally: + sys.stdout = _orig_stdout + + # Patch FastMCP AFTER the provider is live so auto-traced tool calls + # have a real exporter destination. + McpInstrumentor().instrument() + + _TELEMETRY_INITIALIZED = True + _TELEMETRY_ENABLED = True + logger.info( + "OpenLLMetry enabled: endpoint=%s service=%s batch=%s", + endpoint, + service_name, + not disable_batch, + ) + return True + + except Exception: + logger.exception("Failed to initialize OpenLLMetry -- telemetry disabled") + _TELEMETRY_INITIALIZED = True + return False + + +@contextmanager +def trace_operation(name: str, attributes: dict[str, Any] | None = None): + """Context manager that wraps a block in a child INTERNAL span. + + Uses the active OpenTelemetry TracerProvider (configured by Traceloop.init()). + Falls back to a no-op span when opentelemetry is not installed or telemetry + is not configured — safe to use unconditionally. + + Args: + name: Span name, e.g. "prepare_model". + attributes: Optional initial attributes (values truncated to _MAX_ATTR_LEN). + + Yields: + The active Span (may be a NonRecordingSpan or _NoopSpan when telemetry is off). + + Example:: + + with trace_operation("apply_measure", {"measure_dir": measure_dir}) as span: + result = _do_apply(...) + span.set_attribute("ok", str(result.get("ok", False))) + """ + try: + from opentelemetry import trace + from opentelemetry.trace import NonRecordingSpan, SpanKind, StatusCode + except ImportError: + yield _NoopSpan() + return + + tracer = trace.get_tracer("openstudio-mcp") + with tracer.start_as_current_span(name, kind=SpanKind.INTERNAL) as span: + is_recording = not isinstance(span, NonRecordingSpan) + if is_recording and attributes: + for key, val in attributes.items(): + span.set_attribute(key, _truncate(val)) + try: + yield span + except Exception as exc: + if is_recording: + span.record_exception(exc) + span.set_status(StatusCode.ERROR, str(exc)) + raise + + +def traced(op_name: str | None = None) -> Callable[[F], F]: + """Decorator that wraps a synchronous operation in a trace span. + + Uses trace_operation() context manager to create a span. Only active when + telemetry has been successfully enabled via init_telemetry(). This avoids + traceloop stdout warnings when the SDK is installed but no endpoint is set. + + Marks the span ERROR when the function returns a dict with ok=False. + + Args: + op_name: Span name override. Defaults to the function name. + + Example:: + + @traced() + def run_simulation(osm_path: str, ...) -> dict: ... + """ + import functools + + def decorator(fn: F) -> F: + span_name = op_name or fn.__name__ + + @functools.wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + if not _TELEMETRY_ENABLED: + return fn(*args, **kwargs) + + with trace_operation(span_name) as span: + result = fn(*args, **kwargs) + if isinstance(result, dict) and result.get("ok") is False: + _mark_span_error(span, result) + return result + + return wrapper # type: ignore[return-value] + + return decorator + + +def _mark_span_error(span: Any, result: dict[str, Any]) -> None: + """Set ERROR status on the given span.""" + try: + from opentelemetry.trace import NonRecordingSpan, StatusCode + + if isinstance(span, NonRecordingSpan): + return + error_msg = result.get("error") or result.get("message") or "tool returned ok=False" + span.set_status(StatusCode.ERROR, str(error_msg)) + span.set_attribute("error.message", str(error_msg)[:_MAX_ATTR_LEN]) + except Exception: + pass + + +def _truncate(value: Any) -> str: + """Serialize a value to a JSON string capped at _MAX_ATTR_LEN chars.""" + try: + s = json.dumps(value, default=str) + except Exception: + s = str(value) + if len(s) > _MAX_ATTR_LEN: + return s[:_MAX_ATTR_LEN] + "..." + return s + diff --git a/pyproject.toml b/pyproject.toml index 3e67bef..5e32ca8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,10 @@ dev = [ "pytest-timeout>=2.3.1", "mcp", "pyyaml>=6.0", + "opentelemetry-sdk>=1.38.0", +] +telemetry = [ + "traceloop-sdk>=0.49.2", ] [project.scripts] diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py new file mode 100644 index 0000000..8ea05cd --- /dev/null +++ b/tests/test_telemetry.py @@ -0,0 +1,576 @@ +"""Unit tests for mcp_server/telemetry.py (OpenLLMetry / traceloop-sdk integration). + +These tests run without OpenStudio and without Docker. + +Validates: +- Telemetry is a no-op when TRACELOOP_BASE_URL is not set +- init_telemetry() calls McpInstrumentor().instrument() and Traceloop.init() +- init_telemetry() is idempotent and returns correct value on second call +- McpInstrumentor is only called when endpoint is configured +- traced() is a no-op when telemetry is not enabled +- traced() creates a span and marks ERROR on ok=False when telemetry is enabled +- trace_operation() creates a child span when a TracerProvider is configured +- _truncate() caps values at _MAX_ATTR_LEN +- sys.stdout is restored even when Traceloop.init() raises +- init_telemetry() handles Traceloop.init() exceptions gracefully + +Regression: these tests guard against the telemetry module breaking server +startup or silently swallowing init errors. +""" +from __future__ import annotations + +import sys +from contextlib import contextmanager +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +@contextmanager +def _reset_telemetry_module(): + """Force a clean re-import of telemetry so _TELEMETRY_INITIALIZED resets.""" + mod_name = "mcp_server.telemetry" + old = sys.modules.pop(mod_name, None) + try: + yield + finally: + sys.modules.pop(mod_name, None) + if old is not None: + sys.modules[mod_name] = old + + +def _make_in_memory_setup(): + """Return (provider, exporter, tracer) using the OTel SDK InMemorySpanExporter.""" + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = provider.get_tracer("test") + return provider, exporter, tracer + + +# --------------------------------------------------------------------------- +# init_telemetry tests +# --------------------------------------------------------------------------- + +@pytest.mark.unit +def test_init_no_endpoint_returns_false(monkeypatch): + # Validates: init_telemetry returns False when TRACELOOP_BASE_URL is unset, and does NOT call McpInstrumentor or Traceloop.init. + monkeypatch.delenv("TRACELOOP_BASE_URL", raising=False) + + mock_traceloop = MagicMock() + mock_instrumentor = MagicMock() + + with _reset_telemetry_module(): + with patch.dict("sys.modules", { + "traceloop": MagicMock(), + "traceloop.sdk": mock_traceloop, + "traceloop.sdk.decorators": MagicMock(), + "opentelemetry.instrumentation.mcp": mock_instrumentor, + }): + import mcp_server.telemetry as tel + tel._SDK_AVAILABLE = True + result = tel.init_telemetry() + + assert result is False + mock_traceloop.Traceloop.init.assert_not_called() + # McpInstrumentor should NOT be called when no endpoint is set + mock_instrumentor.McpInstrumentor.assert_not_called() + + +@pytest.mark.unit +def test_init_no_sdk_returns_false(monkeypatch): + # Validates: init_telemetry returns False (no warning) when SDK is absent and no endpoint is configured. + monkeypatch.delenv("TRACELOOP_BASE_URL", raising=False) + + with _reset_telemetry_module(): + import mcp_server.telemetry as tel + tel._SDK_AVAILABLE = False + result = tel.init_telemetry() + + assert result is False + + +@pytest.mark.unit +def test_init_sdk_missing_with_endpoint_logs_warning(monkeypatch, caplog): + # Validates: a warning is logged when endpoint is set but SDK is not installed. + import logging + monkeypatch.setenv("TRACELOOP_BASE_URL", "http://localhost:4318") + + with _reset_telemetry_module(): + import mcp_server.telemetry as tel + tel._SDK_AVAILABLE = False + with caplog.at_level(logging.WARNING, logger="mcp_server.telemetry"): + tel.init_telemetry() + + assert any("traceloop-sdk is not installed" in r.message for r in caplog.records) + + +@pytest.mark.unit +def test_init_instruments_mcp_and_calls_traceloop_init(monkeypatch): + # Validates: when endpoint is set, McpInstrumentor().instrument() and Traceloop.init() are both called. + monkeypatch.setenv("TRACELOOP_BASE_URL", "http://localhost:4318") + monkeypatch.setenv("OTEL_SERVICE_NAME", "test-svc") + + mock_traceloop_class = MagicMock() + mock_instrumentor_class = MagicMock() + mock_instrumentor_instance = MagicMock() + mock_instrumentor_class.return_value = mock_instrumentor_instance + + mock_otel_mcp_mod = MagicMock() + mock_otel_mcp_mod.McpInstrumentor = mock_instrumentor_class + + mock_traceloop_mod = MagicMock() + mock_traceloop_mod.Traceloop = mock_traceloop_class + + with _reset_telemetry_module(): + with patch.dict("sys.modules", { + "traceloop": MagicMock(), + "traceloop.sdk": mock_traceloop_mod, + "traceloop.sdk.decorators": MagicMock(), + "opentelemetry.instrumentation.mcp": mock_otel_mcp_mod, + }): + import mcp_server.telemetry as tel + tel._SDK_AVAILABLE = True + result = tel.init_telemetry() + + assert result is True + mock_instrumentor_instance.instrument.assert_called_once() + mock_traceloop_class.init.assert_called_once() + _, kwargs = mock_traceloop_class.init.call_args + assert kwargs["app_name"] == "test-svc" + assert kwargs["api_endpoint"] == "http://localhost:4318" + + +@pytest.mark.unit +def test_init_idempotent(monkeypatch): + # Validates: calling init_telemetry twice only initializes once. + monkeypatch.setenv("TRACELOOP_BASE_URL", "http://localhost:4318") + + mock_traceloop_class = MagicMock() + mock_instrumentor_class = MagicMock() + mock_instrumentor_instance = MagicMock() + mock_instrumentor_class.return_value = mock_instrumentor_instance + + mock_otel_mcp_mod = MagicMock() + mock_otel_mcp_mod.McpInstrumentor = mock_instrumentor_class + + mock_traceloop_mod = MagicMock() + mock_traceloop_mod.Traceloop = mock_traceloop_class + + with _reset_telemetry_module(): + with patch.dict("sys.modules", { + "traceloop": MagicMock(), + "traceloop.sdk": mock_traceloop_mod, + "traceloop.sdk.decorators": MagicMock(), + "opentelemetry.instrumentation.mcp": mock_otel_mcp_mod, + }): + import mcp_server.telemetry as tel + tel._SDK_AVAILABLE = True + r1 = tel.init_telemetry() + r2 = tel.init_telemetry() + + assert mock_traceloop_class.init.call_count == 1 + assert r1 is True + assert r2 is True + + +@pytest.mark.unit +def test_init_idempotent_returns_false_when_disabled(monkeypatch): + # Validates: second call returns False (not True) when first call disabled telemetry. + monkeypatch.delenv("TRACELOOP_BASE_URL", raising=False) + + with _reset_telemetry_module(): + import mcp_server.telemetry as tel + tel._SDK_AVAILABLE = True + r1 = tel.init_telemetry() + r2 = tel.init_telemetry() + + assert r1 is False + assert r2 is False + + +@pytest.mark.unit +def test_init_disable_batch_flag(monkeypatch): + # Validates: OTEL_EXPORT_BATCH=false sets disable_batch=True in Traceloop.init. + monkeypatch.setenv("TRACELOOP_BASE_URL", "http://localhost:4318") + monkeypatch.setenv("OTEL_EXPORT_BATCH", "false") + + mock_traceloop_class = MagicMock() + mock_traceloop_mod = MagicMock() + mock_traceloop_mod.Traceloop = mock_traceloop_class + mock_otel_mcp_mod = MagicMock() + mock_otel_mcp_mod.McpInstrumentor.return_value = MagicMock() + + with _reset_telemetry_module(): + with patch.dict("sys.modules", { + "traceloop": MagicMock(), + "traceloop.sdk": mock_traceloop_mod, + "traceloop.sdk.decorators": MagicMock(), + "opentelemetry.instrumentation.mcp": mock_otel_mcp_mod, + }): + import mcp_server.telemetry as tel + tel._SDK_AVAILABLE = True + tel.init_telemetry() + + _, kwargs = mock_traceloop_class.init.call_args + assert kwargs["disable_batch"] is True + + +@pytest.mark.unit +def test_init_restores_stdout_on_exception(monkeypatch): + # Validates: sys.stdout is restored even when Traceloop.init() raises. + monkeypatch.setenv("TRACELOOP_BASE_URL", "http://localhost:4318") + + mock_traceloop_class = MagicMock() + mock_traceloop_class.init.side_effect = RuntimeError("init boom") + mock_traceloop_mod = MagicMock() + mock_traceloop_mod.Traceloop = mock_traceloop_class + mock_otel_mcp_mod = MagicMock() + mock_otel_mcp_mod.McpInstrumentor.return_value = MagicMock() + + original_stdout = sys.stdout + + with _reset_telemetry_module(): + with patch.dict("sys.modules", { + "traceloop": MagicMock(), + "traceloop.sdk": mock_traceloop_mod, + "traceloop.sdk.decorators": MagicMock(), + "opentelemetry.instrumentation.mcp": mock_otel_mcp_mod, + }): + import mcp_server.telemetry as tel + tel._SDK_AVAILABLE = True + result = tel.init_telemetry() + + assert result is False + assert sys.stdout is original_stdout + + +# --------------------------------------------------------------------------- +# _truncate tests +# --------------------------------------------------------------------------- + +@pytest.mark.unit +def test_truncate_short_value(): + # Validates: short values are returned unchanged. + from mcp_server.telemetry import _MAX_ATTR_LEN, _truncate + assert _truncate("hello") == '"hello"' + assert len(_truncate("hello")) < _MAX_ATTR_LEN + + +@pytest.mark.unit +def test_truncate_long_value(): + # Validates: values longer than _MAX_ATTR_LEN are capped with ellipsis. + from mcp_server.telemetry import _MAX_ATTR_LEN, _truncate + long_val = "x" * 2000 + result = _truncate(long_val) + assert len(result) <= _MAX_ATTR_LEN + 10 # small slack for the suffix + assert result.endswith("...") + + +# --------------------------------------------------------------------------- +# trace_operation tests +# --------------------------------------------------------------------------- + +@pytest.mark.unit +def test_trace_operation_noop_when_no_provider(): + # Validates: trace_operation is safe to call when no provider is configured. + from mcp_server.telemetry import trace_operation + ran = [] + with trace_operation("test_op") as span: + ran.append(True) + # NonRecordingSpan.set_attribute is a no-op -- this must not crash + span.set_attribute("key", "value") + assert ran == [True] + + +@pytest.mark.unit +def test_trace_operation_noop_span_on_import_error(): + # Validates: trace_operation yields a _NoopSpan when opentelemetry is absent. + # Regression: trace_operation() must not raise ImportError in production + # environments where dev extras (opentelemetry-api) are not installed. + import sys + + from mcp_server.telemetry import _NoopSpan, trace_operation + + # Simulate opentelemetry being absent + otel_keys = {"opentelemetry", "opentelemetry.trace"} + with patch.dict(sys.modules, dict.fromkeys(otel_keys, None)): + ran = [] + with trace_operation("test_noop") as span: + ran.append(True) + assert isinstance(span, _NoopSpan) + span.set_attribute("key", "value") + span.set_status("ok") + span.record_exception(None) + assert ran == [True] + + +@pytest.mark.unit +def test_trace_operation_child_span(): + # Validates: trace_operation creates a named child span when a TracerProvider is configured. + from opentelemetry import trace as otel_trace + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + + from mcp_server.telemetry import trace_operation + with patch.object(otel_trace, "get_tracer", return_value=provider.get_tracer("t")): + with trace_operation("my_op", {"key": "val"}): + pass + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == "my_op" + assert spans[0].attributes.get("key") == '"val"' + + +@pytest.mark.unit +def test_trace_operation_records_exception(): + # Validates: trace_operation sets ERROR status when an exception is raised. + from opentelemetry import trace as otel_trace + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + from opentelemetry.trace import StatusCode + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + + from mcp_server.telemetry import trace_operation + with patch.object(otel_trace, "get_tracer", return_value=provider.get_tracer("t")): + with pytest.raises(ValueError): + with trace_operation("failing_op"): + raise ValueError("boom") + + spans = exporter.get_finished_spans() + assert spans[0].status.status_code == StatusCode.ERROR + + +# --------------------------------------------------------------------------- +# traced() decorator tests +# --------------------------------------------------------------------------- + +@pytest.mark.unit +def test_traced_noop_when_telemetry_disabled(): + # Validates: traced() wrapper calls the original function directly when _TELEMETRY_ENABLED is False (SDK installed but no endpoint configured). + with _reset_telemetry_module(): + import mcp_server.telemetry as tel + tel._TELEMETRY_ENABLED = False + + call_log: list[str] = [] + + def my_fn(x: int) -> dict: + call_log.append("called") + return {"ok": True, "value": x} + + wrapped = tel.traced()(my_fn) + result = wrapped(42) + + assert result == {"ok": True, "value": 42} + assert call_log == ["called"] + + +@pytest.mark.unit +def test_traced_creates_span_when_enabled(): + # Validates: traced() creates a span via trace_operation when _TELEMETRY_ENABLED is True. + from opentelemetry import trace as otel_trace + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + + with _reset_telemetry_module(): + import mcp_server.telemetry as tel + tel._TELEMETRY_ENABLED = True + + @tel.traced(op_name="custom_name") + def my_fn() -> dict: + return {"ok": True} + + with patch.object(otel_trace, "get_tracer", return_value=provider.get_tracer("t")): + result = my_fn() + + assert result == {"ok": True} + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == "custom_name" + + +@pytest.mark.unit +def test_traced_marks_error_on_ok_false(): + # Validates: traced() marks the span ERROR when result has ok=False. + from opentelemetry import trace as otel_trace + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + from opentelemetry.trace import StatusCode + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + + with _reset_telemetry_module(): + import mcp_server.telemetry as tel + tel._TELEMETRY_ENABLED = True + + @tel.traced() + def failing_op() -> dict: + return {"ok": False, "error": "something failed"} + + with patch.object(otel_trace, "get_tracer", return_value=provider.get_tracer("t")): + result = failing_op() + + assert result == {"ok": False, "error": "something failed"} + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].status.status_code == StatusCode.ERROR + assert spans[0].attributes.get("error.message") == "something failed" + + +@pytest.mark.unit +def test_traced_uses_function_name_as_default_span_name(): + # Validates: traced() uses the function name when op_name is not specified. + from opentelemetry import trace as otel_trace + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + + with _reset_telemetry_module(): + import mcp_server.telemetry as tel + tel._TELEMETRY_ENABLED = True + + @tel.traced() + def run_simulation(osm_path: str) -> dict: + return {"ok": True} + + with patch.object(otel_trace, "get_tracer", return_value=provider.get_tracer("t")): + run_simulation("/tmp/test.osm") + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == "run_simulation" + + +# --------------------------------------------------------------------------- +# Startup wiring regression — init_telemetry() must appear before FastMCP +# is instantiated in server.main(). Checked via AST so this runs without +# importing server.py (which requires /runs to exist and openstudio). +# --------------------------------------------------------------------------- + +@pytest.mark.unit +def test_server_main_calls_init_telemetry_before_fastmcp(): + # Regression: init_telemetry() must be called before FastMCP is + # instantiated inside server.main(). If the order is reversed, + # McpInstrumentor cannot patch FastMCP.__init__ and all auto- + # instrumentation silently stops working. + import ast + from pathlib import Path + + server_src = (Path(__file__).parent.parent / "mcp_server" / "server.py").read_text() + tree = ast.parse(server_src) + + main_fn = next( + (n for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == "main"), + None, + ) + assert main_fn is not None, "main() not found in mcp_server/server.py" + + init_telemetry_line = None + fastmcp_line = None + + for node in ast.walk(main_fn): + if ( + isinstance(node, ast.Expr) + and isinstance(node.value, ast.Call) + and isinstance(node.value.func, ast.Name) + and node.value.func.id == "init_telemetry" + ): + init_telemetry_line = node.lineno + + if isinstance(node, ast.Call): + func = node.func + if (isinstance(func, ast.Name) and func.id == "FastMCP") or ( + isinstance(func, ast.Attribute) and func.attr == "FastMCP" + ): + fastmcp_line = node.lineno + + assert init_telemetry_line is not None, ( + "init_telemetry() call not found in server.main(). " + "It must be called before FastMCP is instantiated." + ) + assert fastmcp_line is not None, "FastMCP() instantiation not found in server.main()." + assert init_telemetry_line < fastmcp_line, ( + f"init_telemetry() (line {init_telemetry_line}) must come BEFORE " + f"FastMCP() (line {fastmcp_line}) in server.main(). " + "McpInstrumentor patches FastMCP.__init__ during init_telemetry(); " + "reversing the order silently disables all auto-instrumentation." + ) + + +# --------------------------------------------------------------------------- +# Decorator coverage regression — @traced() must remain on all promised ops. +# Checked via AST to avoid importing skill modules that require openstudio. +# --------------------------------------------------------------------------- + +@pytest.mark.unit +def test_traced_decorator_applied_to_promised_operations(): + # Regression: the CHANGELOG and README both promise that specific operations + # emit spans. This test guards against accidental decorator removal by + # inspecting the source AST of each operations file. + import ast + from pathlib import Path + + repo_root = Path(__file__).parent.parent + expected = [ + ("mcp_server/skills/simulation/operations.py", "run_simulation"), + ("mcp_server/skills/measures/operations.py", "apply_measure"), + ("mcp_server/skills/measure_authoring/operations.py", "create_measure_op"), + ("mcp_server/skills/comstock/operations.py", "create_typical_building"), + ("mcp_server/skills/comstock/operations.py", "create_bar_building"), + ("mcp_server/skills/comstock/operations.py", "create_new_building"), + ("mcp_server/skills/common_measures/wrappers.py", "run_qaqc_checks_op"), + ] + + missing = [] + for rel_path, func_name in expected: + src = (repo_root / rel_path).read_text() + tree = ast.parse(src) + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name: + has_traced = any( + (isinstance(d, ast.Call) and isinstance(d.func, ast.Name) and d.func.id == "traced") + or (isinstance(d, ast.Call) and isinstance(d.func, ast.Attribute) and d.func.attr == "traced") + for d in node.decorator_list + ) + if not has_traced: + missing.append(f"{rel_path}::{func_name}") + break + else: + missing.append(f"{rel_path}::{func_name} (function not found)") + + assert not missing, ( + f"The following operations are missing @traced() decoration: {missing}. " + "Every operation listed in the CHANGELOG/README tracing section must be " + "wrapped with @traced() so it emits a span when telemetry is enabled." + )