ioa-observe-sdk 1.0.28__py3-none-any.whl → 1.0.30__py3-none-any.whl

ioa_observe/sdk/instrumentations/a2a.py

@@ -19,6 +19,72 @@ _instruments = ("a2a-sdk >= 0.3.0",)
 _global_tracer = None
 _kv_lock = threading.RLock()  # Add thread-safety for kv_store operations
 
+# Track original methods for uninstrumentation
+_original_methods = {}
+
+
+def _inject_observe_metadata(request, span_name: str):
+    """
+    Helper function to inject observability metadata into request.
+    This is shared between legacy and new client instrumentation.
+    """
+    traceparent = get_current_traceparent()
+    session_id = None
+    if traceparent:
+        session_id = kv_store.get(f"execution.{traceparent}")
+        if not session_id:
+            session_id = get_value("session.id")
+            if session_id:
+                kv_store.set(f"execution.{traceparent}", session_id)
+
+    # Ensure metadata dict exists - handle both request.params.metadata and request.metadata
+    try:
+        if hasattr(request, "params") and hasattr(request.params, "metadata"):
+            md = getattr(request.params, "metadata", None)
+        else:
+            md = getattr(request, "metadata", None)
+    except AttributeError:
+        md = None
+    metadata = md if isinstance(md, dict) else {}
+
+    observe_meta = dict(metadata.get("observe", {}))
+
+    # Inject W3C trace context + baggage into observe_meta
+    TraceContextTextMapPropagator().inject(carrier=observe_meta)
+    W3CBaggagePropagator().inject(carrier=observe_meta)
+
+    if traceparent:
+        observe_meta["traceparent"] = traceparent
+    if session_id:
+        observe_meta["session_id"] = session_id
+        baggage.set_baggage(f"execution.{traceparent}", session_id)
+
+    metadata["observe"] = observe_meta
+
+    # Write back metadata (pydantic models are mutable by default in v2)
+    try:
+        if hasattr(request, "params") and hasattr(request.params, "metadata"):
+            request.params.metadata = metadata
+        else:
+            request.metadata = metadata
+    except Exception:
+        # Fallback - create a new request with updated metadata
+        try:
+            if hasattr(request, "params"):
+                request = request.model_copy(
+                    update={
+                        "params": request.params.model_copy(
+                            update={"metadata": metadata}
+                        )
+                    }
+                )
+            else:
+                request = request.model_copy(update={"metadata": metadata})
+        except Exception:
+            pass  # If all else fails, continue without metadata injection
+
+    return request
+
 
 class A2AInstrumentor(BaseInstrumentor):
     def __init__(self):
@@ -35,10 +101,105 @@ class A2AInstrumentor(BaseInstrumentor):
         if importlib.util.find_spec("a2a") is None:
             raise ImportError("No module named 'a2a-sdk'. Please install it first.")
 
