ioa-observe-sdk 1.0.20__py3-none-any.whl → 1.0.22__py3-none-any.whl

ioa_observe/sdk/instrumentations/nats.py

@@ -0,0 +1,324 @@
+# Copyright AGNTCY Contributors (https://github.com/agntcy)
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Collection
+import functools
+import json
+import base64
+import threading
+
+from opentelemetry import baggage, context
+from opentelemetry.baggage.propagation import W3CBaggagePropagator
+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
+
+_instruments = ("nats-py >= 2.10.0",)
+_global_tracer = None
+_kv_lock = threading.RLock()  # Add thread-safety for kv_store operations
+
+
+class NATSInstrumentor(BaseInstrumentor):
+    def __init__(self):
+        super().__init__()
+        global _global_tracer
+        _global_tracer = TracerWrapper().get_tracer()
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        try:
+            import nats
+        except ImportError:
+            raise ImportError("No module named 'nats'. Please install it first.")
+
+        # Instrument `publish` method
+        original_publish = nats.NATS.publish
+
+        @functools.wraps(original_publish)
+        async def instrumented_publish(self, *args, **kwargs):
+            if _global_tracer:
+                with _global_tracer.start_as_current_span("nats.publish") as span:
+                    traceparent = get_current_traceparent()
+                    span.set_attribute("nats.topic", args[0] if args else None)
+            else:
+                traceparent = get_current_traceparent()
+
+            # Thread-safe access to kv_store
+            session_id = None
+            if traceparent:
+                with _kv_lock:
+                    session_id = kv_store.get(f"execution.{traceparent}")
+                    if session_id:
+                        kv_store.set(f"execution.{traceparent}", session_id)
+
+            headers = {
+                "session_id": session_id if session_id else None,
+                "traceparent": traceparent,
+            }
+
+            # Set baggage context
+            if traceparent and session_id:
+                baggage.set_baggage(f"execution.{traceparent}", session_id)
+
+            # Wrap message with headers - handle different message positions
+            message_arg_index = 1  # message will typically be the second argument
+            if len(args) > message_arg_index:
+                original_args = list(args)
+                message = original_args[message_arg_index]
+                wrapped_message = NATSInstrumentor._wrap_message_with_headers(
+                    self, message, headers
+                )
+
+                # Convert wrapped message back to bytes if needed
+                if isinstance(wrapped_message, dict):
+                    message_to_send = json.dumps(wrapped_message).encode("utf-8")
+                else:
+                    message_to_send = wrapped_message
+
+                original_args[message_arg_index] = message_to_send
+                args = tuple(original_args)
+
+            return await original_publish(self, *args, **kwargs)
+
+        nats.NATS.publish = instrumented_publish
+
+        # Instrument `request` method
+        original_request = nats.NATS.request
+
+        @functools.wraps(original_request)
+        async def instrumented_request(self, *args, **kwargs):
+            if _global_tracer:
+                with _global_tracer.start_as_current_span("nats.request") as span:
+                    traceparent = get_current_traceparent()
+                    span.set_attribute("nats.topic", args[0] if args else None)
+            else:
+                traceparent = get_current_traceparent()
+
+            # Thread-safe access to kv_store
+            session_id = None
+            if traceparent:
+                with _kv_lock:
+                    session_id = kv_store.get(f"execution.{traceparent}")
+                    if session_id:
+                        kv_store.set(f"execution.{traceparent}", session_id)
+
+            headers = {
+                "session_id": session_id if session_id else None,
+                "traceparent": traceparent,
+            }
+
+            # Set baggage context
+            if traceparent and session_id:
+                baggage.set_baggage(f"execution.{traceparent}", session_id)
+
+            # Wrap message with headers - handle different message positions
+            message_arg_index = 1  # message will typically be the second argument
+            if len(args) > message_arg_index:
+                original_args = list(args)
+                message = original_args[message_arg_index]
+                wrapped_message = NATSInstrumentor._wrap_message_with_headers(
+                    self, message, headers
+                )
+
+                # Convert wrapped message back to bytes if needed
+                if isinstance(wrapped_message, dict):
+                    message_to_send = json.dumps(wrapped_message).encode("utf-8")
+                else:
+                    message_to_send = wrapped_message
+
+                original_args[message_arg_index] = message_to_send
+                args = tuple(original_args)
+
+            return await original_request(self, *args, **kwargs)
+
+        nats.NATS.request = instrumented_request
+
+        # Instrument `subscribe` method
+        original_subscribe = nats.NATS.subscribe
+
+        @functools.wraps(original_subscribe)
+        async def instrumented_subscribe(self, subject, cb=None, *args, **kwargs):
+            # Wrap the callback to add tracing spans for message handling
+            if (
+                cb is not None
+                and _global_tracer
+                and not getattr(cb, "_is_instrumented", False)
+            ):
+                user_cb = cb  # SAVE the original callback
+
+                @functools.wraps(user_cb)
+                async def traced_callback(msg):
+                    try:
+                        message_dict = json.loads(msg.data.decode())
+                        headers = message_dict.get("headers", {})
+
+                        # Extract traceparent and session info from headers
+                        traceparent = headers.get("traceparent")
+                        session_id = headers.get("session_id")
+
+                        # Create carrier for context propagation
+                        carrier = {}
+                        for key in ["traceparent", "Traceparent", "baggage", "Baggage"]:
+                            if key.lower() in [k.lower() for k in headers.keys()]:
+                                for k in headers.keys():
+                                    if k.lower() == key.lower():
+                                        carrier[key.lower()] = headers[k]
+
+                        # Restore trace context
+                        ctx = None
+                        if carrier and traceparent:
+                            ctx = TraceContextTextMapPropagator().extract(
+                                carrier=carrier
+                            )
+                            ctx = W3CBaggagePropagator().extract(
+                                carrier=carrier, context=ctx
+                            )
+
+                            # Activate the restored context
+                            token = context.attach(ctx)
+
+                            try:
+                                # Set execution ID with the restored context
+                                if session_id and session_id != "None":
+                                    set_session_id(session_id, traceparent=traceparent)
+
+                                    # Store in kv_store with thread safety
+                                    with _kv_lock:
+                                        kv_store.set(
+                                            f"execution.{traceparent}", session_id
+                                        )
+
+                                # DON'T detach the context yet - we need it to persist for the callback
+                                # The context will be cleaned up later or by the garbage collector
+
+                            except Exception as e:
+                                # Only detach on error
+                                context.detach(token)
+                                raise e
+                        elif traceparent and session_id and session_id != "None":
+                            # Even without carrier context, set session ID if we have the data
+                            set_session_id(session_id, traceparent=traceparent)
+
+                        # Fallback: check stored execution ID if not found in headers
+                        if traceparent and (not session_id or session_id == "None"):
+                            with _kv_lock:
+                                stored_session_id = kv_store.get(
+                                    f"execution.{traceparent}"
+                                )
+                                if stored_session_id:
+                                    session_id = stored_session_id
+                                    set_session_id(session_id, traceparent=traceparent)
+
+                        # Process and clean the message
+                        message_to_return = message_dict.copy()
+                        if "headers" in message_to_return:
+                            headers_copy = message_to_return["headers"].copy()
+                            # Remove tracing-specific headers but keep other headers
+                            headers_copy.pop("traceparent", None)
+                            headers_copy.pop("session_id", None)
+                            if headers_copy:
+                                message_to_return["headers"] = headers_copy
+                            else:
+                                message_to_return.pop("headers", None)
+
+                        # Return processed message, update msg.data
+                        if isinstance(message_to_return, str):
+                            msg.data = message_to_return.encode("utf-8")
+                        else:
+                            msg.data = json.dumps(message_to_return).encode("utf-8")
+
+                        # Now call the original user callback with the modified msg
+                        ctx = {} if ctx is None else ctx
+                        if _global_tracer:
+                            with _global_tracer.start_as_current_span(
+                                "nats.subscribe.callback", context=ctx
+                            ) as span:
+                                span.set_attribute("nats.subject", subject)
+                                span.set_attribute("nats.session_id", session_id)
+                                await user_cb(msg)
+                        else:
+                            await user_cb(msg)
+                    except Exception as e:
+                        print(f"Error processing message in traced_callback: {e}")
+                        await user_cb(msg)  # Call original callback even on error
+
+                traced_callback._is_instrumented = True  # mark as instrumented
+                cb = traced_callback
+
+            return await original_subscribe(self, subject, cb=cb, *args, **kwargs)
+
+        nats.NATS.subscribe = instrumented_subscribe
+
+    def _wrap_message_with_headers(self, message, headers):
+        """Helper method to wrap messages with headers consistently"""
+        if isinstance(message, bytes):
+            try:
+                decoded_message = message.decode("utf-8")
+                try:
+                    original_message = json.loads(decoded_message)
+                    if isinstance(original_message, dict):
+                        wrapped_message = original_message.copy()
+                        existing_headers = wrapped_message.get("headers", {})
+                        existing_headers.update(headers)
+                        wrapped_message["headers"] = existing_headers
+                    else:
+                        wrapped_message = {
+                            "headers": headers,
+                            "payload": original_message,
+                        }
+                except json.JSONDecodeError:
+                    wrapped_message = {"headers": headers, "payload": decoded_message}
+            except UnicodeDecodeError:
+                # Fix type annotation issue by ensuring message is bytes
+                encoded_message = (
+                    message if isinstance(message, bytes) else message.encode("utf-8")
+                )
+                wrapped_message = {
+                    "headers": headers,
+                    "payload": base64.b64encode(encoded_message).decode("utf-8"),
+                }
+        elif isinstance(message, str):
+            try:
+                original_message = json.loads(message)
+                if isinstance(original_message, dict):
+                    wrapped_message = original_message.copy()
+                    existing_headers = wrapped_message.get("headers", {})
+                    existing_headers.update(headers)
+                    wrapped_message["headers"] = existing_headers
+                else:
+                    wrapped_message = {"headers": headers, "payload": original_message}
+            except json.JSONDecodeError:
+                wrapped_message = {"headers": headers, "payload": message}
+        elif isinstance(message, dict):
+            wrapped_message = message.copy()
+            existing_headers = wrapped_message.get("headers", {})
+            existing_headers.update(headers)
+            wrapped_message["headers"] = existing_headers
+        else:
+            wrapped_message = {"headers": headers, "payload": json.dumps(message)}
+
+        return wrapped_message
+
+    def _uninstrument(self, **kwargs):
+        try:
+            import nats
+        except ImportError:
+            raise ImportError("No module named 'nats'. Please install it first.")
+
+        # Restore the original methods
+        methods_to_restore = [
+            "publish",
+            "request",
+            "subscribe",
+        ]
+
+        for method_name in methods_to_restore:
+            if hasattr(nats.NATS, method_name):
+                original_method = getattr(nats.NATS, method_name)
+                if hasattr(original_method, "__wrapped__"):
+                    setattr(nats.NATS, method_name, original_method.__wrapped__)

