struct-sdk 0.1.0__py3-none-any.whl

struct_sdk/claude_agent.py

@@ -0,0 +1,85 @@
+"""Claude Agent SDK auto-instrumentation.
+
+Patches ``ClaudeAgentOptions`` to automatically inject OTel env vars into
+every Claude Code subprocess, regardless of whether the user calls
+``query()`` or ``ClaudeSDKClient``.
+
+Subagents inherit the parent's env vars automatically.
+
+Auto-applied by struct.init() when the claude-agent-sdk package is installed.
+"""
+
+from __future__ import annotations
+
+import functools
+import logging
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from struct_sdk.core import StructSDK
+
+logger = logging.getLogger("struct_sdk.claude_agent")
+
+
+def patch(sdk: StructSDK) -> None:
+    """Patch claude_agent_sdk to inject OTel env vars. Raises ImportError if not installed."""
+    import claude_agent_sdk
+    from claude_agent_sdk import ClaudeAgentOptions
+
+    if getattr(claude_agent_sdk, "__struct_patched", False):
+        return
+
+    # Patch ClaudeAgentOptions.__init__ to merge OTel env vars into every instance.
+    # This covers both query() and ClaudeSDKClient usage patterns.
+    original_init = ClaudeAgentOptions.__init__
+
+    otel_env = _build_otel_env(sdk)
+
+    @functools.wraps(original_init)
+    def patched_init(self: Any, *args: Any, **kwargs: Any) -> None:
+        original_init(self, *args, **kwargs)
+        # Merge OTel env vars into the options' env dict
+        existing_env = getattr(self, "env", None) or {}
+        self.env = {**existing_env, **otel_env}
+
+    patched_init.__struct_wrapped__ = True  # type: ignore[attr-defined]
+    patched_init.__struct_original__ = original_init  # type: ignore[attr-defined]
+    ClaudeAgentOptions.__init__ = patched_init  # type: ignore[method-assign]
+    claude_agent_sdk.__struct_patched = True  # type: ignore[attr-defined]
+
+
+def unpatch() -> None:
+    """Remove patches."""
+    try:
+        import claude_agent_sdk
+        from claude_agent_sdk import ClaudeAgentOptions
+    except ImportError:
+        return
+
+    if not getattr(claude_agent_sdk, "__struct_patched", False):
+        return
+
+    if getattr(ClaudeAgentOptions.__init__, "__struct_wrapped__", False):
+        ClaudeAgentOptions.__init__ = ClaudeAgentOptions.__init__.__struct_original__  # type: ignore[method-assign,union-attr]
+
+    claude_agent_sdk.__struct_patched = False  # type: ignore[attr-defined]
+
+
+def _build_otel_env(sdk: StructSDK) -> dict[str, str]:
+    """Build env vars for Claude Code's built-in OTel export."""
+    env = {
+        "CLAUDE_CODE_ENABLE_TELEMETRY": "1",
+        "CLAUDE_CODE_ENHANCED_TELEMETRY_BETA": "1",
+        "OTEL_METRICS_EXPORTER": "otlp",
+        "OTEL_LOGS_EXPORTER": "otlp",
+        "OTEL_TRACES_EXPORTER": "otlp",
+        "OTEL_EXPORTER_OTLP_PROTOCOL": "http/protobuf",
+        "OTEL_EXPORTER_OTLP_ENDPOINT": sdk._endpoint,
+        "OTEL_EXPORTER_OTLP_HEADERS": f"x-struct-ingest-key={sdk._ingest_key}",
+        "OTEL_LOG_TOOL_DETAILS": "1",
+    }
+    if sdk.capture_content:
+        env["OTEL_LOG_TOOL_CONTENT"] = "1"
+        env["OTEL_LOG_USER_PROMPTS"] = "1"
+        env["OTEL_LOG_RAW_API_BODIES"] = "1"
+    return env