-        # Instrument client send_message
-        from a2a.client import A2AClient
+        # Instrument new A2AClient (BaseClient) from a2a.client.base_client (v0.3.0+)
+        self._instrument_base_client()
+
+        # Instrument legacy A2AClient for backward compatibility
+        self._instrument_legacy_client()
+
+        # Instrument server handler
+        self._instrument_server_handler()
+
+        # Instrument slima2a if available
+        self._instrument_slima2a()
+
+    def _instrument_base_client(self):
+        """Instrument the new A2AClient (BaseClient) pattern introduced in v0.3.0."""
+        import importlib
+
+        # Check if the new base_client module exists
+        if importlib.util.find_spec("a2a.client.base_client") is None:
+            return
+
+        try:
+            from a2a.client.base_client import A2AClient as BaseA2AClient
+        except ImportError:
+            return
+
+        # Instrument send_message
+        if hasattr(BaseA2AClient, "send_message"):
+            original_send_message = BaseA2AClient.send_message
+            _original_methods["BaseA2AClient.send_message"] = original_send_message
+
+            @functools.wraps(original_send_message)
+            async def instrumented_send_message(self, request, *args, **kwargs):
+                nonlocal original_send_message
+                with _global_tracer.start_as_current_span("a2a.client.send_message"):
+                    request = _inject_observe_metadata(
+                        request, "a2a.client.send_message"
+                    )
+                return await original_send_message(self, request, *args, **kwargs)
+
+            BaseA2AClient.send_message = instrumented_send_message
+
+        # Instrument send_message_streaming
+        if hasattr(BaseA2AClient, "send_message_streaming"):
+            original_send_message_streaming = BaseA2AClient.send_message_streaming
+            _original_methods["BaseA2AClient.send_message_streaming"] = (
+                original_send_message_streaming
+            )
+
+            @functools.wraps(original_send_message_streaming)
+            async def instrumented_send_message_streaming(
+                self, request, *args, **kwargs
+            ):
+                nonlocal original_send_message_streaming
+                with _global_tracer.start_as_current_span(
+                    "a2a.client.send_message_streaming"
+                ):
+                    request = _inject_observe_metadata(
+                        request, "a2a.client.send_message_streaming"
+                    )
+                # This is an async generator, so we need to yield from it
+                async for response in original_send_message_streaming(
+                    self, request, *args, **kwargs
+                ):
+                    yield response
+
+            BaseA2AClient.send_message_streaming = instrumented_send_message_streaming
+
+    def _instrument_legacy_client(self):
+        """Instrument the legacy A2AClient for backward compatibility."""
+
+        # Try to import from legacy location first (v0.3.0+)
+        legacy_client = None
+        legacy_module_path = None
 
-        original_send_message = A2AClient.send_message
+        try:
+            from a2a.client.legacy import A2AClient as LegacyA2AClient
+
+            legacy_client = LegacyA2AClient
+            legacy_module_path = "a2a.client.legacy"
+        except ImportError:
+            # Fall back to old import path (pre-v0.3.0)
+            try:
+                from a2a.client import A2AClient as LegacyA2AClient
+
+                # Check if this is the legacy client (has send_message method directly)
+                if hasattr(LegacyA2AClient, "send_message"):
+                    legacy_client = LegacyA2AClient
+                    legacy_module_path = "a2a.client"
+            except ImportError:
+                pass
+
+        if legacy_client is None:
+            return
+
+        # Store original methods for uninstrumentation
+        original_send_message = legacy_client.send_message
+        _original_methods[f"{legacy_module_path}.A2AClient.send_message"] = (
+            original_send_message
+        )
 
         @functools.wraps(original_send_message)
         async def instrumented_send_message(self, request, *args, **kwargs):
@@ -90,11 +251,14 @@ class A2AInstrumentor(BaseInstrumentor):
             # Call through without transport-specific kwargs
             return await original_send_message(self, request, *args, **kwargs)
 
-        A2AClient.send_message = instrumented_send_message
+        legacy_client.send_message = instrumented_send_message
 
-        # Instrument broadcast_message
-        if hasattr(A2AClient, "broadcast_message"):
-            original_broadcast_message = A2AClient.broadcast_message
+        # Instrument broadcast_message if it exists
+        if hasattr(legacy_client, "broadcast_message"):
+            original_broadcast_message = legacy_client.broadcast_message
+            _original_methods[f"{legacy_module_path}.A2AClient.broadcast_message"] = (
+                original_broadcast_message
+            )
 
             @functools.wraps(original_broadcast_message)
             async def instrumented_broadcast_message(self, request, *args, **kwargs):
@@ -146,12 +310,16 @@ class A2AInstrumentor(BaseInstrumentor):
                 # Call through without transport-specific kwargs
                 return await original_broadcast_message(self, request, *args, **kwargs)
 
-            A2AClient.broadcast_message = instrumented_broadcast_message
+            legacy_client.broadcast_message = instrumented_broadcast_message
 
-        # Instrument server handler
+    def _instrument_server_handler(self):
+        """Instrument the server-side request handlers."""
         from a2a.server.request_handlers import DefaultRequestHandler
 
         original_server_on_message_send = DefaultRequestHandler.on_message_send
+        _original_methods["DefaultRequestHandler.on_message_send"] = (
+            original_server_on_message_send
+        )
 
         @functools.wraps(original_server_on_message_send)
         async def instrumented_on_message_send(self, params, context):
@@ -190,173 +358,314 @@ class A2AInstrumentor(BaseInstrumentor):
             finally:
                 if token is not None:
                     try:
+                        from opentelemetry import context as otel_ctx
+
                         otel_ctx.detach(token)
                     except Exception:
                         pass
 
         DefaultRequestHandler.on_message_send = instrumented_on_message_send
 
-        # from a2a.client import A2AClient
+        # Instrument on_message_send_stream for streaming message reception
+        if hasattr(DefaultRequestHandler, "on_message_send_stream"):
+            original_on_message_send_stream = (
+                DefaultRequestHandler.on_message_send_stream
+            )
+            _original_methods["DefaultRequestHandler.on_message_send_stream"] = (
+                original_on_message_send_stream
+            )
 
-        # A2AClient.send_message = instrumented_send_message
+            @functools.wraps(original_on_message_send_stream)
+            async def instrumented_on_message_send_stream(self, params, context):
+                # Read context from A2A message metadata (transport-agnostic)
+                try:
+                    metadata = getattr(params, "metadata", {}) or {}
+                except Exception:
+                    metadata = {}
+
+                carrier = {}
+                observe_meta = metadata.get("observe", {}) or {}
+                for k in ("traceparent", "baggage", "session_id"):
+                    if k in observe_meta:
+                        carrier[k] = observe_meta[k]
+
+                token = None
+                if carrier.get("traceparent"):
+                    ctx = TraceContextTextMapPropagator().extract(carrier=carrier)
+                    ctx = W3CBaggagePropagator().extract(carrier=carrier, context=ctx)
+                    try:
+                        from opentelemetry import context as otel_ctx
 
-        # from a2a.server.request_handlers import DefaultRequestHandler
+                        token = otel_ctx.attach(ctx)
+                    except Exception:
+                        token = None
 
-        # original_server_on_message_send = DefaultRequestHandler.on_message_send
+                    session_id = observe_meta.get("session_id")
+                    if session_id and session_id != "None":
+                        set_session_id(
+                            session_id, traceparent=carrier.get("traceparent")
+                        )
+                        kv_store.set(
+                            f"execution.{carrier.get('traceparent')}", session_id
+                        )
 
-        # Instrumentation for slima2a
-        if importlib.util.find_spec("slima2a"):
-            from slima2a.client_transport import SRPCTransport
+                try:
+                    # This is an async generator
+                    async for event in original_on_message_send_stream(
+                        self, params, context
+                    ):
+                        yield event
+                finally:
+                    if token is not None:
+                        try:
+                            from opentelemetry import context as otel_ctx
+
+                            otel_ctx.detach(token)
+                        except Exception:
+                            pass
+
+            DefaultRequestHandler.on_message_send_stream = (
+                instrumented_on_message_send_stream
+            )
 
-            # send_message
-            original_srpc_transport_send_message = SRPCTransport.send_message
+    def _instrument_slima2a(self):
+        """Instrument slima2a transport if available."""
+        import importlib
 
-            @functools.wraps(original_srpc_transport_send_message)
-            async def instrumented_srpc_transport_send_message(
-                self, request, *args, **kwargs
-            ):
-                # Put context into A2A message metadata instead of HTTP headers
-                with _global_tracer.start_as_current_span("slima2a.send_message"):
-                    traceparent = get_current_traceparent()
-                    session_id = None
-                    if traceparent:
-                        session_id = kv_store.get(f"execution.{traceparent}")
-                        if not session_id:
-                            session_id = get_value("session.id")
-                            if session_id:
-                                kv_store.set(f"execution.{traceparent}", session_id)
+        if importlib.util.find_spec("slima2a") is None:
+            return
 
-                    # Ensure metadata dict exists
-                    try:
-                        md = getattr(request.params, "metadata", None)
-                    except AttributeError:
-                        md = None
-                    metadata = md if isinstance(md, dict) else {}
+        from slima2a.client_transport import SRPCTransport
 
-                    observe_meta = dict(metadata.get("observe", {}))
+        # send_message
+        original_srpc_transport_send_message = SRPCTransport.send_message
+        _original_methods["SRPCTransport.send_message"] = (
+            original_srpc_transport_send_message
+        )
 