ioa_observe/sdk/instrumentations/slim.py

@@ -16,7 +16,7 @@ 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
 
-_instruments = ("slim-bindings >= 0.4",)
+_instruments = ("slim-bindings >= 0.4.0",)
 _global_tracer = None
 _kv_lock = threading.RLock()  # Add thread-safety for kv_store operations
 
@@ -39,7 +39,9 @@ class SLIMInstrumentor(BaseInstrumentor):
             )
 
         # Instrument `publish` method - handles multiple signatures
+        # In v0.6.0+, publish moved from Slim class to Session objects
         if hasattr(slim_bindings.Slim, "publish"):
+            # Legacy v0.5.x app-level publish method
             original_publish = slim_bindings.Slim.publish
 
             @functools.wraps(original_publish)
@@ -166,7 +168,7 @@ class SLIMInstrumentor(BaseInstrumentor):
 
             slim_bindings.Slim.publish_to = instrumented_publish_to
 
-        # Instrument `request_reply` (new v0.4.0+ method)
+        # Instrument `request_reply` (v0.4.0+ to v0.5.x method, removed in v0.6.0)
         if hasattr(slim_bindings.Slim, "request_reply"):
             original_request_reply = slim_bindings.Slim.request_reply
 
@@ -290,35 +292,261 @@ class SLIMInstrumentor(BaseInstrumentor):
 
             slim_bindings.Slim.set_route = instrumented_set_route
 
