From f34a5aa0e4ed360c82ba38d35bfd6ba0b8ee3cd4 Mon Sep 17 00:00:00 2001 From: Shreyas Nagaraj Date: Fri, 24 Jul 2026 17:40:48 +0530 Subject: [PATCH] feat: make instrumentation explicit opt-in with HARNESS_ env prefix Instrumentation is now opt-in instead of opt-out. Nothing is instrumented unless enabled via env flags: - HARNESS_ENABLE_API: all non-AI instrumentation (HTTP servers/clients, gRPC, DB, botocore, generic OTel contrib fallback) - HARNESS_ENABLE_AI_{OPENAI,ANTHROPIC,LITELLM,GOOGLE_GENAI,MCP}: per-provider AI Migrate the SDK env prefix to HARNESS_ while keeping HA_/AT_/TA_ as backwards-compatible aliases (precedence HARNESS_ > HA_ > AT_ > TA_). The new enable flags are strict HARNESS_-only with no legacy aliases, and the legacy HA_GEN_AI_ENABLED master switch no longer controls instrumentation. skip_libraries still takes precedence over enable flags. Co-authored-by: Cursor --- AGENTS.md | 42 +++-- README.md | 36 +++- docs/sdk-quickstart.md | 153 +++++++++++++++++ src/harness_sdk/agent.py | 11 +- src/harness_sdk/agent_init.py | 7 +- src/harness_sdk/config/config.py | 20 +-- src/harness_sdk/config/environment.py | 9 +- src/harness_sdk/env.py | 22 ++- .../instrumentation/anthropic/__init__.py | 19 --- .../instrumentation/google_genai/__init__.py | 15 -- .../instrumentation_definitions.py | 41 +++++ .../instrumentation/litellm/__init__.py | 24 +-- .../instrumentation/mcp/__init__.py | 14 +- .../instrumentation/mcp/gen_ai_mirror.py | 3 - .../instrumentation/openai/__init__.py | 17 -- src/harness_sdk/plugins/builtin/pipeline.py | 8 +- test/config/environment_test.py | 33 +++- test/conftest.py | 12 +- .../anthropic/test_anthropic.py | 5 +- .../instrumentation/botocore/test_botocore.py | 3 + .../litellm/litellm_instrumentation_test.py | 5 +- .../instrumentation/mcp/test_gen_ai_mirror.py | 9 - test/instrumentation/mcp/test_mcp_wrapper.py | 29 +--- .../openai/openai_instrumentation_test.py | 17 +- test/instrumentation/test_opt_in_gating.py | 161 ++++++++++++++++++ 25 files changed, 537 insertions(+), 178 deletions(-) create mode 100644 docs/sdk-quickstart.md create mode 100644 test/instrumentation/test_opt_in_gating.py diff --git a/AGENTS.md b/AGENTS.md index 0197b06..e3757c6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -21,24 +21,38 @@ RUN_SDK_INTEGRATION_TESTS=1 ./scripts/run-unit-tests.sh ## Environment variable naming -All SDK config uses the `HA_` prefix. `AT_` and `TA_` are legacy aliases accepted during migration. +All SDK config uses the `HARNESS_` prefix. `HA_`, `AT_`, and `TA_` are legacy aliases accepted for backwards compatibility. When a setting is defined under multiple prefixes, precedence is `HARNESS_` > `HA_` > `AT_` > `TA_`. Resolution lives in `src/harness_sdk/env.py` (`get_env_value`, `is_env_var_present`, `is_harness_flag_enabled`). Key variables: | Variable | Purpose | |---|---| -| `HA_SERVICE_NAME` | Service name on all spans | -| `HA_REPORTING_ENDPOINT` | OTLP endpoint URL | -| `HA_REPORTING_TOKEN` | Auth token (`x-harness-service-token` header) | -| `HA_REPORTING_TRACE_REPORTER_TYPE` | `OTLP` (gRPC) or `OTLP_HTTP` | -| `HA_REPORTING_SECURE` | `true`/`false` for TLS | -| `HA_REPORTING_COMPRESSION` | `gzip` or empty | -| `HA_CONTROL_PLUGINS` | Comma-separated control plugin names | -| `HA_OBSERVABILITY_PLUGINS` | Comma-separated observability plugin names | -| `HA_ENABLE_CONSOLE_SPAN_EXPORTER` | Set to any value to dump spans to stdout | -| `HA_CONFIG_FILE` | Path to YAML config file (overrides env) | -| `HA_GEN_AI_ENABLED` | Enable/disable GenAI instrumentation | -| `HA_GEN_AI_PAYLOAD_CAPTURE_ENABLED` | Capture LLM prompt/response payloads | -| `HA_GEN_AI_PAYLOAD_EVALUATION_ENABLED` | Run control plugins on GenAI spans | +| `HARNESS_SERVICE_NAME` | Service name on all spans | +| `HARNESS_REPORTING_ENDPOINT` | OTLP endpoint URL | +| `HARNESS_REPORTING_TOKEN` | Auth token (`x-harness-service-token` header) | +| `HARNESS_REPORTING_TRACE_REPORTER_TYPE` | `OTLP` (gRPC) or `OTLP_HTTP` | +| `HARNESS_REPORTING_SECURE` | `true`/`false` for TLS | +| `HARNESS_REPORTING_COMPRESSION` | `gzip` or empty | +| `HARNESS_CONTROL_PLUGINS` | Comma-separated control plugin names | +| `HARNESS_OBSERVABILITY_PLUGINS` | Comma-separated observability plugin names | +| `HARNESS_ENABLE_CONSOLE_SPAN_EXPORTER` | Set to any value to dump spans to stdout | +| `HARNESS_CONFIG_FILE` | Path to YAML config file (overrides env) | +| `HARNESS_GEN_AI_PAYLOAD_CAPTURE_ENABLED` | Capture LLM prompt/response payloads | +| `HARNESS_GEN_AI_PAYLOAD_EVALUATION_ENABLED` | Run control plugins on GenAI spans | + +### Instrumentation opt-in (strict `HARNESS_` prefix, no legacy aliases) + +Instrumentation is opt-in: `Agent().instrument()` instruments nothing unless a flag below is set to `true`. Categorization and gating live in `src/harness_sdk/instrumentation/instrumentation_definitions.py` (`is_library_enabled`, `is_api_instrumentation_enabled`, `any_ai_provider_enabled`), enforced in `Agent.instrument()`. + +| Variable | Enables | +|---|---| +| `HARNESS_ENABLE_API` | All non-AI instrumentation: HTTP servers/clients, gRPC, MySQL/PostgreSQL, botocore, and generic OTel contrib fallback | +| `HARNESS_ENABLE_AI_OPENAI` | OpenAI | +| `HARNESS_ENABLE_AI_ANTHROPIC` | Anthropic | +| `HARNESS_ENABLE_AI_LITELLM` | LiteLLM | +| `HARNESS_ENABLE_AI_GOOGLE_GENAI` | Google GenAI (Gemini / Vertex AI) | +| `HARNESS_ENABLE_AI_MCP` | Model Context Protocol | + +The legacy `HA_GEN_AI_ENABLED` master switch no longer controls instrumentation. `skip_libraries=[...]` on `instrument()` still takes precedence over enable flags. ## Plugin system diff --git a/README.md b/README.md index fcdb8f9..e45e46c 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ pip install -e ".[dev,anthropic,google-genai,openai,litellm]" ./scripts/run-unit-tests.sh ``` -Environment variables for SDK configuration use the `HA_` prefix (for example `HA_SERVICE_NAME`, `HA_REPORTING_ENDPOINT`). +Environment variables for SDK configuration use the `HARNESS_` prefix (for example `HARNESS_SERVICE_NAME`, `HARNESS_REPORTING_ENDPOINT`). The legacy `HA_`, `AT_`, and `TA_` prefixes remain supported for backwards compatibility; when the same setting is defined under multiple prefixes, `HARNESS_` wins, then `HA_`, then `AT_`, then `TA_`. ### Integration tests (MySQL / PostgreSQL) @@ -108,10 +108,42 @@ agent.instrument() Auto-instrumentation: ```bash -export HA_CONFIG_FILE=/path/to/config.yaml +export HARNESS_CONFIG_FILE=/path/to/config.yaml harness-instrument python app.py ``` +## Enabling instrumentation (opt-in) + +Instrumentation is **opt-in**. By default `agent.instrument()` and +`harness-instrument` instrument nothing; you must explicitly enable the +categories you want. A flag is "on" only when its value is (case-insensitively) +`true`; any other value — or an unset variable — leaves it off. + +| Category | Environment variable | Enables | +|----------|----------------------|---------| +| API / HTTP | `HARNESS_ENABLE_API` | All non-AI instrumentation: HTTP servers (Django, Flask, FastAPI), HTTP clients (requests, HTTPX, aiohttp), gRPC, databases (MySQL, PostgreSQL), botocore, and any installed OpenTelemetry contrib instrumentors (generic fallback) | +| AI: OpenAI | `HARNESS_ENABLE_AI_OPENAI` | OpenAI instrumentation | +| AI: Anthropic | `HARNESS_ENABLE_AI_ANTHROPIC` | Anthropic instrumentation | +| AI: LiteLLM | `HARNESS_ENABLE_AI_LITELLM` | LiteLLM instrumentation | +| AI: Google GenAI | `HARNESS_ENABLE_AI_GOOGLE_GENAI` | Google GenAI (Gemini / Vertex AI) instrumentation | +| AI: MCP | `HARNESS_ENABLE_AI_MCP` | Model Context Protocol instrumentation | + +Each AI provider is enabled independently. The `HARNESS_ENABLE_*` opt-in flags +are **only** read under the `HARNESS_` prefix — they have no `HA_`/`AT_`/`TA_` +aliases. The legacy `HA_GEN_AI_ENABLED` master switch no longer enables or +disables AI instrumentation; use the per-provider flags instead. + +```bash +# Enable HTTP/API + OpenAI only +export HARNESS_ENABLE_API=true +export HARNESS_ENABLE_AI_OPENAI=true +harness-instrument python app.py +``` + +`agent.instrument(skip_libraries=[...])` still works and takes precedence: a +library listed in `skip_libraries` is never instrumented even if its category +is enabled. + ## Plugins The SDK loads extensions via [setuptools entry points](https://setuptools.pypa.io/en/latest/userguide/entry_point.html). Each plugin has a **name** (the entry-point key). Names are listed in config or environment variables; only installed plugins are loaded, in the order you configure. diff --git a/docs/sdk-quickstart.md b/docs/sdk-quickstart.md new file mode 100644 index 0000000..7b91c8f --- /dev/null +++ b/docs/sdk-quickstart.md @@ -0,0 +1,153 @@ +# Send traces to Harness UDP ingest + +A short integration: install the SDK, opt into the instrumentation you want, +call `Agent().instrument()` once at startup, and point a few `HARNESS_*` +environment variables at Harness UDP ingest. Instrumentation is opt-in — you +choose HTTP/API instrumentation and each AI provider explicitly — and the SDK +exports OTLP/HTTP spans to ingest. + +> Environment variables use the `HARNESS_` prefix. The legacy `HA_`, `AT_`, and +> `TA_` prefixes still work for backwards compatibility (precedence: +> `HARNESS_` > `HA_` > `AT_` > `TA_`). + +## 1. Install + +The SDK is published on PyPI as `harness-sdk`. Pin the published release in +applications and CI builds: + +```bash +pip install "harness-sdk==1.0.1" +``` + +Add the extra for the LLM client you use: + +```bash +pip install "harness-sdk[litellm]==1.0.1" # for LiteLLM +pip install "harness-sdk[anthropic]==1.0.1" # for the Anthropic client +``` + +> Anthropic in-process spans additionally require the OTel GenAI helpers, which +> are not bundled: +> `pip install opentelemetry-instrumentation-anthropic opentelemetry-util-genai` + +## 2. Configure (environment) + +| Variable | Value | +|---|---| +| `HARNESS_SERVICE_NAME` | your service name | +| `HARNESS_REPORTING_ENDPOINT` | `https://app.harness.io/udp-ingest/otel/v1/traces?accountIdentifier=&routingId=` | +| `HARNESS_REPORTING_TRACE_REPORTER_TYPE` | `OTLP_HTTP` | +| `HARNESS_REPORTING_TOKEN` | a Harness **service account token** | + +```bash +export HARNESS_SERVICE_NAME="my-service" +export HARNESS_REPORTING_ENDPOINT="https://app.harness.io/udp-ingest/otel/v1/traces?accountIdentifier=&routingId=" +export HARNESS_REPORTING_TRACE_REPORTER_TYPE=OTLP_HTTP +export HARNESS_REPORTING_TOKEN="" +``` + +The token is sent to ingest as the `x-harness-service-token` header. + +## 3. Opt into instrumentation + +Instrumentation is opt-in: nothing is instrumented until you enable it. Set one +flag for all HTTP/API instrumentation, and one flag per AI provider you use. A +flag is on only when its value is `true`. + +| Variable | Enables | +|---|---| +| `HARNESS_ENABLE_API` | All non-AI instrumentation (HTTP servers/clients, gRPC, DB, botocore, generic OTel contrib) | +| `HARNESS_ENABLE_AI_OPENAI` | OpenAI | +| `HARNESS_ENABLE_AI_ANTHROPIC` | Anthropic | +| `HARNESS_ENABLE_AI_LITELLM` | LiteLLM | +| `HARNESS_ENABLE_AI_GOOGLE_GENAI` | Google GenAI (Gemini / Vertex AI) | +| `HARNESS_ENABLE_AI_MCP` | Model Context Protocol | + +```bash +# Example: HTTP/API instrumentation plus LiteLLM +export HARNESS_ENABLE_API=true +export HARNESS_ENABLE_AI_LITELLM=true +``` + +## 4. Instrument (two lines) + +Call this once, as early as possible in your process startup — before you make +LLM calls. + +```python +from harness_sdk.agent import Agent + +Agent().instrument() +``` + +That's the entire code change. Everything below is just your normal app code. + +--- + +## Example: LiteLLM + +```python +# Requires: export HARNESS_ENABLE_AI_LITELLM=true +from harness_sdk.agent import Agent +Agent().instrument() # do this first, at startup + +import litellm + +resp = litellm.completion( + model="anthropic/claude-3-5-sonnet-20241022", + messages=[{"role": "user", "content": "Reply with one short sentence."}], + max_tokens=64, +) +print(resp.choices[0].message.content) +``` + +Each `litellm.completion` / `acompletion` / `embedding` / `aembedding` call now +emits a `litellm_request` span with `gen_ai.*` attributes, exported to ingest. + +## Example: Anthropic Python client + +```python +# Requires: export HARNESS_ENABLE_AI_ANTHROPIC=true +from harness_sdk.agent import Agent +Agent().instrument() # do this first, at startup + +import anthropic + +client = anthropic.Anthropic() # reads ANTHROPIC_API_KEY +msg = client.messages.create( + model="claude-3-5-sonnet-20241022", + max_tokens=64, + messages=[{"role": "user", "content": "Reply with one short sentence."}], +) +print(msg.content[0].text) +``` + +`Messages.create` / `stream` (sync and async) are wrapped automatically and emit +GenAI spans to ingest. + +--- + +## Scope the instrumentation (optional) + +Because instrumentation is opt-in, you scope it by enabling only the flags you +need — for example, set `HARNESS_ENABLE_AI_LITELLM=true` and leave the other +`HARNESS_ENABLE_*` flags unset to instrument LiteLLM only. + +For finer control you can also exclude an otherwise-enabled library at the call +site; `skip_libraries` takes precedence over the enable flags: + +```python +# HTTP/API enabled, but never instrument requests. +# export HARNESS_ENABLE_API=true +Agent().instrument(skip_libraries=["requests"]) +``` + +## Useful toggles (optional) + +| Variable | Effect | +|---|---| +| `HARNESS_GEN_AI_PAYLOAD_CAPTURE_ENABLED` | `false` to omit prompt/response bodies from spans | +| `HARNESS_SPAN_ATTRIBUTES` | extra attributes on every span, e.g. `env=prod,team=ai` | +| `HARNESS_OBSERVABILITY_PLUGINS` | `builtin_span_attributes` to keep instrumentation but disable the SDK's own OTLP exporter (if you already export spans yourself) | +| `HARNESS_ENABLED` | `false` to disable the SDK entirely (no code change) | + diff --git a/src/harness_sdk/agent.py b/src/harness_sdk/agent.py index 3dcce58..0012be6 100644 --- a/src/harness_sdk/agent.py +++ b/src/harness_sdk/agent.py @@ -17,6 +17,9 @@ DJANGO_KEY, FAST_API_KEY, instrument_supported_contrib_without_wrapper, + is_library_enabled, + is_api_instrumentation_enabled, + any_ai_provider_enabled, ) from harness_sdk.instrumentation.genai_env import maybe_set_genai_payload_capture_env_vars from harness_sdk.plugins.control import ControlPlugin, get_control_registry @@ -85,16 +88,20 @@ def instrument(self, app=None, skip_libraries=None, auto_instrument=False): logger.debug('agent is not initialized, not instrumenting') return - maybe_set_genai_payload_capture_env_vars() + if any_ai_provider_enabled(): + maybe_set_genai_payload_capture_env_vars() for library_key in SUPPORTED_LIBRARIES: if library_key in skip_libraries: logger.debug('not attempting to instrument %s', library_key) continue + if not is_library_enabled(library_key): + logger.debug('%s not enabled via opt-in env flag, skipping', library_key) + continue logger.debug("attempting to instrument %s", library_key) self._instrument(library_key, app, auto_instrument) - if self._should_instrument_generic_contrib(): + if is_api_instrumentation_enabled() and self._should_instrument_generic_contrib(): instrument_supported_contrib_without_wrapper(skip_libraries) else: logger.debug("Skipping generic contrib fallback instrumentation") diff --git a/src/harness_sdk/agent_init.py b/src/harness_sdk/agent_init.py index 2d93797..8ff94b6 100644 --- a/src/harness_sdk/agent_init.py +++ b/src/harness_sdk/agent_init.py @@ -17,6 +17,7 @@ from opentelemetry.sdk.resources import Resource from harness_sdk import constants from harness_sdk.config import config_pb2 +from harness_sdk.env import is_env_var_present from harness_sdk.otlp_reporting import ( compression_type_to_otlp_grpc, compression_type_to_otlp_http, @@ -42,11 +43,7 @@ def apply_config(self, agent_config): self._config = agent_config self.init_trace_provider() self.init_propagation() - if ( - "HA_ENABLE_CONSOLE_SPAN_EXPORTER" in os.environ - or "AT_ENABLE_CONSOLE_SPAN_EXPORTER" in os.environ - or "TA_ENABLE_CONSOLE_SPAN_EXPORTER" in os.environ - ): + if is_env_var_present("ENABLE_CONSOLE_SPAN_EXPORTER"): self.set_console_span_processor() def init_trace_provider(self) -> None: diff --git a/src/harness_sdk/config/config.py b/src/harness_sdk/config/config.py index 4a00d9f..9ad2888 100644 --- a/src/harness_sdk/config/config.py +++ b/src/harness_sdk/config/config.py @@ -9,14 +9,16 @@ from harness_sdk.config import config_pb2 from harness_sdk.config.default import DEFAULT, SDK_CONFIG_KEYS from harness_sdk.config.environment import overwrite_with_environment +from harness_sdk.env import get_env_value from harness_sdk.otlp_reporting import normalize_reporting_dict from harness_sdk.custom_logger import get_custom_logger logger = get_custom_logger(__name__) -def _parse_plugin_env(env_key: str) -> list | None: - raw = os.environ.get(env_key) +def _parse_plugin_env(target_key: str) -> list | None: + """Parse a comma-separated plugin list, honoring HARNESS_/HA_/AT_/TA_ precedence.""" + raw = get_env_value(target_key) if not raw: return None return [name.strip() for name in raw.split(',') if name.strip()] @@ -94,9 +96,7 @@ def build_config(): plugins_config = config_dict.pop('plugins', {}) - enabled_control_plugins = _parse_plugin_env('HA_CONTROL_PLUGINS') - if enabled_control_plugins is None: - enabled_control_plugins = _parse_plugin_env('AT_CONTROL_PLUGINS') + enabled_control_plugins = _parse_plugin_env('CONTROL_PLUGINS') if enabled_control_plugins is None: enabled_control_plugins = _ordered_plugin_names_from_section( plugins_config.get('control') @@ -104,9 +104,7 @@ def build_config(): if not enabled_control_plugins: enabled_control_plugins = [] - enabled_observability_plugins = _parse_plugin_env('HA_OBSERVABILITY_PLUGINS') - if enabled_observability_plugins is None: - enabled_observability_plugins = _parse_plugin_env('AT_OBSERVABILITY_PLUGINS') + enabled_observability_plugins = _parse_plugin_env('OBSERVABILITY_PLUGINS') if enabled_observability_plugins is None: enabled_observability_plugins = _ordered_plugin_names_from_section( plugins_config.get('observability') @@ -149,8 +147,4 @@ def read_from_file(): def _config_file_path(): - return ( - os.environ.get('HA_CONFIG_FILE') - or os.environ.get('AT_CONFIG_FILE') - or os.environ.get('TA_CONFIG_FILE') - ) + return get_env_value('CONFIG_FILE') diff --git a/src/harness_sdk/config/environment.py b/src/harness_sdk/config/environment.py index eca3cf2..fd8dc37 100644 --- a/src/harness_sdk/config/environment.py +++ b/src/harness_sdk/config/environment.py @@ -1,15 +1,14 @@ -"""Apply HA_* environment variables to harness-sdk configuration.""" -import os - +"""Apply harness-sdk environment variables to configuration.""" from google.protobuf import wrappers_pb2 from harness_sdk.config import config_pb2 +from harness_sdk.env import get_env_value from harness_sdk.otlp_reporting import compression_type_to_enum def _env(name): - """Prefer HA_ prefix; accept legacy AT_/TA_ for SDK settings during migration.""" - return os.environ.get(f"HA_{name}") or os.environ.get(f"AT_{name}") or os.environ.get(f"TA_{name}") + """Prefer HARNESS_ prefix; accept legacy HA_/AT_/TA_ for SDK settings during migration.""" + return get_env_value(name) def overwrite_with_environment(config, reporting_encoding=None): # pylint: disable=R0912,R0914,R0915 diff --git a/src/harness_sdk/env.py b/src/harness_sdk/env.py index fb7f2c6..e84d53f 100644 --- a/src/harness_sdk/env.py +++ b/src/harness_sdk/env.py @@ -1,10 +1,28 @@ import os +# Prefix precedence for resolving SDK settings: new HARNESS_ first, then the +# legacy HA_/AT_/TA_ aliases for backwards compatibility. +_PREFIXES = ("HARNESS_", "HA_", "AT_", "TA_") + def get_env_value(target_key): - """Read SDK env vars: HA_, then legacy AT_/TA_ aliases.""" - for prefix in ("HA_", "AT_", "TA_"): + """Read SDK env vars honoring prefix precedence: HARNESS_, then legacy HA_/AT_/TA_.""" + for prefix in _PREFIXES: env_var_key = f"{prefix}{target_key}" if env_var_key in os.environ: return os.environ[env_var_key] return None + + +def is_env_var_present(target_key): + """Return True if the key is set under any supported prefix (presence check).""" + return any(f"{prefix}{target_key}" in os.environ for prefix in _PREFIXES) + + +def is_harness_flag_enabled(env_var_name): + """Strict opt-in flag under the HARNESS_ prefix only (no legacy aliases). + + Returns True only when the value is present and case-insensitively 'true'. + """ + value = os.environ.get(env_var_name) + return value is not None and value.strip().lower() == "true" diff --git a/src/harness_sdk/instrumentation/anthropic/__init__.py b/src/harness_sdk/instrumentation/anthropic/__init__.py index 8afcbd0..9f9e274 100644 --- a/src/harness_sdk/instrumentation/anthropic/__init__.py +++ b/src/harness_sdk/instrumentation/anthropic/__init__.py @@ -137,10 +137,6 @@ def _wrapper( args: tuple[Any, ...], kwargs: dict[str, Any], ) -> Any: - if not Config().config.gen_ai.enabled.value: - logger.debug("Anthropic Messages.create (sync): gen_ai disabled, passthrough") - return wrapped(*args, **kwargs) - is_streaming = kwargs.get("stream") is True logger.debug("Anthropic Messages.create (sync): streaming=%s", is_streaming) params = extract_params(**kwargs) @@ -181,10 +177,6 @@ async def _wrapper( args: tuple[Any, ...], kwargs: dict[str, Any], ) -> Any: - if not Config().config.gen_ai.enabled.value: - logger.debug("Anthropic Messages.create (async): gen_ai disabled, passthrough") - return await wrapped(*args, **kwargs) - is_streaming = kwargs.get("stream") is True logger.debug("Anthropic Messages.create (async): streaming=%s", is_streaming) params = extract_params(**kwargs) @@ -277,10 +269,6 @@ def _wrapper( args: tuple[Any, ...], kwargs: dict[str, Any], ) -> Any: - if not Config().config.gen_ai.enabled.value: - logger.debug("Anthropic Messages.stream (sync): gen_ai disabled, passthrough") - return wrapped(*args, **kwargs) - merged_kwargs = dict(kwargs) merged_kwargs["stream"] = True params = extract_params(**merged_kwargs) @@ -314,10 +302,6 @@ def _wrapper( args: tuple[Any, ...], kwargs: dict[str, Any], ) -> Any: - if not Config().config.gen_ai.enabled.value: - logger.debug("Anthropic Messages.stream (async): gen_ai disabled, passthrough") - return wrapped(*args, **kwargs) - merged_kwargs = dict(kwargs) merged_kwargs["stream"] = True params = extract_params(**merged_kwargs) @@ -353,9 +337,6 @@ def instrument(self, **_kwargs: Any) -> None: if self._applied: logger.debug("Anthropic instrumentation already applied.") return - if not Config().config.gen_ai.enabled.value: - logger.debug("Gen AI instrumentation disabled; skip Anthropic wraps.") - return try: handler = _get_handler() wrapt.wrap_function_wrapper( diff --git a/src/harness_sdk/instrumentation/google_genai/__init__.py b/src/harness_sdk/instrumentation/google_genai/__init__.py index 4bf29a4..6e02cd8 100644 --- a/src/harness_sdk/instrumentation/google_genai/__init__.py +++ b/src/harness_sdk/instrumentation/google_genai/__init__.py @@ -312,8 +312,6 @@ def _make_generate_sync(handler: TelemetryHandler) -> Callable[..., Any]: capture_content = should_capture_content_on_spans_in_experimental_mode() def _wrapper(wrapped, instance, args, kwargs): - if not Config().config.gen_ai.enabled.value: - return wrapped(*args, **kwargs) invocation = _build_invocation( handler, _resolve_provider(instance), _extract_model(args, kwargs), _extract_contents(args, kwargs), @@ -342,8 +340,6 @@ def _make_generate_async(handler: TelemetryHandler) -> Callable[..., Any]: capture_content = should_capture_content_on_spans_in_experimental_mode() async def _wrapper(wrapped, instance, args, kwargs): - if not Config().config.gen_ai.enabled.value: - return await wrapped(*args, **kwargs) invocation = _build_invocation( handler, _resolve_provider(instance), _extract_model(args, kwargs), _extract_contents(args, kwargs), @@ -417,8 +413,6 @@ def _make_generate_stream_sync(handler: TelemetryHandler) -> Callable[..., Any]: capture_content = should_capture_content_on_spans_in_experimental_mode() def _wrapper(wrapped, instance, args, kwargs): - if not Config().config.gen_ai.enabled.value: - return wrapped(*args, **kwargs) invocation = _build_invocation( handler, _resolve_provider(instance), _extract_model(args, kwargs), _extract_contents(args, kwargs), @@ -443,8 +437,6 @@ def _make_generate_stream_async(handler: TelemetryHandler) -> Callable[..., Any] capture_content = should_capture_content_on_spans_in_experimental_mode() async def _wrapper(wrapped, instance, args, kwargs): - if not Config().config.gen_ai.enabled.value: - return await wrapped(*args, **kwargs) invocation = _build_invocation( handler, _resolve_provider(instance), _extract_model(args, kwargs), _extract_contents(args, kwargs), @@ -483,8 +475,6 @@ def _apply_embedding_response(invocation: Any, result: Any) -> None: def _make_embed_sync(handler: TelemetryHandler) -> Callable[..., Any]: def _wrapper(wrapped, instance, args, kwargs): - if not Config().config.gen_ai.enabled.value: - return wrapped(*args, **kwargs) invocation = handler.start_embedding( _resolve_provider(instance), request_model=_extract_model(args, kwargs), @@ -510,8 +500,6 @@ def _wrapper(wrapped, instance, args, kwargs): def _make_embed_async(handler: TelemetryHandler) -> Callable[..., Any]: async def _wrapper(wrapped, instance, args, kwargs): - if not Config().config.gen_ai.enabled.value: - return await wrapped(*args, **kwargs) invocation = handler.start_embedding( _resolve_provider(instance), request_model=_extract_model(args, kwargs), @@ -560,9 +548,6 @@ def instrument(self, **_kwargs: Any) -> None: if self._applied: logger.debug("google-genai instrumentation already applied.") return - if not Config().config.gen_ai.enabled.value: - logger.debug("Gen AI instrumentation disabled; skip google-genai wraps.") - return try: handler = _get_handler() for cls_name, method, factory in _WRAPPED_METHODS: diff --git a/src/harness_sdk/instrumentation/instrumentation_definitions.py b/src/harness_sdk/instrumentation/instrumentation_definitions.py index 91cf359..47e6f9a 100644 --- a/src/harness_sdk/instrumentation/instrumentation_definitions.py +++ b/src/harness_sdk/instrumentation/instrumentation_definitions.py @@ -3,6 +3,7 @@ from importlib import metadata as importlib_metadata from harness_sdk.custom_logger import get_custom_logger +from harness_sdk.env import is_harness_flag_enabled FLASK_KEY = 'flask' DJANGO_KEY = 'django' @@ -34,6 +35,46 @@ MCP_KEY, ] +# Instrumentation is strictly opt-in. Everything that is not AI (HTTP servers +# and clients, RPC, databases, AWS) is gated behind a single API toggle. +API_ENABLE_ENV = "HARNESS_ENABLE_API" + +API_LIBRARIES = frozenset({ + FLASK_KEY, DJANGO_KEY, FAST_API_KEY, + GRPC_SERVER_KEY, GRPC_CLIENT_KEY, + POSTGRESQL_KEY, MYSQL_KEY, + REQUESTS_KEY, HTTPX_KEY, AIOHTTP_CLIENT_KEY, + BOTOCORE, +}) + +# Each AI provider is enabled independently via its own HARNESS_ENABLE_AI_* flag. +AI_LIBRARY_ENV_FLAGS = { + OPENAI_KEY: "HARNESS_ENABLE_AI_OPENAI", + ANTHROPIC_KEY: "HARNESS_ENABLE_AI_ANTHROPIC", + LITELLM_KEY: "HARNESS_ENABLE_AI_LITELLM", + GOOGLE_GENAI_KEY: "HARNESS_ENABLE_AI_GOOGLE_GENAI", + MCP_KEY: "HARNESS_ENABLE_AI_MCP", +} + + +def is_api_instrumentation_enabled(): + """True when the single API/HTTP instrumentation toggle is explicitly on.""" + return is_harness_flag_enabled(API_ENABLE_ENV) + + +def any_ai_provider_enabled(): + """True when at least one AI provider is explicitly opted in.""" + return any(is_harness_flag_enabled(flag) for flag in AI_LIBRARY_ENV_FLAGS.values()) + + +def is_library_enabled(library_key): + """Decide whether a supported library should be instrumented based on opt-in env flags.""" + if library_key in AI_LIBRARY_ENV_FLAGS: + return is_harness_flag_enabled(AI_LIBRARY_ENV_FLAGS[library_key]) + if library_key in API_LIBRARIES: + return is_api_instrumentation_enabled() + return False + # map of library_key => instrumentation wrapper instance _INSTRUMENTATION_STATE = {} _GENERIC_INSTRUMENTATION_STATE = {} diff --git a/src/harness_sdk/instrumentation/litellm/__init__.py b/src/harness_sdk/instrumentation/litellm/__init__.py index 13e24fb..eaf7fae 100644 --- a/src/harness_sdk/instrumentation/litellm/__init__.py +++ b/src/harness_sdk/instrumentation/litellm/__init__.py @@ -25,7 +25,6 @@ import contextvars import json -import os from dataclasses import dataclass from typing import Any, Callable, Optional @@ -37,6 +36,7 @@ from harness_sdk.config.config import Config from harness_sdk.custom_logger import get_custom_logger +from harness_sdk.env import get_env_value from harness_sdk.plugins.control import get_control_registry from harness_sdk.gen_ai.exceptions import ControlEvaluationBlocked from harness_sdk.instrumentation import BaseInstrumentorWrapper @@ -50,10 +50,11 @@ "harness_litellm_span_active", default=False ) -# Env var to opt into bounded raw response/usage capture as a resilience fallback -# when LiteLLM changes response object shape or field names. -_RAW_CAPTURE_ENV = "HA_GEN_AI_RAW_CAPTURE_ENABLED" -_RAW_CAPTURE_MAX_BYTES_ENV = "HA_GEN_AI_RAW_CAPTURE_MAX_BYTES" +# Env keys (resolved with HARNESS_/HA_/AT_/TA_ precedence) to opt into bounded raw +# response/usage capture as a resilience fallback when LiteLLM changes response +# object shape or field names. +_RAW_CAPTURE_ENV = "GEN_AI_RAW_CAPTURE_ENABLED" +_RAW_CAPTURE_MAX_BYTES_ENV = "GEN_AI_RAW_CAPTURE_MAX_BYTES" _RAW_CAPTURE_DEFAULT_MAX_BYTES = 8192 # Fields that must never be serialized into the raw usage/response fallback. @@ -326,11 +327,11 @@ def _is_bedrock_model_id_header(key: Any) -> bool: def _raw_capture_enabled() -> bool: - return os.getenv(_RAW_CAPTURE_ENV, "").strip().lower() == "true" + return (get_env_value(_RAW_CAPTURE_ENV) or "").strip().lower() == "true" def _raw_capture_max_bytes() -> int: - raw = os.getenv(_RAW_CAPTURE_MAX_BYTES_ENV, "").strip() + raw = (get_env_value(_RAW_CAPTURE_MAX_BYTES_ENV) or "").strip() if not raw: return _RAW_CAPTURE_DEFAULT_MAX_BYTES try: @@ -787,9 +788,6 @@ def _sync_wrapper( args: tuple[Any, ...], kwargs: dict[str, Any], ) -> Any: - if not Config().config.gen_ai.enabled.value: - return wrapped(*args, **kwargs) - # Skip nested re-dispatch (e.g. aembedding -> embedding) so each provider # call yields exactly one litellm_request span. if _LITELLM_SPAN_ACTIVE.get(): @@ -808,9 +806,6 @@ async def _async_wrapper( args: tuple[Any, ...], kwargs: dict[str, Any], ) -> Any: - if not Config().config.gen_ai.enabled.value: - return await wrapped(*args, **kwargs) - if _LITELLM_SPAN_ACTIVE.get(): return await wrapped(*args, **kwargs) @@ -836,9 +831,6 @@ def instrument(self, **_kwargs: Any) -> None: if self._applied: logger.debug("LiteLLM instrumentation already applied.") return - if not Config().config.gen_ai.enabled.value: - logger.debug("Gen AI instrumentation disabled; skip LiteLLM wraps.") - return try: import litellm # pylint: disable=import-outside-toplevel diff --git a/src/harness_sdk/instrumentation/mcp/__init__.py b/src/harness_sdk/instrumentation/mcp/__init__.py index c3a1a71..5d34352 100644 --- a/src/harness_sdk/instrumentation/mcp/__init__.py +++ b/src/harness_sdk/instrumentation/mcp/__init__.py @@ -1,7 +1,6 @@ '''Traceable wrapper around OpenTelemetry MCP instrumentation (Model Context Protocol).''' from opentelemetry.instrumentation.mcp import McpInstrumentor -from harness_sdk.config.config import Config from harness_sdk.custom_logger import get_custom_logger from harness_sdk.instrumentation import BaseInstrumentorWrapper from harness_sdk.instrumentation.mcp.gen_ai_mirror import ( @@ -16,10 +15,10 @@ class McpInstrumentorWrapper(McpInstrumentor, BaseInstrumentorWrapper): """ Instruments the MCP Python SDK with GenAI-aware configuration. - Uses the same ``gen_ai`` block as the rest of the agent (``TA_GEN_AI_ENABLED``, - ``TA_GEN_AI_PAYLOAD_CAPTURE_ENABLED``, ``TA_GEN_AI_PAYLOAD_EVALUATION_ENABLED``): - when GenAI is disabled, MCP instrumentation is not applied; payload capture mirrors - ``TRACELOOP_TRACE_CONTENT`` for FastMCP; tool spans additionally record + MCP instrumentation is opt-in via ``HARNESS_ENABLE_AI_MCP`` (enforced by the + Agent before this wrapper is loaded). Payload capture mirrors + ``TRACELOOP_TRACE_CONTENT`` for FastMCP based on + ``HARNESS_GEN_AI_PAYLOAD_CAPTURE_ENABLED``; tool spans additionally record ``gen_ai.operation.name``, ``gen_ai.system``, ``gen_ai.tool.name``, ``mcp.method.name``, and opt-in ``gen_ai.tool.call.arguments`` / ``gen_ai.tool.call.result``. """ @@ -29,11 +28,6 @@ def __init__(self): BaseInstrumentorWrapper.__init__(self) def _instrument(self, **kwargs) -> None: - gen = Config().config.gen_ai - if not gen.enabled.value: - logger.debug("gen_ai disabled in config; skipping MCP instrumentation") - return - apply_gen_ai_env_for_mcp() with mcp_instrumentation_get_tracer_patched(): diff --git a/src/harness_sdk/instrumentation/mcp/gen_ai_mirror.py b/src/harness_sdk/instrumentation/mcp/gen_ai_mirror.py index fc8b6b9..85b6bc3 100644 --- a/src/harness_sdk/instrumentation/mcp/gen_ai_mirror.py +++ b/src/harness_sdk/instrumentation/mcp/gen_ai_mirror.py @@ -99,9 +99,6 @@ def _mirror_entity_output(span, key: str, value: Any, current_kind: Optional[str def mirror_traceloop_to_gen_ai(span, key: str, value: Any, current_kind: Optional[str]) -> None: """Copy select Traceloop attributes to gen_ai.* / mcp.* on the same recording span.""" gen = Config().config.gen_ai - if not gen.enabled.value: - return - tool_kind = TraceloopSpanKindValues.TOOL.value _mirror_span_kind(span, key, value, tool_kind) _mirror_entity_name(span, key, value, tool_kind, current_kind) diff --git a/src/harness_sdk/instrumentation/openai/__init__.py b/src/harness_sdk/instrumentation/openai/__init__.py index 06a50a3..7d92cd5 100644 --- a/src/harness_sdk/instrumentation/openai/__init__.py +++ b/src/harness_sdk/instrumentation/openai/__init__.py @@ -88,10 +88,6 @@ def _wrapper( args: tuple[Any, ...], kwargs: dict[str, Any], ) -> Any: - if not Config().config.gen_ai.enabled.value: - logger.debug("OpenAI Completions.create (sync): gen_ai disabled, passthrough") - return wrapped(*args, **kwargs) - invocation = handler.start_llm( create_chat_invocation(kwargs, instance, capture_content=capture_content) ) @@ -130,10 +126,6 @@ async def _wrapper( args: tuple[Any, ...], kwargs: dict[str, Any], ) -> Any: - if not Config().config.gen_ai.enabled.value: - logger.debug("OpenAI AsyncCompletions.create (async): gen_ai disabled, passthrough") - return await wrapped(*args, **kwargs) - invocation = handler.start_llm( create_chat_invocation(kwargs, instance, capture_content=capture_content) ) @@ -170,9 +162,6 @@ def _wrapper( args: tuple[Any, ...], kwargs: dict[str, Any], ) -> Any: - if not Config().config.gen_ai.enabled.value: - return wrapped(*args, **kwargs) - address, port = get_server_address_and_port(instance) invocation = handler.start_embedding( "openai", @@ -219,9 +208,6 @@ async def _wrapper( args: tuple[Any, ...], kwargs: dict[str, Any], ) -> Any: - if not Config().config.gen_ai.enabled.value: - return await wrapped(*args, **kwargs) - address, port = get_server_address_and_port(instance) invocation = handler.start_embedding( "openai", @@ -273,9 +259,6 @@ def instrument(self, **_kwargs: Any) -> None: if self._applied: logger.debug("OpenAI instrumentation already applied.") return - if not Config().config.gen_ai.enabled.value: - logger.debug("Gen AI instrumentation disabled; skip OpenAI wraps.") - return try: handler = _get_handler() latest_experimental = is_experimental_mode() diff --git a/src/harness_sdk/plugins/builtin/pipeline.py b/src/harness_sdk/plugins/builtin/pipeline.py index 46a3f45..a59522d 100644 --- a/src/harness_sdk/plugins/builtin/pipeline.py +++ b/src/harness_sdk/plugins/builtin/pipeline.py @@ -1,5 +1,4 @@ """Default OTLP export pipeline with db-filter and attribute processors.""" -import os from typing import Any, List from opentelemetry.sdk.trace import SpanProcessor @@ -7,6 +6,7 @@ from harness_sdk.agent_init import AgentInit from harness_sdk.custom_logger import get_custom_logger +from harness_sdk.env import is_env_var_present from harness_sdk.excluded_by_attribute_span_processor import ExcludeByAttributeSpanProcessor from harness_sdk.db_control_span_processor import DbControlSpanProcessor @@ -24,11 +24,7 @@ def on_init(self, config: Any) -> None: self._agent_init = AgentInit(config) def create_span_processors(self, config: Any) -> List[SpanProcessor]: - if ( - "HA_ENABLE_CONSOLE_SPAN_EXPORTER" in os.environ - or "AT_ENABLE_CONSOLE_SPAN_EXPORTER" in os.environ - or "TA_ENABLE_CONSOLE_SPAN_EXPORTER" in os.environ - ): + if is_env_var_present("ENABLE_CONSOLE_SPAN_EXPORTER"): self._agent_init.set_console_span_processor() return [] diff --git a/test/config/environment_test.py b/test/config/environment_test.py index 5fb6651..b2f717b 100644 --- a/test/config/environment_test.py +++ b/test/config/environment_test.py @@ -57,7 +57,38 @@ def test_gen_ai_env_overrides() -> None: Config._instance = None +def test_harness_prefix_takes_precedence_over_legacy() -> None: + Config._instance = None + os.environ["HARNESS_SERVICE_NAME"] = "harness_service" + os.environ["HA_SERVICE_NAME"] = "legacy_ha_service" + os.environ["AT_SERVICE_NAME"] = "legacy_at_service" + os.environ["TA_SERVICE_NAME"] = "legacy_ta_service" + config = Config().config + assert config.service_name.value == "harness_service" + unset_env_variables() + Config._instance = None + + +def test_legacy_ha_prefix_applies_when_harness_absent() -> None: + Config._instance = None + os.environ["HA_SERVICE_NAME"] = "legacy_ha_service" + config = Config().config + assert config.service_name.value == "legacy_ha_service" + unset_env_variables() + Config._instance = None + + +def test_legacy_at_prefix_applies_when_ha_and_harness_absent() -> None: + Config._instance = None + os.environ["AT_SERVICE_NAME"] = "legacy_at_service" + config = Config().config + assert config.service_name.value == "legacy_at_service" + unset_env_variables() + Config._instance = None + + def unset_env_variables(): - keys_to_delete = [key for key in os.environ if key.startswith("HA_")] + prefixes = ("HARNESS_", "HA_", "AT_", "TA_") + keys_to_delete = [key for key in os.environ if key.startswith(prefixes)] for key in keys_to_delete: del os.environ[key] diff --git a/test/conftest.py b/test/conftest.py index 5deb7cd..993a3f3 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -54,7 +54,8 @@ def reset_singletons(): del os.environ["_HANDLER"] keys_to_delete = [ key for key in os.environ - if key.startswith("HA_") and key != "HA_ENABLE_CONSOLE_SPAN_EXPORTER" + if (key.startswith("HA_") or key.startswith("HARNESS_")) + and key != "HA_ENABLE_CONSOLE_SPAN_EXPORTER" ] for key in keys_to_delete: del os.environ[key] @@ -74,10 +75,17 @@ def pytest_sessionfinish(session, exitstatus): # pylint: disable=unused-argumen @pytest.fixture def agent(): os.environ['HA_ENABLE_CONSOLE_SPAN_EXPORTER'] = 'true' - os.environ['HA_GEN_AI_ENABLED'] = 'true' os.environ['HA_GEN_AI_PAYLOAD_CAPTURE_ENABLED'] = 'true' os.environ['HA_GEN_AI_PAYLOAD_EVALUATION_ENABLED'] = 'true' os.environ['HA_CONTROL_PLUGINS'] = '' + # Instrumentation is opt-in; enable every category so shared fixtures cover + # both API and AI instrumentation suites. + os.environ['HARNESS_ENABLE_API'] = 'true' + os.environ['HARNESS_ENABLE_AI_OPENAI'] = 'true' + os.environ['HARNESS_ENABLE_AI_ANTHROPIC'] = 'true' + os.environ['HARNESS_ENABLE_AI_LITELLM'] = 'true' + os.environ['HARNESS_ENABLE_AI_GOOGLE_GENAI'] = 'true' + os.environ['HARNESS_ENABLE_AI_MCP'] = 'true' _uninstrument_all() maybe_set_genai_payload_capture_env_vars() diff --git a/test/instrumentation/anthropic/test_anthropic.py b/test/instrumentation/anthropic/test_anthropic.py index 0a94767..cc604b8 100644 --- a/test/instrumentation/anthropic/test_anthropic.py +++ b/test/instrumentation/anthropic/test_anthropic.py @@ -481,7 +481,7 @@ def counting_fake(_self, *_a, **_k): # ── gen_ai disabled passthrough ───────────────────────────────────────────── -def test_gen_ai_disabled_passthrough(agent, exporter, anthropic_instrumentor): +def test_legacy_gen_ai_master_ignored(agent, exporter, anthropic_instrumentor): call_count = {"n": 0} def counting_fake(_self, *_a, **_k): @@ -489,6 +489,7 @@ def counting_fake(_self, *_a, **_k): return _FakeMessage() with patch.object(Messages, "create", new=counting_fake): + # Legacy master env must not disable instrumentation once opted in. os.environ["HA_GEN_AI_ENABLED"] = "false" from harness_sdk.config.config import Config Config._instance = None # force re-read of env @@ -503,7 +504,7 @@ def counting_fake(_self, *_a, **_k): spans = exporter.get_finished_spans() exporter.clear() - assert len(spans) == 0 + assert len(spans) == 1 assert call_count["n"] == 1 # reached the real (fake) implementation diff --git a/test/instrumentation/botocore/test_botocore.py b/test/instrumentation/botocore/test_botocore.py index b4f8bf9..f2e843e 100644 --- a/test/instrumentation/botocore/test_botocore.py +++ b/test/instrumentation/botocore/test_botocore.py @@ -1,6 +1,7 @@ import io import json import logging +import os import sys import traceback import zipfile @@ -20,6 +21,8 @@ def test_run(): logger = setup_custom_logger(__name__) + # API/HTTP instrumentation (incl. botocore) is opt-in. + os.environ["HARNESS_ENABLE_API"] = "true" with mock_iam(), mock_lambda(): agent = Agent() agent.instrument(None) diff --git a/test/instrumentation/litellm/litellm_instrumentation_test.py b/test/instrumentation/litellm/litellm_instrumentation_test.py index 800d19a..8b1b490 100644 --- a/test/instrumentation/litellm/litellm_instrumentation_test.py +++ b/test/instrumentation/litellm/litellm_instrumentation_test.py @@ -384,9 +384,10 @@ def test_litellm_raw_usage_capture_disabled_by_default(agent, exporter, litellm_ assert "gen_ai.response.usage.raw" not in attrs -def test_litellm_gen_ai_disabled_passthrough(agent, exporter, litellm_instrumentor): # pylint: disable=unused-argument +def test_litellm_legacy_gen_ai_master_ignored(agent, exporter, litellm_instrumentor): # pylint: disable=unused-argument import os + # Legacy master env must not disable instrumentation once opted in. os.environ["HA_GEN_AI_ENABLED"] = "false" from harness_sdk.config.config import Config @@ -410,7 +411,7 @@ def counting_fake(*_a, **_k): assert calls["n"] == 1 spans = exporter.get_finished_spans() exporter.clear() - assert len(spans) == 0 + assert len(spans) == 1 def test_litellm_mock_response_with_wrapper_enrichment(agent, exporter, litellm_instrumentor): # pylint: disable=unused-argument diff --git a/test/instrumentation/mcp/test_gen_ai_mirror.py b/test/instrumentation/mcp/test_gen_ai_mirror.py index 08b43bd..fcc7d34 100644 --- a/test/instrumentation/mcp/test_gen_ai_mirror.py +++ b/test/instrumentation/mcp/test_gen_ai_mirror.py @@ -48,15 +48,6 @@ def test_mirror_tool_name_and_mcp_method(mock_gen_ai_config): span.set_attribute.assert_any_call(AiSpanAttributes.MCP_METHOD_NAME, "tools/call") -def test_mirror_skips_when_gen_ai_disabled(mock_gen_ai_config): - mock_gen_ai_config.enabled.value = False - span = MagicMock() - mirror_traceloop_to_gen_ai( - span, AiSpanAttributes.TRACELOOP_SPAN_KIND, TraceloopSpanKindValues.TOOL.value, None - ) - span.set_attribute.assert_not_called() - - def test_mirror_arguments_when_capture_and_eval(mock_gen_ai_config): span = MagicMock() kind = TraceloopSpanKindValues.TOOL.value diff --git a/test/instrumentation/mcp/test_mcp_wrapper.py b/test/instrumentation/mcp/test_mcp_wrapper.py index 76271e1..fc98ceb 100644 --- a/test/instrumentation/mcp/test_mcp_wrapper.py +++ b/test/instrumentation/mcp/test_mcp_wrapper.py @@ -1,5 +1,5 @@ """Tests for :class:`harness_sdk.instrumentation.mcp.McpInstrumentorWrapper`.""" -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest from opentelemetry.instrumentation.mcp import McpInstrumentor @@ -9,30 +9,9 @@ @pytest.fixture def gen_ai_enabled(): - gen = MagicMock() - gen.enabled.value = True - gen.payload_capture_enabled.value = True - gen.payload_evaluation_enabled.value = True - root = MagicMock() - root.config.gen_ai = gen - with patch("harness_sdk.instrumentation.mcp.Config") as mock_cfg: - mock_cfg.return_value = root - yield gen - - -def test_mcp_wrapper_skips_when_gen_ai_disabled(): - gen = MagicMock() - gen.enabled.value = False - root = MagicMock() - root.config.gen_ai = gen - with patch("harness_sdk.instrumentation.mcp.Config") as mock_cfg: - mock_cfg.return_value = root - with patch("harness_sdk.instrumentation.mcp.apply_gen_ai_env_for_mcp") as mock_apply: - with patch.object(McpInstrumentor, "_instrument") as parent_instr: - w = McpInstrumentorWrapper() - w._instrument() - mock_apply.assert_not_called() - parent_instr.assert_not_called() + # MCP instrumentation is now gated by HARNESS_ENABLE_AI_MCP at the Agent level; + # the wrapper itself no longer consults Config for an enable flag. + yield def test_mcp_wrapper_restores_get_tracer_after_instrument(gen_ai_enabled): diff --git a/test/instrumentation/openai/openai_instrumentation_test.py b/test/instrumentation/openai/openai_instrumentation_test.py index b52a3fd..6d7047c 100644 --- a/test/instrumentation/openai/openai_instrumentation_test.py +++ b/test/instrumentation/openai/openai_instrumentation_test.py @@ -317,8 +317,9 @@ def test_double_instrument_is_noop(agent, exporter, openai_instrumentor): # pyl assert len(spans) == 1 -def test_gen_ai_disabled_passthrough_chat(agent, exporter, openai_instrumentor): # pylint: disable=unused-argument +def test_legacy_gen_ai_master_ignored_chat(agent, exporter, openai_instrumentor): # pylint: disable=unused-argument import os + # Legacy master env must not disable instrumentation once the provider is opted in. os.environ["HA_GEN_AI_ENABLED"] = "false" from harness_sdk.config.config import Config Config._instance = None @@ -337,11 +338,11 @@ def counting_fake(_self, *_a, **_k): assert calls["n"] == 1 spans = exporter.get_finished_spans() exporter.clear() - assert len(spans) == 0 + assert len(spans) == 1 @pytest.mark.asyncio -async def test_gen_ai_disabled_passthrough_async_chat(agent, exporter, openai_instrumentor): # pylint: disable=unused-argument +async def test_legacy_gen_ai_master_ignored_async_chat(agent, exporter, openai_instrumentor): # pylint: disable=unused-argument import os os.environ["HA_GEN_AI_ENABLED"] = "false" from harness_sdk.config.config import Config @@ -361,10 +362,10 @@ async def counting_fake(_self, *_a, **_k): assert calls["n"] == 1 spans = exporter.get_finished_spans() exporter.clear() - assert len(spans) == 0 + assert len(spans) == 1 -def test_gen_ai_disabled_passthrough_embeddings(agent, exporter, openai_instrumentor): # pylint: disable=unused-argument +def test_legacy_gen_ai_master_ignored_embeddings(agent, exporter, openai_instrumentor): # pylint: disable=unused-argument import os os.environ["HA_GEN_AI_ENABLED"] = "false" from harness_sdk.config.config import Config @@ -384,11 +385,11 @@ def counting_fake(_self, *_a, **_k): assert calls["n"] == 1 spans = exporter.get_finished_spans() exporter.clear() - assert len(spans) == 0 + assert len(spans) == 1 @pytest.mark.asyncio -async def test_gen_ai_disabled_passthrough_async_embeddings(agent, exporter, openai_instrumentor): # pylint: disable=unused-argument +async def test_legacy_gen_ai_master_ignored_async_embeddings(agent, exporter, openai_instrumentor): # pylint: disable=unused-argument import os os.environ["HA_GEN_AI_ENABLED"] = "false" from harness_sdk.config.config import Config @@ -408,7 +409,7 @@ async def counting_fake(_self, *_a, **_k): assert calls["n"] == 1 spans = exporter.get_finished_spans() exporter.clear() - assert len(spans) == 0 + assert len(spans) == 1 def test_completions_exception_records_error(agent, exporter, openai_instrumentor): # pylint: disable=unused-argument diff --git a/test/instrumentation/test_opt_in_gating.py b/test/instrumentation/test_opt_in_gating.py new file mode 100644 index 0000000..19c537d --- /dev/null +++ b/test/instrumentation/test_opt_in_gating.py @@ -0,0 +1,161 @@ +"""Opt-in instrumentation gating: default-off, one API toggle, per-provider AI toggles.""" +import os + +import pytest + +from harness_sdk.agent import Agent +from harness_sdk.instrumentation.genai_env import maybe_set_genai_payload_capture_env_vars +from harness_sdk.instrumentation.instrumentation_definitions import ( + _uninstrument_all, + is_api_instrumentation_enabled, + any_ai_provider_enabled, + is_library_enabled, + API_LIBRARIES, + AI_LIBRARY_ENV_FLAGS, + OPENAI_KEY, + REQUESTS_KEY, +) + + +# --------------------------------------------------------------------------- # +# Pure categorization helpers +# --------------------------------------------------------------------------- # +def test_default_off_nothing_enabled(): + assert is_api_instrumentation_enabled() is False + assert any_ai_provider_enabled() is False + for key in list(API_LIBRARIES) + list(AI_LIBRARY_ENV_FLAGS): + assert is_library_enabled(key) is False + + +def test_api_flag_enables_all_non_ai_only(): + os.environ["HARNESS_ENABLE_API"] = "true" + assert is_api_instrumentation_enabled() is True + for key in API_LIBRARIES: + assert is_library_enabled(key) is True + for key in AI_LIBRARY_ENV_FLAGS: + assert is_library_enabled(key) is False + assert any_ai_provider_enabled() is False + + +@pytest.mark.parametrize("provider_key,flag", list(AI_LIBRARY_ENV_FLAGS.items())) +def test_ai_provider_enabled_independently(provider_key, flag): + os.environ[flag] = "true" + assert is_library_enabled(provider_key) is True + assert any_ai_provider_enabled() is True + for other_key, other_flag in AI_LIBRARY_ENV_FLAGS.items(): + if other_key != provider_key: + assert is_library_enabled(other_key) is False + assert is_api_instrumentation_enabled() is False + for key in API_LIBRARIES: + assert is_library_enabled(key) is False + + +def test_flag_requires_exact_true_value(): + os.environ["HARNESS_ENABLE_API"] = "1" + assert is_api_instrumentation_enabled() is False + os.environ["HARNESS_ENABLE_API"] = "yes" + assert is_api_instrumentation_enabled() is False + os.environ["HARNESS_ENABLE_API"] = "false" + assert is_api_instrumentation_enabled() is False + os.environ["HARNESS_ENABLE_API"] = "TRUE" + assert is_api_instrumentation_enabled() is True + os.environ["HARNESS_ENABLE_API"] = " true " + assert is_api_instrumentation_enabled() is True + + +def test_ai_flags_have_no_legacy_aliases(): + os.environ["HA_ENABLE_AI_OPENAI"] = "true" + os.environ["AT_ENABLE_AI_OPENAI"] = "true" + os.environ["TA_ENABLE_AI_OPENAI"] = "true" + try: + assert is_library_enabled(OPENAI_KEY) is False + finally: + for legacy in ("AT_ENABLE_AI_OPENAI", "TA_ENABLE_AI_OPENAI"): + os.environ.pop(legacy, None) + + +def test_api_flag_has_no_legacy_aliases(): + os.environ["HA_ENABLE_API"] = "true" + os.environ["AT_ENABLE_API"] = "true" + try: + assert is_api_instrumentation_enabled() is False + finally: + os.environ.pop("AT_ENABLE_API", None) + + +# --------------------------------------------------------------------------- # +# Agent.instrument() gating integration +# --------------------------------------------------------------------------- # +def _build_agent(): + _uninstrument_all() + maybe_set_genai_payload_capture_env_vars() + ag = Agent() + ag._init.init_trace_provider() + return ag + + +def _record_instrumented(monkeypatch): + instrumented = [] + monkeypatch.setattr( + Agent, + "_instrument", + lambda self, library_key, app=None, auto_instrument=False: instrumented.append(library_key), + ) + return instrumented + + +def test_instrument_default_off_instruments_nothing(monkeypatch): + ag = _build_agent() + instrumented = _record_instrumented(monkeypatch) + contrib_calls = [] + monkeypatch.setattr( + "harness_sdk.agent.instrument_supported_contrib_without_wrapper", + lambda skip_libraries=None: contrib_calls.append(skip_libraries), + ) + ag.instrument() + assert instrumented == [] + assert contrib_calls == [] + + +def test_instrument_single_ai_provider_only(monkeypatch): + os.environ["HARNESS_ENABLE_AI_OPENAI"] = "true" + ag = _build_agent() + instrumented = _record_instrumented(monkeypatch) + contrib_calls = [] + monkeypatch.setattr( + "harness_sdk.agent.instrument_supported_contrib_without_wrapper", + lambda skip_libraries=None: contrib_calls.append(skip_libraries), + ) + ag.instrument() + assert instrumented == [OPENAI_KEY] + # API-only generic contrib must stay off when only an AI provider is enabled. + assert contrib_calls == [] + + +def test_instrument_api_flag_enables_api_and_generic_contrib(monkeypatch): + os.environ["HARNESS_ENABLE_API"] = "true" + ag = _build_agent() + instrumented = _record_instrumented(monkeypatch) + contrib_calls = [] + monkeypatch.setattr( + "harness_sdk.agent.instrument_supported_contrib_without_wrapper", + lambda skip_libraries=None: contrib_calls.append(skip_libraries), + ) + ag.instrument() + assert set(instrumented) == set(API_LIBRARIES) + for key in AI_LIBRARY_ENV_FLAGS: + assert key not in instrumented + assert len(contrib_calls) == 1 + + +def test_skip_libraries_overrides_enabled_category(monkeypatch): + os.environ["HARNESS_ENABLE_API"] = "true" + ag = _build_agent() + instrumented = _record_instrumented(monkeypatch) + monkeypatch.setattr( + "harness_sdk.agent.instrument_supported_contrib_without_wrapper", + lambda skip_libraries=None: None, + ) + ag.instrument(skip_libraries=[REQUESTS_KEY]) + assert REQUESTS_KEY not in instrumented + assert set(instrumented) == set(API_LIBRARIES) - {REQUESTS_KEY}