-                    # Inject W3C trace context + baggage into observe_meta
-                    TraceContextTextMapPropagator().inject(carrier=observe_meta)
-                    W3CBaggagePropagator().inject(carrier=observe_meta)
+        @functools.wraps(original_srpc_transport_send_message)
+        async def instrumented_srpc_transport_send_message(
+            self, request, *args, **kwargs
+        ):
+            # Put context into A2A message metadata instead of HTTP headers
+            with _global_tracer.start_as_current_span("slima2a.send_message"):
+                traceparent = get_current_traceparent()
+                session_id = None
+                if traceparent:
+                    session_id = kv_store.get(f"execution.{traceparent}")
+                    if not session_id:
+                        session_id = get_value("session.id")
+                        if session_id:
+                            kv_store.set(f"execution.{traceparent}", session_id)
 
-                    if traceparent:
-                        observe_meta["traceparent"] = traceparent
-                    if session_id:
-                        observe_meta["session_id"] = session_id
-                        baggage.set_baggage(f"execution.{traceparent}", session_id)
+                # Ensure metadata dict exists
+                try:
+                    md = getattr(request.params, "metadata", None)
+                except AttributeError:
+                    md = None
+                metadata = md if isinstance(md, dict) else {}
 
-                    metadata["observe"] = observe_meta
+                observe_meta = dict(metadata.get("observe", {}))
 
-                    # Write back metadata (pydantic models are mutable by default in v2)
-                    try:
-                        request.metadata = metadata
-                    except Exception:
-                        # Fallback
-                        request = request.model_copy(update={"metadata": metadata})
+                # Inject W3C trace context + baggage into observe_meta
+                TraceContextTextMapPropagator().inject(carrier=observe_meta)
+                W3CBaggagePropagator().inject(carrier=observe_meta)
 
-                # Call through without transport-specific kwargs
-                return await original_srpc_transport_send_message(
-                    self, request, *args, **kwargs
-                )
+                if traceparent:
+                    observe_meta["traceparent"] = traceparent
+                if session_id:
+                    observe_meta["session_id"] = session_id
+                    baggage.set_baggage(f"execution.{traceparent}", session_id)
 
-            SRPCTransport.send_message = instrumented_srpc_transport_send_message
+                metadata["observe"] = observe_meta
 
-            # send_message_streaming
-            original_srpc_transport_send_message_streaming = (
-                SRPCTransport.send_message_streaming
-            )
+                # Write back metadata (pydantic models are mutable by default in v2)
+                try:
+                    request.metadata = metadata
+                except Exception:
+                    # Fallback
+                    request = request.model_copy(update={"metadata": metadata})
 
-            @functools.wraps(original_srpc_transport_send_message_streaming)
-            async def instrumented_srpc_transport_send_message_streaming(
+            # Call through without transport-specific kwargs
+            return await original_srpc_transport_send_message(
                 self, request, *args, **kwargs
-            ):
-                # Put context into A2A message metadata instead of HTTP headers
-                with _global_tracer.start_as_current_span(
-                    "slima2a.send_message_streaming"
-                ):
-                    traceparent = get_current_traceparent()
-                    session_id = None
-                    if traceparent:
-                        session_id = kv_store.get(f"execution.{traceparent}")
-                        if not session_id:
-                            session_id = get_value("session.id")
-                            if session_id:
-                                kv_store.set(f"execution.{traceparent}", session_id)
+            )
 
-                    # Ensure metadata dict exists
-                    try:
-                        md = getattr(request.params, "metadata", None)
-                    except AttributeError:
-                        md = None
-                    metadata = md if isinstance(md, dict) else {}
+        SRPCTransport.send_message = instrumented_srpc_transport_send_message
 
-                    observe_meta = dict(metadata.get("observe", {}))
+        # send_message_streaming
+        original_srpc_transport_send_message_streaming = (
+            SRPCTransport.send_message_streaming
+        )
+        _original_methods["SRPCTransport.send_message_streaming"] = (
+            original_srpc_transport_send_message_streaming
+        )
 