-        # Instrument `receive`
-        original_receive = slim_bindings.Slim.receive
-
-        @functools.wraps(original_receive)
-        async def instrumented_receive(
-            self, session=None, timeout=None, *args, **kwargs
-        ):
-            # Handle both old and new API patterns
-            if session is not None or timeout is not None:
-                # New API pattern with session parameter
-                kwargs_with_params = kwargs.copy()
-                if session is not None:
-                    kwargs_with_params["session"] = session
-                if timeout is not None:
-                    kwargs_with_params["timeout"] = timeout
-                recv_session, raw_message = await original_receive(
-                    self, **kwargs_with_params
-                )
+        # Instrument `receive` - only if it exists (removed in v0.6.0)
+        if hasattr(slim_bindings.Slim, "receive"):
+            original_receive = slim_bindings.Slim.receive
+
+            @functools.wraps(original_receive)
+            async def instrumented_receive(
+                self, session=None, timeout=None, *args, **kwargs
+            ):
+                # Handle both old and new API patterns
+                if session is not None or timeout is not None:
+                    # New API pattern with session parameter
+                    kwargs_with_params = kwargs.copy()
+                    if session is not None:
+                        kwargs_with_params["session"] = session
+                    if timeout is not None:
+                        kwargs_with_params["timeout"] = timeout
+                    recv_session, raw_message = await original_receive(
+                        self, **kwargs_with_params
+                    )
+                else:
+                    # Legacy API pattern
+                    recv_session, raw_message = await original_receive(
+                        self, *args, **kwargs
+                    )
+
+                if raw_message is None:
+                    return recv_session, raw_message
+
+                try:
+                    message_dict = json.loads(raw_message.decode())
+                    headers = message_dict.get("headers", {})
+
+                    # Extract traceparent and session info from headers
+                    traceparent = headers.get("traceparent")
+                    session_id = headers.get("session_id")
+
+                    # Create carrier for context propagation
+                    carrier = {}
+                    for key in ["traceparent", "Traceparent", "baggage", "Baggage"]:
+                        if key.lower() in [k.lower() for k in headers.keys()]:
+                            for k in headers.keys():
+                                if k.lower() == key.lower():
+                                    carrier[key.lower()] = headers[k]
+
+                    # Restore trace context
+                    if carrier and traceparent:
+                        ctx = TraceContextTextMapPropagator().extract(carrier=carrier)
+                        ctx = W3CBaggagePropagator().extract(
+                            carrier=carrier, context=ctx
+                        )
+
+                        # Activate the restored context
+                        token = context.attach(ctx)
+
+                        try:
+                            # Set execution ID with the restored context
+                            if session_id and session_id != "None":
+                                set_session_id(session_id, traceparent=traceparent)
+
+                                # Store in kv_store with thread safety
+                                with _kv_lock:
+                                    kv_store.set(f"execution.{traceparent}", session_id)
+
+                            # DON'T detach the context yet - we need it to persist for the callback
+                            # The context will be cleaned up later or by the garbage collector
+
+                        except Exception as e:
+                            # Only detach on error
+                            context.detach(token)
+                            raise e
+                    elif traceparent and session_id and session_id != "None":
+                        # Even without carrier context, set session ID if we have the data
+                        set_session_id(session_id, traceparent=traceparent)
+
+                    # Fallback: check stored execution ID if not found in headers
+                    if traceparent and (not session_id or session_id == "None"):
+                        with _kv_lock:
+                            stored_session_id = kv_store.get(f"execution.{traceparent}")
+                            if stored_session_id:
+                                session_id = stored_session_id
+                                set_session_id(session_id, traceparent=traceparent)
+
+                    # Process and clean the message
+                    message_to_return = message_dict.copy()
+                    if "headers" in message_to_return:
+                        headers_copy = message_to_return["headers"].copy()
+                        # Remove tracing-specific headers but keep other headers
+                        headers_copy.pop("traceparent", None)
+                        headers_copy.pop("session_id", None)
+                        headers_copy.pop("slim_session_id", None)
+                        if headers_copy:
+                            message_to_return["headers"] = headers_copy
+                        else:
+                            message_to_return.pop("headers", None)
+
+                    # Return processed message
+                    if len(message_to_return) == 1 and "payload" in message_to_return:
+                        payload = message_to_return["payload"]
+                        if isinstance(payload, str):
+                            try:
+                                payload_dict = json.loads(payload)
+                                return recv_session, json.dumps(payload_dict).encode(
+                                    "utf-8"
+                                )
+                            except json.JSONDecodeError:
+                                return recv_session, payload.encode(
+                                    "utf-8"
+                                ) if isinstance(payload, str) else payload
+                        return recv_session, json.dumps(payload).encode(
+                            "utf-8"
+                        ) if isinstance(payload, (dict, list)) else payload
+                    else:
+                        return recv_session, json.dumps(message_to_return).encode(
+                            "utf-8"
+                        )
+
+                except Exception as e:
+                    print(f"Error processing message: {e}")
+                    return recv_session, raw_message
+
+            slim_bindings.Slim.receive = instrumented_receive
+
+        # Instrument `connect` - only if it exists
+        if hasattr(slim_bindings.Slim, "connect"):
+            original_connect = slim_bindings.Slim.connect
+
+            @functools.wraps(original_connect)
+            async def instrumented_connect(self, *args, **kwargs):
+                if _global_tracer:
+                    with _global_tracer.start_as_current_span("slim.connect"):
+                        return await original_connect(self, *args, **kwargs)
+                else:
+                    return await original_connect(self, *args, **kwargs)
+
+            slim_bindings.Slim.connect = instrumented_connect
+
+        # Instrument `create_session` (new v0.4.0+ method)
+        if hasattr(slim_bindings.Slim, "create_session"):
+            original_create_session = slim_bindings.Slim.create_session
+
+            @functools.wraps(original_create_session)
+            async def instrumented_create_session(self, config, *args, **kwargs):
+                if _global_tracer:
+                    with _global_tracer.start_as_current_span(
+                        "slim.create_session"
+                    ) as span:
+                        session_info = await original_create_session(
+                            self, config, *args, **kwargs
+                        )
+
+                        # Add session attributes to span
+                        if hasattr(session_info, "id"):
+                            span.set_attribute("slim.session.id", str(session_info.id))
+
+                        return session_info
+                else:
+                    return await original_create_session(self, config, *args, **kwargs)
+
+            slim_bindings.Slim.create_session = instrumented_create_session
+
+        # Instrument new v0.6.0+ session-level methods
+        # These methods are available on Session objects, not the Slim app
+        self._instrument_session_methods(slim_bindings)
+
+        # Instrument new v0.6.0+ app-level methods
+        # listen_for_session replaces app.receive() for new sessions in v0.6.0+
+        if hasattr(slim_bindings.Slim, "listen_for_session"):
+            original_listen_for_session = slim_bindings.Slim.listen_for_session
+
+            @functools.wraps(original_listen_for_session)
+            async def instrumented_listen_for_session(self, *args, **kwargs):
+                if _global_tracer:
+                    with _global_tracer.start_as_current_span(
+                        "slim.listen_for_session"
+                    ):
+                        session = await original_listen_for_session(
+                            self, *args, **kwargs
+                        )
+
+                        return session
+                else:
+                    return await original_listen_for_session(self, *args, **kwargs)
+
+            slim_bindings.Slim.listen_for_session = instrumented_listen_for_session
+
+    def _instrument_session_methods(self, slim_bindings):
+        # In v0.6.0+, we need to instrument session classes dynamically
+        # Try to find session-related classes in the slim_bindings module
+        session_classes = []
+
+        # Look for common session class names
+        for attr_name in ["Session", "P2PSession", "GroupSession"]:
+            if hasattr(slim_bindings, attr_name):
+                session_class = getattr(slim_bindings, attr_name)
+                session_classes.append((attr_name, session_class))
+
+        # Also look for any class that has session-like methods
+        for attr_name in dir(slim_bindings):
+            attr = getattr(slim_bindings, attr_name)
+            if isinstance(attr, type) and (
+                hasattr(attr, "get_message") or hasattr(attr, "publish")
+            ):
+                session_classes.append((attr_name, attr))
+
+        # Instrument session methods for found classes
+        for _, session_class in session_classes:
+            # Instrument get_message (v0.6.0+ replacement for receive)
+            if hasattr(session_class, "get_message"):
+                self._instrument_session_get_message(session_class)
+
+            # Instrument session publish methods
+            if hasattr(session_class, "publish"):
+                self._instrument_session_publish(session_class, "publish")
+
+            if hasattr(session_class, "publish_to"):
+                self._instrument_session_publish(session_class, "publish_to")
+
+    def _instrument_session_get_message(self, session_class):
+        """Instrument session.get_message method (v0.6.0+)"""
+        original_get_message = session_class.get_message
+
+        @functools.wraps(original_get_message)
+        async def instrumented_get_message(self, timeout=None, *args, **kwargs):
+            # Handle the message reception similar to the old receive method
+            if timeout is not None:
+                kwargs["timeout"] = timeout
+
+            result = await original_get_message(self, **kwargs)
+
+            # Handle different return types from get_message
+            if result is None:
+                return result
+
+            # Check if get_message returns a tuple (context, message) or just message
+            if isinstance(result, tuple) and len(result) == 2:
+                message_context, raw_message = result
             else:
