gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/tracing/instrumentors/claude_agent_sdk.py

@@ -0,0 +1,766 @@
+# Copyright (c) Gateway. All rights reserved.
+
+"""Instrumentor for claude-agent-sdk.
+
+This module provides automatic instrumentation for the Claude Agent SDK,
+capturing conversation turns, tool calls, and response metadata.
+
+Spans are annotated with semantic attributes following Gateway conventions,
+ensuring proper processing and display in the Gateway dashboard.
+"""
+
+from __future__ import annotations
+
+import functools
+import json
+import logging
+from typing import TYPE_CHECKING, Any, AsyncIterator, Optional
+
+from opentelemetry import trace
+from opentelemetry.context import Context
+from opentelemetry.trace import Span, SpanKind, Status, StatusCode
+
+from ..semconv import (  # Gateway observation attributes (primary namespace, resolved by attributeAliases.ts); OpenTelemetry GenAI attributes (widely supported); OpenInference attributes; GatewaySDK custom attributes; Claude-specific attributes
+    CLAUDE_API_ERROR,
+    CLAUDE_ERRORS,
+    CLAUDE_MESSAGE_UUID,
+    CLAUDE_MODEL_USAGE,
+    CLAUDE_PARENT_TOOL_USE_ID,
+    CLAUDE_RATE_LIMIT_STATUS,
+    CLAUDE_RATE_LIMIT_TYPE,
+    CLAUDE_STOP_REASON,
+    CLAUDE_STRUCTURED_OUTPUT,
+    CLAUDE_SYSTEM_DATA,
+    CLAUDE_SYSTEM_SUBTYPE,
+    CLAUDE_THINKING_SIGNATURE,
+    GEN_AI_CONVERSATION_ID,
+    GEN_AI_OPERATION_NAME,
+    GEN_AI_REQUEST_MODEL,
+    GEN_AI_RESPONSE_FINISH_REASONS,
+    GEN_AI_RESPONSE_ID,
+    GEN_AI_RESPONSE_MODEL,
+    GEN_AI_SYSTEM,
+    GEN_AI_TOOL_CALL_ARGUMENTS,
+    GEN_AI_TOOL_CALL_RESULT,
+    GEN_AI_TOOL_NAME,
+    GEN_AI_USAGE_COST,
+    GEN_AI_USAGE_INPUT_TOKENS,
+    GEN_AI_USAGE_OUTPUT_TOKENS,
+    INPUT_MIME_TYPE,
+    INPUT_VALUE,
+    LLM_MODEL_NAME,
+    LLM_PROVIDER,
+    LLM_SYSTEM,
+    LLM_TOKEN_COUNT_COMPLETION,
+    LLM_TOKEN_COUNT_PROMPT,
+    LLM_TOKEN_COUNT_TOTAL,
+    GATEWAY_DURATION_API_MS,
+    GATEWAY_DURATION_MS,
+    GATEWAY_LLM_COST_USD,
+    GATEWAY_LLM_THINKING,
+    GATEWAY_NUM_TURNS,
+    GATEWAY_TOOL_IS_ERROR,
+    MIME_TYPE_JSON,
+    MIME_TYPE_TEXT,
+    MODEL_NAME,
+    OBSERVATION_INPUT,
+    OBSERVATION_OUTPUT,
+    OPENINFERENCE_SPAN_KIND,
+    OUTPUT_MIME_TYPE,
+    OUTPUT_VALUE,
+    SESSION_ID,
+    SPAN_KIND_LLM,
+    SPAN_KIND_TOOL,
+    TOOL_NAME,
+    TRACE_NAME,
+)
+from .registry import BaseInstrumentor, _is_package_available
+
+if TYPE_CHECKING:
+    from openinference.instrumentation import TraceConfig
+
+logger = logging.getLogger(__name__)
+
+# Tracer name for claude-agent-sdk spans
+TRACER_NAME = "gatewaysdk.instrumentation.claude_agent_sdk"
+
+# Attribute to store span on client instance
+_CLIENT_SPAN_ATTR = "_gatewaysdk_conversation_span"
+
+# Session IDs that indicate "no real session was provided" — ResultMessage session_id
+# should only override the span when it is NOT one of these placeholder values.
+# "default" is the instrumented_query() default parameter; placeholder values from
+# test helpers or SDK internals should not clobber a user-provided session_id.
+_SESSION_ID_PLACEHOLDERS = frozenset({"default"})
+
+# Max length for auto-generated trace names derived from the prompt.
+_TRACE_NAME_MAX_LEN = 80
+
+
+class ClaudeAgentSDKInstrumentor(BaseInstrumentor):
+    """Instrumentor for claude-agent-sdk.
+
+    This instrumentor wraps the ClaudeSDKClient to automatically create
+    OpenTelemetry spans for:
+    - Conversation turns (query calls)
+    - Tool executions
+    - Response processing
+
+    Example::
+
+        import gatewaysdk.tracing as tracing
+
+        # Auto-instruments claude-agent-sdk if installed
+        tracing.init()
+
+        # Now all ClaudeSDKClient usage is traced
+        from claude_agent_sdk import ClaudeSDKClient
+        client = ClaudeSDKClient()
+        # ...
+    """
+
+    def __init__(self):
+        self._instrumented = False
+        self._original_query: Optional[Any] = None
+        self._original_receive_messages: Optional[Any] = None
+        self._original_receive_response: Optional[Any] = None
+
+    @property
+    def name(self) -> str:
+        return "claude_agent_sdk"
+
+    @property
+    def package_name(self) -> str:
+        return "claude_agent_sdk"
+
+    def is_available(self) -> bool:
+        """Check if claude-agent-sdk is installed."""
+        return _is_package_available("claude_agent_sdk")
+
+    def instrument(self, trace_config: Optional["TraceConfig"] = None) -> None:
+        """Activate instrumentation for claude-agent-sdk.
+
+        This wraps ClaudeSDKClient.query() and receive_messages() to
+        automatically create spans for conversation turns.
+        """
+        if self._instrumented:
+            logger.debug("claude_agent_sdk already instrumented")
+            return
+
+        if not self.is_available():
+            logger.debug("claude_agent_sdk not available, skipping instrumentation")
+            return
+
+        try:
+            from claude_agent_sdk import ClaudeSDKClient
+
+            # Store original methods
+            self._original_query = ClaudeSDKClient.query
+            self._original_receive_messages = ClaudeSDKClient.receive_messages
+
+            # Create instrumented versions
+            instrumented_query = self._create_instrumented_query(self._original_query)
+            instrumented_receive_messages = self._create_instrumented_receive_messages(self._original_receive_messages)
+
+            # Patch the methods
+            ClaudeSDKClient.query = instrumented_query
+            ClaudeSDKClient.receive_messages = instrumented_receive_messages
+
+            # NOTE: do NOT also wrap receive_response(). It is a thin convenience
+            # method that internally does `async for m in self.receive_messages()`
+            # — i.e. it delegates to the (already patched) receive_messages. Also
+            # wrapping it double-instruments every message when a caller uses
+            # receive_response(): duplicated TOOL/GENERATION spans and a nested
+            # attach()/detach() that raises "Token was created in a different
+            # Context". Patching receive_messages alone covers both entry points.
+            self._original_receive_response = None
+
+            self._instrumented = True
+            logger.info("Instrumented claude_agent_sdk")
+
+        except Exception as e:
+            logger.warning("Failed to instrument claude_agent_sdk: %s", e)
+            raise
+
+    def uninstrument(self) -> None:
+        """Deactivate instrumentation for claude-agent-sdk."""
+        if not self._instrumented:
+            return
+
+        try:
+            from claude_agent_sdk import ClaudeSDKClient
+
+            # Restore original methods
+            if self._original_query is not None:
+                ClaudeSDKClient.query = self._original_query
+            if self._original_receive_messages is not None:
+                ClaudeSDKClient.receive_messages = self._original_receive_messages
+            if self._original_receive_response is not None:
+                ClaudeSDKClient.receive_response = self._original_receive_response
+
+            self._instrumented = False
+            self._original_query = None
+            self._original_receive_messages = None
+            self._original_receive_response = None
+            logger.info("Uninstrumented claude_agent_sdk")
+
+        except Exception as e:
+            logger.warning("Failed to uninstrument claude_agent_sdk: %s", e)
+
+    def _create_instrumented_query(self, original_query):
+        """Create an instrumented version of ClaudeSDKClient.query().
+
+        This creates a span and stores it on the client instance so that
+        receive_messages() can add output attributes to it later.
+        """
+
+        @functools.wraps(original_query)
+        async def instrumented_query(
+            client_self,
+            prompt: str | Any,
+            session_id: str = "default",
+        ) -> None:
+            tracer = trace.get_tracer(TRACER_NAME)
+
+            # Get model from client options
+            model = getattr(client_self.options, "model", None) or "claude"
+
+            # Create span as a NEW TRACE (fresh context = each turn is its own trace)
+            span = tracer.start_span(
+                "llm.conversation_turn",
+                kind=SpanKind.CLIENT,
+                context=Context(),  # Empty context creates a new root trace
+            )
+
+            # Record input using multiple conventions for compatibility
+            input_value = prompt if isinstance(prompt, str) else str(prompt)
+
+            # Langfuse attributes (REQUIRED for OtelIngestionProcessor to extract I/O)
+            span.set_attribute(OBSERVATION_INPUT, input_value)
+            span.set_attribute(MODEL_NAME, model)
+
+            # OpenInference attributes
+            span.set_attribute(OPENINFERENCE_SPAN_KIND, SPAN_KIND_LLM)
+            span.set_attribute(INPUT_VALUE, input_value)
+            span.set_attribute(INPUT_MIME_TYPE, MIME_TYPE_TEXT)
+            span.set_attribute(LLM_MODEL_NAME, model)
+            span.set_attribute(LLM_SYSTEM, "anthropic")
+            span.set_attribute(LLM_PROVIDER, "anthropic")
+
+            # GenAI attributes
+            span.set_attribute(GEN_AI_REQUEST_MODEL, model)
+            span.set_attribute(GEN_AI_SYSTEM, "anthropic")
+            span.set_attribute(GEN_AI_OPERATION_NAME, "chat")
+
+            # Session ID — only set when the caller provided a real value.
+            # Placeholders like "default" are skipped to avoid creating ghost
+            # sessions in Gateway that group all sessionless traces together.
+            if session_id not in _SESSION_ID_PLACEHOLDERS:
+                span.set_attribute(SESSION_ID, session_id)
+                span.set_attribute(GEN_AI_CONVERSATION_ID, session_id)
+            else:
+                # Auto-generate a session ID per client instance so conversations
+                # are grouped without requiring the user to pass one manually.
+                auto_session = getattr(client_self, "_gatewaysdk_auto_session_id", None)
+                if auto_session is None:
+                    import uuid
+
+                    auto_session = f"claude-{uuid.uuid4().hex[:12]}"
+                    client_self._gatewaysdk_auto_session_id = auto_session
+                span.set_attribute(SESSION_ID, auto_session)
+                span.set_attribute(GEN_AI_CONVERSATION_ID, auto_session)
+
+            # Trace name — derive from the prompt so traces are distinguishable
+            # in the Gateway UI. Truncate to keep names readable.
+            trace_label = input_value.strip().replace("\n", " ")
+            if len(trace_label) > _TRACE_NAME_MAX_LEN:
+                trace_label = trace_label[:_TRACE_NAME_MAX_LEN] + "…"
+            span.set_attribute(TRACE_NAME, trace_label)
+
+            # Store span on client instance for receive_messages to use
+            setattr(client_self, _CLIENT_SPAN_ATTR, span)
+
+            try:
+                # Call original method
+                result = await original_query(client_self, prompt, session_id)
+                return result
+            except Exception as e:
+                # On error, end the span here
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                span.end()
+                # Clear the stored span
+                if hasattr(client_self, _CLIENT_SPAN_ATTR):
+                    delattr(client_self, _CLIENT_SPAN_ATTR)
+                raise
+
+        return instrumented_query
+
+    def _create_instrumented_receive_messages(self, original_receive_messages):
+        """Create an instrumented version of ClaudeSDKClient.receive_messages().
+
+        This uses the span created by query() to record output attributes,
+        then ends the span when ResultMessage is received.
+
+        The output is structured as a sequence of blocks to preserve the order:
+        text → tool_call → text (if that's how the response came)
+
+        Tool calls appear as child spans under the conversation turn span.
+        """
+
+        @functools.wraps(original_receive_messages)
+        async def instrumented_receive_messages(client_self) -> AsyncIterator[Any]:
+            tracer = trace.get_tracer(TRACER_NAME)
+
+            # Import message types
+            from claude_agent_sdk import (
+                AssistantMessage,
+                ResultMessage,
+                SystemMessage,
+                TextBlock,
+                ThinkingBlock,
+                ToolResultBlock,
+                ToolUseBlock,
+                UserMessage,
+            )
+
+            # Import new message types with graceful fallback for older SDK versions
+            try:
+                from claude_agent_sdk import TaskStartedMessage
+            except ImportError:
+                TaskStartedMessage = None
+
+            try:
+                from claude_agent_sdk import TaskProgressMessage
+            except ImportError:
+                TaskProgressMessage = None
+
+            try:
+                from claude_agent_sdk import TaskNotificationMessage
+            except ImportError:
+                TaskNotificationMessage = None
+
+            try:
+                from claude_agent_sdk import RateLimitEvent
+            except ImportError:
+                RateLimitEvent = None
+
+            try:
+                from claude_agent_sdk import StreamEvent
+            except ImportError:
+                StreamEvent = None
+
+            # Get the span created by query()
+            conversation_span: Optional[Span] = getattr(client_self, _CLIENT_SPAN_ATTR, None)
+
+            # Parent context for the spans this instrumentor creates. Those spans
+            # are started with context=parent_ctx explicitly (see below), so we do
+            # NOT attach() it for the generator's lifetime. Holding an attached
+            # context across the generator's yields makes the teardown detach()
+            # run in a different context and log a benign "Failed to detach
+            # context"; and since the Claude Agent SDK executes tools in a separate
+            # subprocess, an attached parent-process context never reached them
+            # anyway. Explicit parent context gives correct nesting with no noise.
+            parent_ctx = None
+            if conversation_span:
+                parent_ctx = trace.set_span_in_context(conversation_span)
+
+            # Track state for this conversation turn
+            output_blocks = []  # Sequence of blocks: text, tool_call, text...
+            collected_thinking = []
+            thinking_signatures = []  # ThinkingBlock signatures
+            tool_spans = {}  # Map tool_use_id to span
+            current_text_buffer = []
+            system_events = []  # SystemMessage events
+
+            try:
+                async for message in original_receive_messages(client_self):
+                    if isinstance(message, AssistantMessage):
+                        model = getattr(message, "model", None)
+                        if conversation_span and conversation_span.is_recording():
+                            # Update model from response (GenAI convention)
+                            if model:
+                                conversation_span.set_attribute(GEN_AI_RESPONSE_MODEL, model)
+                                conversation_span.set_attribute(MODEL_NAME, model)
+                                conversation_span.set_attribute(LLM_MODEL_NAME, model)
+
+                            # Capture API error if present
+                            if message.error:
+                                conversation_span.set_attribute(CLAUDE_API_ERROR, message.error)
+
+                            # Capture parent_tool_use_id for sub-agent tracking
+                            if message.parent_tool_use_id:
+                                conversation_span.set_attribute(CLAUDE_PARENT_TOOL_USE_ID, message.parent_tool_use_id)
+
+                            # Message ID for correlation with Anthropic logs
+                            if hasattr(message, "message_id") and message.message_id:
+                                conversation_span.set_attribute(GEN_AI_RESPONSE_ID, message.message_id)
+
+                            # Per-turn stop reason
+                            if hasattr(message, "stop_reason") and message.stop_reason:
+                                conversation_span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [message.stop_reason])
+
+                            # Per-turn usage (captured as events since there may be multiple turns)
+                            if hasattr(message, "usage") and message.usage:
+                                usage_attrs = {}
+                                input_tokens = message.usage.get("input_tokens")
+                                output_tokens = message.usage.get("output_tokens")
+                                if input_tokens is not None:
+                                    usage_attrs[GEN_AI_USAGE_INPUT_TOKENS] = input_tokens
+                                if output_tokens is not None:
+                                    usage_attrs[GEN_AI_USAGE_OUTPUT_TOKENS] = output_tokens
+                                if usage_attrs:
+                                    conversation_span.add_event("turn_usage", attributes=usage_attrs)
+
+                        for block in message.content:
+                            if isinstance(block, TextBlock):
+                                current_text_buffer.append(block.text)
+
+                            elif isinstance(block, ThinkingBlock):
+                                collected_thinking.append(block.thinking)
+                                # Capture thinking signature if present
+                                if hasattr(block, "signature") and block.signature:
+                                    thinking_signatures.append(block.signature)
+
+                            elif isinstance(block, ToolUseBlock):
+                                # Flush any buffered text before tool call
+                                if current_text_buffer:
+                                    output_blocks.append({"type": "text", "content": "\n".join(current_text_buffer)})
+                                    current_text_buffer = []
+
+                                # Add tool call block (will be updated with result)
+                                output_blocks.append(
+                                    {
+                                        "type": "tool_call",
+                                        "name": block.name,
+                                        "id": block.id,
+                                        "input": block.input,
+                                        "output": None,
+                                    }
+                                )
+
+                                # Create a child span for this tool call
+                                # Note: Span timing is based on when we receive ToolUseBlock/ToolResultBlock messages,
+                                # not when the tool actually executes in Claude CLI. This may cause timing inaccuracies.
+                                tool_span = tracer.start_span(
+                                    f"tool.{block.name}",
+                                    kind=SpanKind.INTERNAL,
+                                    context=parent_ctx,
+                                )
+
+                                # OpenInference attributes
+                                tool_span.set_attribute(OPENINFERENCE_SPAN_KIND, SPAN_KIND_TOOL)
+                                tool_span.set_attribute(TOOL_NAME, block.name)
+
+                                # GenAI attributes
+                                tool_span.set_attribute(GEN_AI_TOOL_NAME, block.name)
+                                tool_span.set_attribute(GEN_AI_OPERATION_NAME, "tool_call")
+
+                                try:
+                                    input_json = json.dumps(block.input)
+                                    # GenAI
+                                    tool_span.set_attribute(GEN_AI_TOOL_CALL_ARGUMENTS, input_json)
+                                    # Langfuse
+                                    tool_span.set_attribute(OBSERVATION_INPUT, input_json)
+                                    # OpenInference
+                                    tool_span.set_attribute(INPUT_VALUE, input_json)
+                                    tool_span.set_attribute(INPUT_MIME_TYPE, MIME_TYPE_JSON)
+                                except (TypeError, ValueError):
+                                    input_str = str(block.input)
+                                    tool_span.set_attribute(GEN_AI_TOOL_CALL_ARGUMENTS, input_str)
+                                    tool_span.set_attribute(OBSERVATION_INPUT, input_str)
+                                    tool_span.set_attribute(INPUT_VALUE, input_str)
+                                    tool_span.set_attribute(INPUT_MIME_TYPE, MIME_TYPE_TEXT)
+
+                                tool_spans[block.id] = tool_span
+
+                    elif isinstance(message, UserMessage):
+                        # Capture UserMessage uuid and parent_tool_use_id if present
+                        if conversation_span and conversation_span.is_recording():
+                            if hasattr(message, "uuid") and message.uuid:
+                                conversation_span.set_attribute(CLAUDE_MESSAGE_UUID, message.uuid)
+                            if hasattr(message, "parent_tool_use_id") and message.parent_tool_use_id:
+                                conversation_span.set_attribute(CLAUDE_PARENT_TOOL_USE_ID, message.parent_tool_use_id)
+
+                        # Check for tool results
+                        if isinstance(message.content, list):
+                            for block in message.content:
+                                if isinstance(block, ToolResultBlock):
+                                    tool_span = tool_spans.get(block.tool_use_id)
+                                    content_str = str(block.content) if block.content else ""
+                                    is_error = getattr(block, "is_error", False) or False
+
+                                    # Update the tool_call block in output_blocks
+                                    for ob in output_blocks:
+                                        if ob.get("type") == "tool_call" and ob.get("id") == block.tool_use_id:
+                                            ob["output"] = content_str
+                                            ob["is_error"] = is_error
+                                            break
+
+                                    if tool_span:
+                                        # GenAI
+                                        tool_span.set_attribute(GEN_AI_TOOL_CALL_RESULT, content_str)
+                                        # Langfuse
+                                        tool_span.set_attribute(OBSERVATION_OUTPUT, content_str)
+                                        # OpenInference
+                                        tool_span.set_attribute(OUTPUT_VALUE, content_str)
+                                        # Determine output mime type based on content (validate JSON)
+                                        output_mime = MIME_TYPE_TEXT
+                                        stripped = content_str.strip()
+                                        if stripped and stripped[0] in ("{", "["):
+                                            try:
+                                                json.loads(content_str)
+                                                output_mime = MIME_TYPE_JSON
+                                            except (json.JSONDecodeError, ValueError):
+                                                pass
+                                        tool_span.set_attribute(OUTPUT_MIME_TYPE, output_mime)
+                                        # GatewaySDK
+                                        tool_span.set_attribute(GATEWAY_TOOL_IS_ERROR, is_error)
+
+                                        if is_error:
+                                            tool_span.set_status(Status(StatusCode.ERROR, content_str))
+                                        else:
+                                            tool_span.set_status(Status(StatusCode.OK))
+
+                                        tool_span.end()
+                                        del tool_spans[block.tool_use_id]
+
+                    # Task messages MUST be checked before SystemMessage because
+                    # they are subclasses of SystemMessage — isinstance(TaskStartedMessage(), SystemMessage) is True.
+                    elif TaskStartedMessage is not None and isinstance(message, TaskStartedMessage):
+                        if conversation_span and conversation_span.is_recording():
+                            event_attrs = {}
+                            if hasattr(message, "task_id") and message.task_id:
+                                event_attrs["claude.task.id"] = message.task_id
+                            if hasattr(message, "description") and message.description:
+                                event_attrs["claude.task.description"] = message.description
+                            if hasattr(message, "tool_use_id") and message.tool_use_id:
+                                event_attrs["claude.task.tool_use_id"] = message.tool_use_id
+                            if hasattr(message, "task_type") and message.task_type:
+                                event_attrs["claude.task.type"] = message.task_type
+                            if hasattr(message, "session_id") and message.session_id:
+                                event_attrs["claude.task.session_id"] = message.session_id
+                            conversation_span.add_event("task.started", attributes=event_attrs)
+
+                    elif TaskProgressMessage is not None and isinstance(message, TaskProgressMessage):
+                        if conversation_span and conversation_span.is_recording():
+                            event_attrs = {}
+                            if hasattr(message, "task_id") and message.task_id:
+                                event_attrs["claude.task.id"] = message.task_id
+                            if hasattr(message, "description") and message.description:
+                                event_attrs["claude.task.description"] = message.description
+                            if hasattr(message, "tool_use_id") and message.tool_use_id:
+                                event_attrs["claude.task.tool_use_id"] = message.tool_use_id
+                            if hasattr(message, "last_tool_name") and message.last_tool_name:
+                                event_attrs["claude.task.last_tool"] = message.last_tool_name
+                            if hasattr(message, "usage") and message.usage:
+                                try:
+                                    event_attrs["claude.task.usage"] = json.dumps(message.usage)
+                                except (TypeError, ValueError):
+                                    event_attrs["claude.task.usage"] = str(message.usage)
+                            conversation_span.add_event("task.progress", attributes=event_attrs)
+
+                    elif TaskNotificationMessage is not None and isinstance(message, TaskNotificationMessage):
+                        if conversation_span and conversation_span.is_recording():
+                            event_attrs = {}
+                            if hasattr(message, "task_id") and message.task_id:
+                                event_attrs["claude.task.id"] = message.task_id
+                            if hasattr(message, "status") and message.status:
+                                event_attrs["claude.task.status"] = message.status
+                            if hasattr(message, "summary") and message.summary:
+                                event_attrs["claude.task.summary"] = message.summary
+                            if hasattr(message, "tool_use_id") and message.tool_use_id:
+                                event_attrs["claude.task.tool_use_id"] = message.tool_use_id
+                            if hasattr(message, "output_file") and message.output_file:
+                                event_attrs["claude.task.output_file"] = message.output_file
+                            if hasattr(message, "usage") and message.usage:
+                                try:
+                                    event_attrs["claude.task.usage"] = json.dumps(message.usage)
+                                except (TypeError, ValueError):
+                                    event_attrs["claude.task.usage"] = str(message.usage)
+                            conversation_span.add_event("task.notification", attributes=event_attrs)
+
+                    elif isinstance(message, SystemMessage):
+                        # Generic SystemMessage handler (non-task system events)
+                        system_events.append({"subtype": message.subtype, "data": message.data})
+
+                        if conversation_span and conversation_span.is_recording():
+                            event_attrs = {CLAUDE_SYSTEM_SUBTYPE: message.subtype}
+                            if message.data:
+                                try:
+                                    event_attrs[CLAUDE_SYSTEM_DATA] = json.dumps(message.data)
+                                except (TypeError, ValueError):
+                                    event_attrs[CLAUDE_SYSTEM_DATA] = str(message.data)
+                            conversation_span.add_event(f"system.{message.subtype}", attributes=event_attrs)
+
+                    elif RateLimitEvent is not None and isinstance(message, RateLimitEvent):
+                        if conversation_span and conversation_span.is_recording():
+                            event_attrs = {}
+                            # Real SDK uses rate_limit_info nested object
+                            rli = getattr(message, "rate_limit_info", None)
+                            if rli:
+                                if hasattr(rli, "rate_limit_type") and rli.rate_limit_type:
+                                    event_attrs[CLAUDE_RATE_LIMIT_TYPE] = rli.rate_limit_type
+                                if hasattr(rli, "status") and rli.status:
+                                    event_attrs[CLAUDE_RATE_LIMIT_STATUS] = rli.status
+                            conversation_span.add_event("rate_limit", attributes=event_attrs)
+
+                    elif StreamEvent is not None and isinstance(message, StreamEvent):
+                        if conversation_span and conversation_span.is_recording():
+                            event_type = getattr(message, "type", "stream")
+                            conversation_span.add_event(f"stream.{event_type}")
+
+                    elif isinstance(message, ResultMessage):
+                        # Flush any remaining text buffer
+                        if current_text_buffer:
+                            output_blocks.append({"type": "text", "content": "\n".join(current_text_buffer)})
+                            current_text_buffer = []
+
+                        # Record final metrics on the conversation span
+                        if conversation_span and conversation_span.is_recording():
+                            # Output attributes (multiple conventions)
+                            if output_blocks:
+                                text_only = all(b["type"] == "text" for b in output_blocks)
+                                if text_only:
+                                    full_response = "\n".join(b["content"] for b in output_blocks)
+                                    # Langfuse
+                                    conversation_span.set_attribute(OBSERVATION_OUTPUT, full_response)
+                                    # OpenInference
+                                    conversation_span.set_attribute(OUTPUT_VALUE, full_response)
+                                    conversation_span.set_attribute(OUTPUT_MIME_TYPE, MIME_TYPE_TEXT)
+                                else:
+                                    output_json = json.dumps(output_blocks, indent=2)
+                                    # Langfuse
+                                    conversation_span.set_attribute(OBSERVATION_OUTPUT, output_json)
+                                    # OpenInference
+                                    conversation_span.set_attribute(OUTPUT_VALUE, output_json)
+                                    conversation_span.set_attribute(OUTPUT_MIME_TYPE, MIME_TYPE_JSON)
+
+                            # Thinking content
+                            if collected_thinking:
+                                conversation_span.set_attribute(GATEWAY_LLM_THINKING, "\n".join(collected_thinking))
+
+                            # Thinking signatures (Claude-specific)
+                            if thinking_signatures:
+                                conversation_span.set_attribute(
+                                    CLAUDE_THINKING_SIGNATURE, "\n".join(thinking_signatures)
+                                )
+
+                            # Structured output (Claude-specific)
+                            if hasattr(message, "structured_output") and message.structured_output is not None:
+                                try:
+                                    conversation_span.set_attribute(
+                                        CLAUDE_STRUCTURED_OUTPUT, json.dumps(message.structured_output)
+                                    )
+                                except (TypeError, ValueError):
+                                    conversation_span.set_attribute(
+                                        CLAUDE_STRUCTURED_OUTPUT, str(message.structured_output)
+                                    )
+
+                            # Cost (GenAI and GatewaySDK)
+                            if message.total_cost_usd is not None:
+                                conversation_span.set_attribute(GEN_AI_USAGE_COST, message.total_cost_usd)
+                                conversation_span.set_attribute(GATEWAY_LLM_COST_USD, message.total_cost_usd)
+
+                            # Token usage (GenAI and OpenInference)
+                            if message.usage:
+                                usage = message.usage
+                                input_tokens = usage.get("input_tokens")
+                                output_tokens = usage.get("output_tokens")
+
+                                if input_tokens is not None:
+                                    # GenAI
+                                    conversation_span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
+                                    # OpenInference
+                                    conversation_span.set_attribute(LLM_TOKEN_COUNT_PROMPT, input_tokens)
+
+                                if output_tokens is not None:
+                                    # GenAI
+                                    conversation_span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens)
+                                    # OpenInference
+                                    conversation_span.set_attribute(LLM_TOKEN_COUNT_COMPLETION, output_tokens)
+
+                                if input_tokens is not None and output_tokens is not None:
+                                    conversation_span.set_attribute(LLM_TOKEN_COUNT_TOTAL, input_tokens + output_tokens)
+
+                            # Duration (GatewaySDK)
+                            conversation_span.set_attribute(GATEWAY_DURATION_MS, message.duration_ms)
+                            conversation_span.set_attribute(GATEWAY_DURATION_API_MS, message.duration_api_ms)
+
+                            # Turns (GatewaySDK)
+                            conversation_span.set_attribute(GATEWAY_NUM_TURNS, message.num_turns)
+
+                            # Session ID: update from ResultMessage only when it provides a
+                            # real session (not a placeholder). In production the SDK always
+                            # returns the authoritative session_id here; placeholders only
+                            # appear in test helpers or when no session was created.
+                            if (
+                                hasattr(message, "session_id")
+                                and message.session_id
+                                and message.session_id not in _SESSION_ID_PLACEHOLDERS
+                            ):
+                                conversation_span.set_attribute(SESSION_ID, message.session_id)
+                                conversation_span.set_attribute(GEN_AI_CONVERSATION_ID, message.session_id)
+
+                            # Stop reason (why the conversation ended)
+                            if hasattr(message, "stop_reason") and message.stop_reason:
+                                conversation_span.set_attribute(CLAUDE_STOP_REASON, message.stop_reason)
+                                conversation_span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [message.stop_reason])
+
+                            # Errors (non-fatal errors during the turn)
+                            if hasattr(message, "errors") and message.errors:
+                                try:
+                                    conversation_span.set_attribute(CLAUDE_ERRORS, json.dumps(message.errors))
+                                except (TypeError, ValueError):
+                                    conversation_span.set_attribute(CLAUDE_ERRORS, str(message.errors))
+
+                            # Per-model usage breakdown
+                            if hasattr(message, "model_usage") and message.model_usage:
+                                try:
+                                    conversation_span.set_attribute(CLAUDE_MODEL_USAGE, json.dumps(message.model_usage))
+                                except (TypeError, ValueError):
+                                    conversation_span.set_attribute(CLAUDE_MODEL_USAGE, str(message.model_usage))
+
+                            # Error status
+                            if message.is_error:
+                                conversation_span.set_status(
+                                    Status(StatusCode.ERROR, message.result or "Unknown error")
+                                )
+                            else:
+                                conversation_span.set_status(Status(StatusCode.OK))
+
+                            # End the conversation span
+                            conversation_span.end()
+
+                        # Clear the stored span
+                        if hasattr(client_self, _CLIENT_SPAN_ATTR):
+                            delattr(client_self, _CLIENT_SPAN_ATTR)
+
+                    # Yield the original message unchanged
+                    yield message
+
+            finally:
+                # Clean up any unclosed tool spans
+                for tool_id, tool_span in tool_spans.items():
+                    tool_span.set_status(Status(StatusCode.ERROR, "Tool span not properly closed"))
+                    tool_span.end()
+
+                # Ensure conversation span is ended if not already
+                if hasattr(client_self, _CLIENT_SPAN_ATTR):
+                    remaining_span = getattr(client_self, _CLIENT_SPAN_ATTR)
+                    if remaining_span and remaining_span.is_recording():
+                        if current_text_buffer:
+                            output_blocks.append({"type": "text", "content": "\n".join(current_text_buffer)})
+                        if output_blocks:
+                            text_only = all(b["type"] == "text" for b in output_blocks)
+                            if text_only:
+                                remaining_span.set_attribute(
+                                    OBSERVATION_OUTPUT, "\n".join(b["content"] for b in output_blocks)
+                                )
+                            else:
+                                remaining_span.set_attribute(OBSERVATION_OUTPUT, json.dumps(output_blocks, indent=2))
+                        remaining_span.set_status(Status(StatusCode.OK))
+                        remaining_span.end()
+                    delattr(client_self, _CLIENT_SPAN_ATTR)
+
+        return instrumented_receive_messages