ioa-observe-sdk 1.0.28__py3-none-any.whl → 1.0.29__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()

ioa_observe/sdk/instrumentations/slim.py

@@ -7,16 +7,15 @@ import json
 import base64
 import threading
 
-from opentelemetry.context import get_value
 from opentelemetry import baggage, context
 from opentelemetry.baggage.propagation import W3CBaggagePropagator
+from opentelemetry.context import get_value
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 
 from ioa_observe.sdk import TracerWrapper
 from ioa_observe.sdk.client import kv_store
 from ioa_observe.sdk.tracing import set_session_id, get_current_traceparent
-from ioa_observe.sdk.tracing.context_manager import get_tracer
 
 _instruments = ("slim-bindings >= 0.4.0",)
 _global_tracer = None
@@ -81,10 +80,6 @@ class SLIMInstrumentor(BaseInstrumentor):
                 if traceparent:
                     with _kv_lock:
                         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)
 
                 headers = {
                     "session_id": session_id if session_id else None,
@@ -142,10 +137,8 @@ class SLIMInstrumentor(BaseInstrumentor):
                 if traceparent:
                     with _kv_lock:
                         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 session_id:
+                            kv_store.set(f"execution.{traceparent}", session_id)
 
                 headers = {
                     "session_id": session_id if session_id else None,
@@ -207,10 +200,6 @@ class SLIMInstrumentor(BaseInstrumentor):
                 if traceparent:
                     with _kv_lock:
                         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)
 
                 headers = {
                     "session_id": session_id if session_id else None,
@@ -531,15 +520,7 @@ class SLIMInstrumentor(BaseInstrumentor):
             if timeout is not None:
                 kwargs["timeout"] = timeout
 
-            if _global_tracer:
-                with _global_tracer.start_as_current_span(
-                    "session.get_message"
-                ) as span:
-                    if hasattr(self, "id"):
-                        span.set_attribute("slim.session.id", str(self.id))
-                    result = await original_get_message(self, **kwargs)
-            else:
-                result = await original_get_message(self, **kwargs)
+            result = await original_get_message(self, **kwargs)
 
             # Handle different return types from get_message
             if result is None:
@@ -688,8 +669,10 @@ class SLIMInstrumentor(BaseInstrumentor):
                                 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 session_id:
+                                        kv_store.set(
+                                            f"execution.{traceparent}", session_id
+                                        )
 
                         headers = {
                             "session_id": session_id if session_id else None,
@@ -800,8 +783,7 @@ class SLIMInstrumentor(BaseInstrumentor):
         @functools.wraps(original_method)
         async def instrumented_session_method(self, *args, **kwargs):
             if _global_tracer:
-                tracer = get_tracer()
-                with tracer.start_as_current_span(f"session.{method_name}"):
+                with _global_tracer.start_as_current_span(f"session.{method_name}"):
                     traceparent = get_current_traceparent()
 
                     # Handle message wrapping for publish methods
@@ -857,10 +839,6 @@ class SLIMInstrumentor(BaseInstrumentor):
                     if traceparent:
                         with _kv_lock:
                             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 or session_id:
                         headers = {

ioa_observe/sdk/tracing/tracing.py

@@ -30,8 +30,15 @@ from opentelemetry.sdk.trace.export import (
     SimpleSpanProcessor,
     BatchSpanProcessor,
 )
-from opentelemetry.trace import get_tracer_provider, ProxyTracerProvider
+from opentelemetry.trace import (
+    get_tracer_provider,
+    ProxyTracerProvider,
+    SpanContext,
+    TraceFlags,
+    NonRecordingSpan,
+)
 from opentelemetry.context import get_value, attach, set_value
+from opentelemetry.trace import set_span_in_context
 from opentelemetry.instrumentation.threading import ThreadingInstrumentor
 from opentelemetry.metrics import get_meter
 from ioa_observe.sdk.metrics.agents.agent_connections import connection_reliability
@@ -137,9 +144,12 @@ class TracerWrapper(object):
             if not TracerWrapper.endpoint:
                 return obj
 
-            # session activity tracking
-            obj._session_last_activity: dict[str, float] = {}
+            # session activity tracking: {session_id: (last_activity_time, trace_id)}
+            obj._session_last_activity: dict[str, tuple[float, int]] = {}
             obj._session_lock = threading.Lock()
+            # Track sessions that have already been ended to prevent duplicates
+            obj._ended_sessions: set[str] = set()
+            obj._ended_sessions_lock = threading.Lock()
             obj.__image_uploader = image_uploader
             # {(agent_name): [success_count, total_count]}
             obj._agent_execution_counts = {}
@@ -300,13 +310,15 @@ class TracerWrapper(object):
         while True:
             time.sleep(SESSION_WATCHER_INTERVAL_SECONDS)
             now = time.time()
-            expired: dict[str, float] = {}
+            expired: dict[str, tuple[float, int]] = {}
 
             # Find idle sessions and remove them from _session_last_activity
             with self._session_lock:
-                for session_id, last_ts in list(self._session_last_activity.items()):
+                for session_id, (last_ts, trace_id) in list(
+                    self._session_last_activity.items()
+                ):
                     if now - last_ts > SESSION_IDLE_TIMEOUT_SECONDS:
-                        expired[session_id] = last_ts
+                        expired[session_id] = (last_ts, trace_id)
                         del self._session_last_activity[session_id]
 
             # Periodic cleanup of processed spans to prevent unbounded memory growth
@@ -314,19 +326,42 @@ class TracerWrapper(object):
                 if len(self._processed_spans) > MAX_PROCESSED_SPANS_SIZE:
                     self._processed_spans.clear()
 
+            # Periodic cleanup of ended sessions to prevent unbounded memory growth
+            with self._ended_sessions_lock:
+                if len(self._ended_sessions) > MAX_PROCESSED_SPANS_SIZE:
+                    self._ended_sessions.clear()
+
             if not expired:
                 continue
 
             tracer = self.get_tracer()
 
             # Iterate over a snapshot and do *not* modify `expired` in the loop
-            for session_id, _last_ts in list(expired.items()):
-                with tracer.start_as_current_span("session.end") as span:
+            for session_id, (last_ts, trace_id) in list(expired.items()):
+                # Check if this session has already been ended to prevent duplicates
+                with self._ended_sessions_lock:
+                    if session_id in self._ended_sessions:
+                        continue
+                    self._ended_sessions.add(session_id)
+
+                # Create a parent context from the stored trace_id to keep session.end under the same trace
+                parent_span_context = SpanContext(
+                    trace_id=trace_id,
+                    span_id=0,  # Use 0 as we don't have a specific parent span
+                    is_remote=False,
+                    trace_flags=TraceFlags(TraceFlags.SAMPLED),
+                )
+                parent_span = NonRecordingSpan(parent_span_context)
+                parent_context = set_span_in_context(parent_span)
+
+                with tracer.start_as_current_span(
+                    "session.end", context=parent_context
+                ) as span:
                     span.set_attribute("session.id", session_id)
                     workflow_name = get_value("workflow_name")
                     if workflow_name:
                         span.set_attribute(OBSERVE_WORKFLOW_NAME, workflow_name)
-                    span.set_attribute("session.ended_at", _last_ts)
+                    span.set_attribute("session.ended_at", last_ts)
 
             # ensure end spans are exported reasonably fast
             self.flush()
@@ -337,11 +372,29 @@ class TracerWrapper(object):
         tracer = self.get_tracer()
 
         with self._session_lock:
-            remaining_ids = list(self._session_last_activity.keys())
+            remaining_sessions = dict(self._session_last_activity)
             self._session_last_activity.clear()
 
-        for session_id in remaining_ids:
-            with tracer.start_as_current_span("session.end") as span:
+        for session_id, (_, trace_id) in remaining_sessions.items():
+            # Check if this session has already been ended to prevent duplicates
+            with self._ended_sessions_lock:
+                if session_id in self._ended_sessions:
+                    continue
+                self._ended_sessions.add(session_id)
+
+            # Create a parent context from the stored trace_id to keep session.end under the same trace
+            parent_span_context = SpanContext(
+                trace_id=trace_id,
+                span_id=0,  # Use 0 as we don't have a specific parent span
+                is_remote=False,
+                trace_flags=TraceFlags(TraceFlags.SAMPLED),
+            )
+            parent_span = NonRecordingSpan(parent_span_context)
+            parent_context = set_span_in_context(parent_span)
+
+            with tracer.start_as_current_span(
+                "session.end", context=parent_context
+            ) as span:
                 span.set_attribute("session.id", session_id)
                 workflow_name = get_value("workflow_name")
                 if workflow_name:
@@ -441,7 +494,11 @@ class TracerWrapper(object):
         session_id = span.attributes.get("session.id")
         if session_id and span.name != "session.end":
             with self._session_lock:
-                self._session_last_activity[session_id] = time.time()
+                # Store both the last activity time and the trace_id for this session
+                self._session_last_activity[session_id] = (
+                    time.time(),
+                    span.context.trace_id,
+                )
 
         determine_reliability_score(span)
         # start_time = span.attributes.get("ioa_start_time")

{ioa_observe_sdk-1.0.28.dist-info → ioa_observe_sdk-1.0.29.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ioa-observe-sdk
-Version: 1.0.28
+Version: 1.0.29
 Summary: IOA Observability SDK
 License-Expression: Apache-2.0
 Requires-Python: >=3.10
@@ -8,37 +8,37 @@ Description-Content-Type: text/markdown
 License-File: LICENSE.md
 Requires-Dist: colorama==0.4.6
 Requires-Dist: requests>=2.32.3
-Requires-Dist: opentelemetry-api==1.37.0
+Requires-Dist: opentelemetry-api==1.39.1
 Requires-Dist: opentelemetry-distro
-Requires-Dist: opentelemetry-exporter-otlp==1.37.0
-Requires-Dist: opentelemetry-exporter-otlp-proto-common==1.37.0
-Requires-Dist: opentelemetry-exporter-otlp-proto-grpc==1.37.0
-Requires-Dist: opentelemetry-exporter-otlp-proto-http==1.37.0
+Requires-Dist: opentelemetry-exporter-otlp==1.39.1
+Requires-Dist: opentelemetry-exporter-otlp-proto-common==1.39.1
+Requires-Dist: opentelemetry-exporter-otlp-proto-grpc==1.39.1
+Requires-Dist: opentelemetry-exporter-otlp-proto-http==1.39.1
 Requires-Dist: opentelemetry-instrumentation
-Requires-Dist: opentelemetry-instrumentation-logging==0.58b0
-Requires-Dist: opentelemetry-instrumentation-openai==0.43.1
-Requires-Dist: opentelemetry-instrumentation-llamaindex==0.43.1
-Requires-Dist: opentelemetry-instrumentation-ollama==0.43.1
-Requires-Dist: opentelemetry-instrumentation-anthropic==0.43.1
-Requires-Dist: opentelemetry-instrumentation-langchain==0.43.1
-Requires-Dist: opentelemetry-instrumentation-bedrock==0.43.1
-Requires-Dist: opentelemetry-instrumentation-cohere==0.43.1
-Requires-Dist: opentelemetry-instrumentation-crewai==0.43.1
-Requires-Dist: opentelemetry-instrumentation-google-generativeai==0.43.1
-Requires-Dist: opentelemetry-instrumentation-groq==0.43.1
-Requires-Dist: opentelemetry-instrumentation-mistralai==0.43.1
-Requires-Dist: opentelemetry-instrumentation-requests==0.58b0
-Requires-Dist: opentelemetry-instrumentation-sagemaker==0.43.1
-Requires-Dist: opentelemetry-instrumentation-threading==0.58b0
-Requires-Dist: opentelemetry-instrumentation-together==0.43.1
-Requires-Dist: opentelemetry-instrumentation-transformers==0.43.1
-Requires-Dist: opentelemetry-instrumentation-urllib3==0.58b0
-Requires-Dist: opentelemetry-instrumentation-vertexai==0.43.1
-Requires-Dist: opentelemetry-proto==1.37.0
-Requires-Dist: opentelemetry-sdk==1.37.0
-Requires-Dist: opentelemetry-semantic-conventions==0.58b0
+Requires-Dist: opentelemetry-instrumentation-logging==0.60b1
+Requires-Dist: opentelemetry-instrumentation-openai==0.52.1
+Requires-Dist: opentelemetry-instrumentation-llamaindex==0.52.1
+Requires-Dist: opentelemetry-instrumentation-ollama==0.52.1
+Requires-Dist: opentelemetry-instrumentation-anthropic==0.52.1
+Requires-Dist: opentelemetry-instrumentation-langchain==0.52.1
+Requires-Dist: opentelemetry-instrumentation-bedrock==0.52.1
+Requires-Dist: opentelemetry-instrumentation-cohere==0.52.1
+Requires-Dist: opentelemetry-instrumentation-crewai==0.52.1
+Requires-Dist: opentelemetry-instrumentation-google-generativeai==0.52.1
+Requires-Dist: opentelemetry-instrumentation-groq==0.52.1
+Requires-Dist: opentelemetry-instrumentation-mistralai==0.52.1
+Requires-Dist: opentelemetry-instrumentation-requests==0.60b1
+Requires-Dist: opentelemetry-instrumentation-sagemaker==0.52.1
+Requires-Dist: opentelemetry-instrumentation-threading==0.60b1
+Requires-Dist: opentelemetry-instrumentation-together==0.52.1
+Requires-Dist: opentelemetry-instrumentation-transformers==0.52.1
+Requires-Dist: opentelemetry-instrumentation-urllib3==0.60b1
+Requires-Dist: opentelemetry-instrumentation-vertexai==0.52.1
+Requires-Dist: opentelemetry-proto==1.39.1
+Requires-Dist: opentelemetry-sdk==1.39.1
+Requires-Dist: opentelemetry-semantic-conventions==0.60b1
 Requires-Dist: opentelemetry-semantic-conventions-ai>=0.4.11
-Requires-Dist: opentelemetry-util-http==0.58b0
+Requires-Dist: opentelemetry-util-http==0.60b1
 Requires-Dist: langgraph>=0.3.2
 Requires-Dist: langchain>=0.3.19
 Requires-Dist: langchain-openai>=0.3.8

{ioa_observe_sdk-1.0.28.dist-info → ioa_observe_sdk-1.0.29.dist-info}/RECORD

@@ -14,10 +14,10 @@ ioa_observe/sdk/decorators/base.py,sha256=gezoLMLOo-CWsxzsXbQMsj5HxmchUEmlhTspbq
 ioa_observe/sdk/decorators/helpers.py,sha256=I9HXMBivkZpGDtPe9Ad_UU35p_m_wEPate4r_fU0oOA,2705
 ioa_observe/sdk/decorators/util.py,sha256=IebvH9gwZN1en3LblYJUh4bAV2STl6xmp8WpZzBDH2g,30068
 ioa_observe/sdk/instrumentations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ioa_observe/sdk/instrumentations/a2a.py,sha256=LHvvQUluPtmvbRNuoqf-EkWlZJa_NicerTpM3t1UZws,15299
+ioa_observe/sdk/instrumentations/a2a.py,sha256=smPMHhFONi9B7r2qVHe-YWoFctjHWehO_MR0kY_DgVI,27897
 ioa_observe/sdk/instrumentations/mcp.py,sha256=UG82fWitFT-K-ecnkTKZa-uFTDAs-j2HymvhUwCBRLY,18721
 ioa_observe/sdk/instrumentations/nats.py,sha256=MEYCFqHyRK5HtkjrbGVV5qIWplZ1ZjJjFCB8vtQl6qE,14736
-ioa_observe/sdk/instrumentations/slim.py,sha256=c_Q1cb6pdjZLOw1MbclUf1eRFrc_sUEe9pQJvAyoD1Y,45712
+ioa_observe/sdk/instrumentations/slim.py,sha256=pwKUnCvDmydysQ9VpqnREf4DyHjN2bZ-e54YyX-de1o,44500
 ioa_observe/sdk/logging/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ioa_observe/sdk/logging/logging.py,sha256=HZxW9s8Due7jgiNkdI38cIjv5rC9D-Flta3RQMOnpow,2891
 ioa_observe/sdk/metrics/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -35,15 +35,15 @@ ioa_observe/sdk/tracing/content_allow_list.py,sha256=1fAkpIwUQ7vDwCTkIVrqeltWQtr
 ioa_observe/sdk/tracing/context_manager.py,sha256=O0JEXYa9h8anhW78R8KKBuqS0j4by1E1KXxNIMPnLr8,400
 ioa_observe/sdk/tracing/context_utils.py,sha256=-sYS9vPLI87davV9ubneq5xqbV583CC_c0SmOQS1TAs,2933
 ioa_observe/sdk/tracing/manual.py,sha256=KS6WN-zw9vAACzXYmnMoJm9d1fenYMfvzeK1GrGDPDE,1937
-ioa_observe/sdk/tracing/tracing.py,sha256=ywvT5Z1DcaLYZQvUP30YKz6m6v76Idlvb2qszupcxS4,50774
+ioa_observe/sdk/tracing/tracing.py,sha256=pGiIa1S3562QRlTG5mANtU94wkFjKspnX21N_bJ7SfY,53359
 ioa_observe/sdk/tracing/transform_span.py,sha256=XTApi_gJxum7ynvhtcoCfDyK8VVOj91Q1DT6hAeLHA8,8419
 ioa_observe/sdk/utils/__init__.py,sha256=UPn182U-UblF_XwXaFpx8F-TmQTbm1LYf9y89uSp5Hw,704
 ioa_observe/sdk/utils/const.py,sha256=d67dUTAH9UpWvUV9GLBUqn1Sc2knJ55dy-e6YoLrvSo,1318
 ioa_observe/sdk/utils/in_memory_span_exporter.py,sha256=H_4TRaThMO1H6vUQ0OpQvzJk_fZH0OOsRAM1iZQXsR8,2112
 ioa_observe/sdk/utils/json_encoder.py,sha256=g4NQ0tTqgWssY6I1D7r4zo0G6PiUo61jhofTAw5-jno,639
 ioa_observe/sdk/utils/package_check.py,sha256=1d1MjxhwoEZIx9dumirT2pRsEWgn-m-SI4npDeEalew,576
-ioa_observe_sdk-1.0.28.dist-info/licenses/LICENSE.md,sha256=55VjUfgjWOS4vv3Cf55gfq-RxjPgRIO2vlgYPUuC5lA,11362
-ioa_observe_sdk-1.0.28.dist-info/METADATA,sha256=xEbYDkX8xS0u72TgyF-DvEM9R7BxEA8wp3khv4MWHAE,7997
-ioa_observe_sdk-1.0.28.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-ioa_observe_sdk-1.0.28.dist-info/top_level.txt,sha256=Yt-6Y1olZEDqCs2REeqI30WjYx0pLGQSVqzYmDd67N8,12
-ioa_observe_sdk-1.0.28.dist-info/RECORD,,
+ioa_observe_sdk-1.0.29.dist-info/licenses/LICENSE.md,sha256=55VjUfgjWOS4vv3Cf55gfq-RxjPgRIO2vlgYPUuC5lA,11362
+ioa_observe_sdk-1.0.29.dist-info/METADATA,sha256=qksjp3ihfWLY1ihgNoZv13Q4mluzMOfTP2kP6Tud1GA,7997
+ioa_observe_sdk-1.0.29.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+ioa_observe_sdk-1.0.29.dist-info/top_level.txt,sha256=Yt-6Y1olZEDqCs2REeqI30WjYx0pLGQSVqzYmDd67N8,12
+ioa_observe_sdk-1.0.29.dist-info/RECORD,,