-                # Legacy API pattern
-                recv_session, raw_message = await original_receive(
-                    self, *args, **kwargs
-                )
+                raw_message = result
+                message_context = None
 
             if raw_message is None:
-                return recv_session, raw_message
+                return result
 
             try:
-                message_dict = json.loads(raw_message.decode())
+                # Handle different message types
+                if isinstance(raw_message, bytes):
+                    message_dict = json.loads(raw_message.decode())
+                elif isinstance(raw_message, str):
+                    message_dict = json.loads(raw_message)
+                elif isinstance(raw_message, dict):
+                    message_dict = raw_message
+                else:
+                    # Unknown type, return as-is
+                    return result
+
                 headers = message_dict.get("headers", {})
 
                 # Extract traceparent and session info from headers
@@ -351,7 +579,6 @@ class SLIMInstrumentor(BaseInstrumentor):
                                 kv_store.set(f"execution.{traceparent}", session_id)
 
                         # DON'T detach the context yet - we need it to persist for the callback
-                        # The context will be cleaned up later or by the garbage collector
 
                     except Exception as e:
                         # Only detach on error
@@ -382,68 +609,257 @@ class SLIMInstrumentor(BaseInstrumentor):
                     else:
                         message_to_return.pop("headers", None)
 
-                # Return processed message
+                # Return processed message, maintaining original return format
                 if len(message_to_return) == 1 and "payload" in message_to_return:
                     payload = message_to_return["payload"]
                     if isinstance(payload, str):
                         try:
                             payload_dict = json.loads(payload)