-                    # Inject W3C trace context + baggage into observe_meta
-                    TraceContextTextMapPropagator().inject(carrier=observe_meta)
-                    W3CBaggagePropagator().inject(carrier=observe_meta)
+        @functools.wraps(original_srpc_transport_send_message_streaming)
+        async def instrumented_srpc_transport_send_message_streaming(
+            self, request, *args, **kwargs
+        ):
+            # Put context into A2A message metadata instead of HTTP headers
+            with _global_tracer.start_as_current_span("slima2a.send_message_streaming"):
+                traceparent = get_current_traceparent()
+                session_id = None
+                if traceparent:
+                    session_id = kv_store.get(f"execution.{traceparent}")
+                    if not session_id:
+                        session_id = get_value("session.id")
+                        if session_id:
+                            kv_store.set(f"execution.{traceparent}", session_id)
 
-                    if traceparent:
-                        observe_meta["traceparent"] = traceparent
-                    if session_id:
-                        observe_meta["session_id"] = session_id
-                        baggage.set_baggage(f"execution.{traceparent}", session_id)
+                # Ensure metadata dict exists
+                try:
+                    md = getattr(request.params, "metadata", None)
+                except AttributeError:
+                    md = None
+                metadata = md if isinstance(md, dict) else {}
 
-                    metadata["observe"] = observe_meta
+                observe_meta = dict(metadata.get("observe", {}))
 
-                    # Write back metadata (pydantic models are mutable by default in v2)
-                    try:
-                        request.metadata = metadata
-                    except Exception:
-                        # Fallback
-                        request = request.model_copy(update={"metadata": metadata})
+                # Inject W3C trace context + baggage into observe_meta
+                TraceContextTextMapPropagator().inject(carrier=observe_meta)
+                W3CBaggagePropagator().inject(carrier=observe_meta)
 
-                # Call through without transport-specific kwargs
-                return await original_srpc_transport_send_message_streaming(
-                    self, request, *args, **kwargs
-                )
+                if traceparent:
+                    observe_meta["traceparent"] = traceparent
+                if session_id:
+                    observe_meta["session_id"] = session_id
+                    baggage.set_baggage(f"execution.{traceparent}", session_id)
 