-                            return recv_session, json.dumps(payload_dict).encode(
-                                "utf-8"
-                            )
+                            processed_message = json.dumps(payload_dict).encode("utf-8")
                         except json.JSONDecodeError:
-                            return recv_session, payload.encode("utf-8") if isinstance(
-                                payload, str
-                            ) else payload
-                    return recv_session, json.dumps(payload).encode(
-                        "utf-8"
-                    ) if isinstance(payload, (dict, list)) else payload
+                            processed_message = (
+                                payload.encode("utf-8")
+                                if isinstance(payload, str)
+                                else payload
+                            )
+                    else:
+                        processed_message = (
+                            json.dumps(payload).encode("utf-8")
+                            if isinstance(payload, (dict, list))
+                            else payload
+                        )
                 else:
-                    return recv_session, json.dumps(message_to_return).encode("utf-8")
+                    processed_message = json.dumps(message_to_return).encode("utf-8")
+
+                # Return in the same format as received
+                if isinstance(result, tuple) and len(result) == 2:
+                    return (message_context, processed_message)
+                else:
+                    return processed_message
 
             except Exception as e:
                 print(f"Error processing message: {e}")
-                return recv_session, raw_message
+                return result
 
-        slim_bindings.Slim.receive = instrumented_receive
+        session_class.get_message = instrumented_get_message
 
-        # Instrument `connect`
-        original_connect = slim_bindings.Slim.connect
+    def _instrument_session_publish(self, session_class, method_name):
+        """Instrument session publish methods"""
+        original_method = getattr(session_class, method_name)
 
-        @functools.wraps(original_connect)
-        async def instrumented_connect(self, *args, **kwargs):
+        @functools.wraps(original_method)
+        async def instrumented_session_publish(self, *args, **kwargs):
             if _global_tracer:
-                with _global_tracer.start_as_current_span("slim.connect"):
-                    return await original_connect(self, *args, **kwargs)
+                with _global_tracer.start_as_current_span(
+                    f"session.{method_name}"
+                ) as span:
+                    traceparent = get_current_traceparent()
+
+                    # Add session context to span
+                    if hasattr(self, "id"):
+                        span.set_attribute("slim.session.id", str(self.id))
+
+                    # Handle message wrapping
+                    if args:
+                        # Thread-safe access to kv_store
+                        session_id = None
+                        if traceparent:
+                            with _kv_lock:
+                                session_id = kv_store.get(f"execution.{traceparent}")
+                                if session_id:
+                                    kv_store.set(f"execution.{traceparent}", session_id)
+
+                        headers = {
+                            "session_id": session_id if session_id else None,
+                            "traceparent": traceparent,
+                            "slim_session_id": str(self.id)
+                            if hasattr(self, "id")
+                            else None,
+                        }
+
+                        # Set baggage context
+                        if traceparent and session_id:
+                            baggage.set_baggage(f"execution.{traceparent}", session_id)
+
+                        # Wrap the message (first argument for publish, second for publish_to)
+                        message_idx = 1 if method_name == "publish_to" else 0
+                        if len(args) > message_idx:
+                            args_list = list(args)
+                            message = args_list[message_idx]
+                            wrapped_message = (
+                                SLIMInstrumentor._wrap_message_with_headers(
+                                    None, message, headers
+                                )
+                            )
+
+                            # Convert wrapped message back to bytes if needed
+                            if isinstance(wrapped_message, dict):
+                                message_to_send = json.dumps(wrapped_message).encode(
+                                    "utf-8"
+                                )
+                            else:
+                                message_to_send = wrapped_message
+
+                            args_list[message_idx] = message_to_send
+                            args = tuple(args_list)
+
+                    return await original_method(self, *args, **kwargs)
             else:
-                return await original_connect(self, *args, **kwargs)
+                # Handle message wrapping even without tracing
+                if args:
+                    traceparent = get_current_traceparent()
+                    session_id = None
+                    if traceparent:
+                        with _kv_lock:
+                            session_id = kv_store.get(f"execution.{traceparent}")
+
+                    if traceparent or session_id:
+                        headers = {
+                            "session_id": session_id if session_id else None,
+                            "traceparent": traceparent,
+                            "slim_session_id": str(self.id)
+                            if hasattr(self, "id")
+                            else None,
+                        }
 
-        slim_bindings.Slim.connect = instrumented_connect
+                        # Wrap the message (first argument for publish, second for publish_to)
+                        message_idx = 1 if method_name == "publish_to" else 0
+                        if len(args) > message_idx:
+                            args_list = list(args)
+                            message = args_list[message_idx]
+                            wrapped_message = (
+                                SLIMInstrumentor._wrap_message_with_headers(
+                                    None, message, headers
+                                )
+                            )
 
-        # Instrument `create_session` (new v0.4.0+ method)
-        if hasattr(slim_bindings.Slim, "create_session"):
-            original_create_session = slim_bindings.Slim.create_session
+                            if isinstance(wrapped_message, dict):
+                                message_to_send = json.dumps(wrapped_message).encode(
+                                    "utf-8"
+                                )
+                            else:
+                                message_to_send = wrapped_message
 
-            @functools.wraps(original_create_session)
-            async def instrumented_create_session(self, config, *args, **kwargs):
-                if _global_tracer:
-                    with _global_tracer.start_as_current_span(
-                        "slim.create_session"
-                    ) as span:
-                        session_info = await original_create_session(
-                            self, config, *args, **kwargs
-                        )
+                            args_list[message_idx] = message_to_send
+                            args = tuple(args_list)
 
-                        # Add session attributes to span
-                        if hasattr(session_info, "id"):
-                            span.set_attribute("slim.session.id", str(session_info.id))
+                return await original_method(self, *args, **kwargs)
 
-                        return session_info
-                else:
-                    return await original_create_session(self, config, *args, **kwargs)
+        setattr(session_class, method_name, instrumented_session_publish)
 