-            SRPCTransport.send_message_streaming = (
-                instrumented_srpc_transport_send_message_streaming
+                metadata["observe"] = observe_meta
+
+                # Write back metadata (pydantic models are mutable by default in v2)
+                try:
+                    request.metadata = metadata
+                except Exception:
+                    # Fallback
+                    request = request.model_copy(update={"metadata": metadata})
+
+            # Call through without transport-specific kwargs
+            return await original_srpc_transport_send_message_streaming(
+                self, request, *args, **kwargs
             )
 
+        SRPCTransport.send_message_streaming = (
+            instrumented_srpc_transport_send_message_streaming
+        )
+
     def _uninstrument(self, **kwargs):
         import importlib
 
         if importlib.util.find_spec("a2a") is None:
             raise ImportError("No module named 'a2a-sdk'. Please install it first.")
 
-        # Uninstrument `send_message`
-        from a2a.client import A2AClient
+        # Uninstrument BaseA2AClient (new pattern v0.3.0+)
+        if importlib.util.find_spec("a2a.client.base_client") is not None:
+            try:
+                from a2a.client.base_client import A2AClient as BaseA2AClient
+
+                if "BaseA2AClient.send_message" in _original_methods:
+                    BaseA2AClient.send_message = _original_methods[
+                        "BaseA2AClient.send_message"
+                    ]
+                elif hasattr(BaseA2AClient.send_message, "__wrapped__"):
+                    BaseA2AClient.send_message = BaseA2AClient.send_message.__wrapped__
+
+                if "BaseA2AClient.send_message_streaming" in _original_methods:
+                    BaseA2AClient.send_message_streaming = _original_methods[
+                        "BaseA2AClient.send_message_streaming"
+                    ]
+                elif hasattr(BaseA2AClient, "send_message_streaming") and hasattr(
+                    BaseA2AClient.send_message_streaming, "__wrapped__"
+                ):
+                    BaseA2AClient.send_message_streaming = (
+                        BaseA2AClient.send_message_streaming.__wrapped__
+                    )
+            except ImportError:
+                pass
 
-        A2AClient.send_message = A2AClient.send_message.__wrapped__
+        # Uninstrument legacy A2AClient
+        legacy_client = None
+        legacy_module_path = None
 
-        # Uninstrument `broadcast_message`
-        if hasattr(A2AClient, "broadcast_message") and hasattr(
-            A2AClient.broadcast_message, "__wrapped__"
-        ):
-            A2AClient.broadcast_message = A2AClient.broadcast_message.__wrapped__
+        try:
+            from a2a.client.legacy import A2AClient as LegacyA2AClient
+
+            legacy_client = LegacyA2AClient
+            legacy_module_path = "a2a.client.legacy"
+        except ImportError:
+            try:
+                from a2a.client import A2AClient as LegacyA2AClient
+
+                if hasattr(LegacyA2AClient, "send_message"):
+                    legacy_client = LegacyA2AClient
+                    legacy_module_path = "a2a.client"
+            except ImportError:
+                pass
+
+        if legacy_client is not None:
+            key = f"{legacy_module_path}.A2AClient.send_message"
+            if key in _original_methods:
+                legacy_client.send_message = _original_methods[key]
+            elif hasattr(legacy_client.send_message, "__wrapped__"):
+                legacy_client.send_message = legacy_client.send_message.__wrapped__
+
+            broadcast_key = f"{legacy_module_path}.A2AClient.broadcast_message"
+            if broadcast_key in _original_methods:
+                legacy_client.broadcast_message = _original_methods[broadcast_key]
+            elif hasattr(legacy_client, "broadcast_message") and hasattr(
+                legacy_client.broadcast_message, "__wrapped__"
+            ):
+                legacy_client.broadcast_message = (
+                    legacy_client.broadcast_message.__wrapped__
+                )
 
         # Uninstrument server handler
         from a2a.server.request_handlers import DefaultRequestHandler
 
-        DefaultRequestHandler.on_message_send = (
-            DefaultRequestHandler.on_message_send.__wrapped__
-        )
+        if "DefaultRequestHandler.on_message_send" in _original_methods:
+            DefaultRequestHandler.on_message_send = _original_methods[
+                "DefaultRequestHandler.on_message_send"
+            ]
+        elif hasattr(DefaultRequestHandler.on_message_send, "__wrapped__"):
+            DefaultRequestHandler.on_message_send = (
+                DefaultRequestHandler.on_message_send.__wrapped__
+            )
+
+        if "DefaultRequestHandler.on_message_send_stream" in _original_methods:
+            DefaultRequestHandler.on_message_send_stream = _original_methods[
+                "DefaultRequestHandler.on_message_send_stream"
+            ]
+        elif hasattr(DefaultRequestHandler, "on_message_send_stream") and hasattr(
+            DefaultRequestHandler.on_message_send_stream, "__wrapped__"
+        ):
+            DefaultRequestHandler.on_message_send_stream = (
+                DefaultRequestHandler.on_message_send_stream.__wrapped__
+            )
 
-        # handle slima2a
+        # Uninstrument slima2a
         if importlib.util.find_spec("slima2a"):
             from slima2a.client_transport import SRPCTransport
 
-            # Uninstrument `send_message`
-            if hasattr(SRPCTransport, "send_message") and hasattr(
+            if "SRPCTransport.send_message" in _original_methods:
+                SRPCTransport.send_message = _original_methods[
+                    "SRPCTransport.send_message"
+                ]
+            elif hasattr(SRPCTransport, "send_message") and hasattr(
                 SRPCTransport.send_message, "__wrapped__"
             ):
                 SRPCTransport.send_message = SRPCTransport.send_message.__wrapped__
 
-            if hasattr(SRPCTransport, "send_message_streaming") and hasattr(
+            if "SRPCTransport.send_message_streaming" in _original_methods:
+                SRPCTransport.send_message_streaming = _original_methods[
+                    "SRPCTransport.send_message_streaming"
+                ]
+            elif hasattr(SRPCTransport, "send_message_streaming") and hasattr(
                 SRPCTransport.send_message_streaming, "__wrapped__"
             ):
                 SRPCTransport.send_message_streaming = (
-                    SRPCTransport.send_message.__wrapped__
+                    SRPCTransport.send_message_streaming.__wrapped__
                 )
+
+        # Clear stored original methods
+        _original_methods.clear()