-            slim_bindings.Slim.create_session = instrumented_create_session
+    def _instrument_session_method_if_exists(self, slim_bindings, method_name):
+        """Helper to instrument a session method if it exists"""
+
+        # Look for session classes that might have this method
+        for attr_name in dir(slim_bindings):
+            attr = getattr(slim_bindings, attr_name)
+            if hasattr(attr, method_name):
+                original_method = getattr(attr, method_name)
+
+                if callable(original_method):
+                    instrumented_method = self._create_session_method_wrapper(
+                        method_name, original_method
+                    )
+                    setattr(attr, method_name, instrumented_method)
+
+    def _create_session_method_wrapper(self, method_name, original_method):
+        """Create an instrumented wrapper for session methods"""
+
+        @functools.wraps(original_method)
+        async def instrumented_session_method(self, *args, **kwargs):
+            if _global_tracer:
+                with _global_tracer.start_as_current_span(f"session.{method_name}"):
+                    traceparent = get_current_traceparent()
+
+                    # Handle message wrapping for publish methods
+                    if method_name in ["publish", "publish_to"] and args:
+                        # Thread-safe access to kv_store
+                        session_id = None
+                        if traceparent:
+                            with _kv_lock:
+                                session_id = kv_store.get(f"execution.{traceparent}")
+                                if session_id:
+                                    kv_store.set(f"execution.{traceparent}", session_id)
+
+                        headers = {
+                            "session_id": session_id if session_id else None,
+                            "traceparent": traceparent,
+                            "slim_session_id": str(self.id)
+                            if hasattr(self, "id")
+                            else None,
+                        }
+
+                        # Set baggage context
+                        if traceparent and session_id:
+                            baggage.set_baggage(f"execution.{traceparent}", session_id)
+
+                        # Wrap the message (first argument for publish, second for publish_to)
+                        message_idx = 1 if method_name == "publish_to" else 0
+                        if len(args) > message_idx:
+                            args_list = list(args)
+                            message = args_list[message_idx]
+                            wrapped_message = SLIMInstrumentor._wrap_message_with_headers(
+                                None,
+                                message,
+                                headers,  # Pass None for self since it's a static method
+                            )
+
+                            # Convert wrapped message back to bytes if needed
+                            if isinstance(wrapped_message, dict):
+                                message_to_send = json.dumps(wrapped_message).encode(
+                                    "utf-8"
+                                )
+                            else:
+                                message_to_send = wrapped_message
+
+                            args_list[message_idx] = message_to_send
+                            args = tuple(args_list)
 
+                    return await original_method(self, *args, **kwargs)
+            else:
+                # Handle message wrapping even without tracing
+                if method_name in ["publish", "publish_to"] and args:
+                    traceparent = get_current_traceparent()
+                    session_id = None
+                    if traceparent:
+                        with _kv_lock:
+                            session_id = kv_store.get(f"execution.{traceparent}")
+
+                    if traceparent or session_id:
+                        headers = {
+                            "session_id": session_id if session_id else None,
+                            "traceparent": traceparent,
+                            "slim_session_id": str(self.id)
+                            if hasattr(self, "id")
+                            else None,
+                        }
+
+                        # Wrap the message
+                        message_idx = 1 if method_name == "publish_to" else 0
+                        if len(args) > message_idx:
+                            args_list = list(args)
+                            message = args_list[message_idx]
+                            wrapped_message = (
+                                SLIMInstrumentor._wrap_message_with_headers(
+                                    None, message, headers
+                                )
+                            )
+
+                            if isinstance(wrapped_message, dict):
+                                message_to_send = json.dumps(wrapped_message).encode(
+                                    "utf-8"
+                                )
+                            else:
+                                message_to_send = wrapped_message
+
+                            args_list[message_idx] = message_to_send
+                            args = tuple(args_list)
+
+                return await original_method(self, *args, **kwargs)
+
+        return instrumented_session_method
+
+    @staticmethod
     def _wrap_message_with_headers(self, message, headers):
         """Helper method to wrap messages with headers consistently"""
         if isinstance(message, bytes):
@@ -506,12 +922,13 @@ class SLIMInstrumentor(BaseInstrumentor):
         methods_to_restore = [
             "publish",
             "publish_to",
-            "request_reply",
+            "request_reply",  # v0.4.0-v0.5.x only
             "receive",
             "connect",
             "create_session",
             "invite",
             "set_route",
+            "listen_for_session",  # v0.6.0+
         ]
 
         for method_name in methods_to_restore:
@@ -521,3 +938,21 @@ class SLIMInstrumentor(BaseInstrumentor):
                     setattr(
                         slim_bindings.Slim, method_name, original_method.__wrapped__
                     )
+
+        # Also try to restore session-level methods (v0.6.0+)
+        # This is best-effort since session classes may vary
+        session_methods_to_restore = [
+            "publish",
+            "publish_to",
+            "get_message",
+            "invite",
+            "remove",
+        ]
+
+        for attr_name in dir(slim_bindings):
+            attr = getattr(slim_bindings, attr_name)
+            for method_name in session_methods_to_restore:
+                if hasattr(attr, method_name):
+                    original_method = getattr(attr, method_name)
+                    if hasattr(original_method, "__wrapped__"):
+                        setattr(attr, method_name, original_method.__wrapped__)

{ioa_observe_sdk-1.0.20.dist-info → ioa_observe_sdk-1.0.22.dist-info}/METADATA

@@ -1,20 +1,21 @@
 Metadata-Version: 2.4
 Name: ioa-observe-sdk
-Version: 1.0.20
+Version: 1.0.22
 Summary: IOA Observability SDK
+License-Expression: Apache-2.0
 Requires-Python: >=3.10
 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.33.1
+Requires-Dist: opentelemetry-api==1.37.0
 Requires-Dist: opentelemetry-distro
-Requires-Dist: opentelemetry-exporter-otlp==1.33.1
-Requires-Dist: opentelemetry-exporter-otlp-proto-common==1.33.1
-Requires-Dist: opentelemetry-exporter-otlp-proto-grpc==1.33.1
-Requires-Dist: opentelemetry-exporter-otlp-proto-http==1.33.1
+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-instrumentation
-Requires-Dist: opentelemetry-instrumentation-logging==0.54b1
+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
@@ -26,18 +27,18 @@ 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.54b1
+Requires-Dist: opentelemetry-instrumentation-requests==0.58b0
 Requires-Dist: opentelemetry-instrumentation-sagemaker==0.43.1
-Requires-Dist: opentelemetry-instrumentation-threading==0.54b1
+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.54b1
+Requires-Dist: opentelemetry-instrumentation-urllib3==0.58b0
 Requires-Dist: opentelemetry-instrumentation-vertexai==0.43.1
-Requires-Dist: opentelemetry-proto==1.33.1
-Requires-Dist: opentelemetry-sdk==1.33.1
-Requires-Dist: opentelemetry-semantic-conventions==0.54b1
+Requires-Dist: opentelemetry-proto==1.37.0
+Requires-Dist: opentelemetry-sdk==1.37.0
+Requires-Dist: opentelemetry-semantic-conventions==0.58b0
 Requires-Dist: opentelemetry-semantic-conventions-ai>=0.4.11
-Requires-Dist: opentelemetry-util-http==0.54b1
+Requires-Dist: opentelemetry-util-http==0.58b0
 Requires-Dist: langgraph>=0.3.2
 Requires-Dist: langchain>=0.3.19
 Requires-Dist: langchain-openai>=0.3.8

{ioa_observe_sdk-1.0.20.dist-info → ioa_observe_sdk-1.0.22.dist-info}/RECORD

@@ -16,7 +16,8 @@ ioa_observe/sdk/decorators/util.py,sha256=IebvH9gwZN1en3LblYJUh4bAV2STl6xmp8WpZz
 ioa_observe/sdk/instrumentations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ioa_observe/sdk/instrumentations/a2a.py,sha256=ZpqvPl4u-yheQzSdBfxnZhWFZ8ntbKni_uaW3IDyjqw,6309
 ioa_observe/sdk/instrumentations/mcp.py,sha256=vRM3ofnn7AMmry2RrfyZnZVPEutLWiDMghx2TSnm0Wk,18569
-ioa_observe/sdk/instrumentations/slim.py,sha256=WIe0E2y6eg5heQ1EA6l9BUDniYpYIjWQ24Pcv_8ZfPk,23147
+ioa_observe/sdk/instrumentations/nats.py,sha256=UOp2AJlm1JkYkwF3xzU_izzohQVQkByjL-AX4n_JRfo,14476
+ioa_observe/sdk/instrumentations/slim.py,sha256=r_vIuYUo-IhSRuYWymqfS26swDZ8_CFDH6p_QbduLKU,43757
 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
@@ -41,8 +42,8 @@ ioa_observe/sdk/utils/const.py,sha256=d67dUTAH9UpWvUV9GLBUqn1Sc2knJ55dy-e6YoLrvS
 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.20.dist-info/licenses/LICENSE.md,sha256=55VjUfgjWOS4vv3Cf55gfq-RxjPgRIO2vlgYPUuC5lA,11362
-ioa_observe_sdk-1.0.20.dist-info/METADATA,sha256=rq00vlT-d5xj5wwvTp0j8LeERhmSX1vC90mDkOi9tKI,7027
-ioa_observe_sdk-1.0.20.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ioa_observe_sdk-1.0.20.dist-info/top_level.txt,sha256=Yt-6Y1olZEDqCs2REeqI30WjYx0pLGQSVqzYmDd67N8,12
-ioa_observe_sdk-1.0.20.dist-info/RECORD,,
+ioa_observe_sdk-1.0.22.dist-info/licenses/LICENSE.md,sha256=55VjUfgjWOS4vv3Cf55gfq-RxjPgRIO2vlgYPUuC5lA,11362
+ioa_observe_sdk-1.0.22.dist-info/METADATA,sha256=Uces4XbDjeIZ2okZImfwPXNHUtyoO7od75dHthOs56Q,7058
+ioa_observe_sdk-1.0.22.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ioa_observe_sdk-1.0.22.dist-info/top_level.txt,sha256=Yt-6Y1olZEDqCs2REeqI30WjYx0pLGQSVqzYmDd67N8,12
+ioa_observe_sdk-1.0.22.dist-info/RECORD,,