ioa-observe-sdk 1.0.19__py3-none-any.whl → 1.0.21__py3-none-any.whl

ioa_observe/sdk/decorators/base.py

@@ -16,10 +16,10 @@ from ioa_observe.sdk.decorators.helpers import (
     _is_async_generator,
 )
 
-
 from langgraph.graph.state import CompiledStateGraph
 from opentelemetry import trace
 from opentelemetry import context as context_api
+from opentelemetry.context import get_value, attach, set_value
 from pydantic_core import PydanticSerializationError
 from typing_extensions import ParamSpec
 
@@ -113,7 +113,49 @@ def _setup_span(
     with get_tracer() as tracer:
         span = tracer.start_span(span_name)
         ctx = trace.set_span_in_context(span)
+
+        # Preserve existing context values before attaching new context
+        session_id = get_value("session.id")
+        current_traceparent = get_value("current_traceparent")
+        agent_id = get_value("agent_id")
+        application_id_ctx = get_value("application_id")
+        association_properties = get_value("association_properties")
+        managed_prompt = get_value("managed_prompt")
+        prompt_key = get_value("prompt_key")
+        prompt_version = get_value("prompt_version")
+        prompt_version_name = get_value("prompt_version_name")
+        prompt_version_hash = get_value("prompt_version_hash")
+        prompt_template = get_value("prompt_template")
+        prompt_template_variables = get_value("prompt_template_variables")
+
         ctx_token = context_api.attach(ctx)
+
+        # Re-attach preserved context values to the new context
+        if session_id is not None:
+            attach(set_value("session.id", session_id))
+        if current_traceparent is not None:
+            attach(set_value("current_traceparent", current_traceparent))
+        if agent_id is not None:
+            attach(set_value("agent_id", agent_id))
+        if application_id_ctx is not None:
+            attach(set_value("application_id", application_id_ctx))
+        if association_properties is not None:
+            attach(set_value("association_properties", association_properties))
+        if managed_prompt is not None:
+            attach(set_value("managed_prompt", managed_prompt))
+        if prompt_key is not None:
+            attach(set_value("prompt_key", prompt_key))
+        if prompt_version is not None:
+            attach(set_value("prompt_version", prompt_version))
+        if prompt_version_name is not None:
+            attach(set_value("prompt_version_name", prompt_version_name))
+        if prompt_version_hash is not None:
+            attach(set_value("prompt_version_hash", prompt_version_hash))
+        if prompt_template is not None:
+            attach(set_value("prompt_template", prompt_template))
+        if prompt_template_variables is not None:
+            attach(set_value("prompt_template_variables", prompt_template_variables))
+
         if tlp_span_kind == ObserveSpanKindValues.AGENT:
             with trace.get_tracer(__name__).start_span(
                 "agent_start_event", context=trace.set_span_in_context(span)
@@ -127,9 +169,6 @@ def _setup_span(
                     },
                 )
             # start_span.end()  # end the span immediately
-        # session_id = get_value("session.id")
-        # if session_id is not None:
-        #     span.set_attribute("session.id", session_id)
         if tlp_span_kind in [
             ObserveSpanKindValues.TASK,
             ObserveSpanKindValues.TOOL,

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
@@ -349,8 +577,16 @@ class SLIMInstrumentor(BaseInstrumentor):
                             # Store in kv_store with thread safety
                             with _kv_lock:
                                 kv_store.set(f"execution.{traceparent}", session_id)
-                    finally:
+
+                        # DON'T detach the context yet - we need it to persist for the callback
+
+                    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"):
@@ -373,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:
+                    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 recv_session, json.dumps(message_to_return).encode("utf-8")
+                    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):
@@ -497,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:
@@ -512,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.19.dist-info → ioa_observe_sdk-1.0.21.dist-info}/METADATA

@@ -1,20 +1,21 @@
 Metadata-Version: 2.4
 Name: ioa-observe-sdk
-Version: 1.0.19
+Version: 1.0.21
 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.19.dist-info → ioa_observe_sdk-1.0.21.dist-info}/RECORD

@@ -10,13 +10,13 @@ ioa_observe/sdk/config/__init__.py,sha256=8aVNaw0yRNLFPxlf97iOZLlJVcV81ivSDnudH2
 ioa_observe/sdk/connectors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ioa_observe/sdk/connectors/slim.py,sha256=NwbKEV7d5NIOqmG8zKqtgGigSJl7kf3QJ65z2gxpsY8,8498
 ioa_observe/sdk/decorators/__init__.py,sha256=qCpJAv98eLKs3I5EMXJVTV0s49Nc6QDSOHNh5rW5vLg,4268
-ioa_observe/sdk/decorators/base.py,sha256=eYep1xHXscdxIGisXgXwSGbkCrdpqqT_2s5BMP07ldg,30646
+ioa_observe/sdk/decorators/base.py,sha256=pnoj73UrpvaZmOPbcpGRRYCwg7LBAgO8PiBDh5ntn0c,32694
 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=ZpqvPl4u-yheQzSdBfxnZhWFZ8ntbKni_uaW3IDyjqw,6309
 ioa_observe/sdk/instrumentations/mcp.py,sha256=vRM3ofnn7AMmry2RrfyZnZVPEutLWiDMghx2TSnm0Wk,18569
-ioa_observe/sdk/instrumentations/slim.py,sha256=UM_XCuzetOzeDFQ-bFY0Q93OLffH7OsvbgNaS58E0RA,22632
+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 +41,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.19.dist-info/licenses/LICENSE.md,sha256=55VjUfgjWOS4vv3Cf55gfq-RxjPgRIO2vlgYPUuC5lA,11362
-ioa_observe_sdk-1.0.19.dist-info/METADATA,sha256=ynlTiy3cP00dX-e_PjBgTP3wB8FoT20ztUbOLeSKQEg,7027
-ioa_observe_sdk-1.0.19.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ioa_observe_sdk-1.0.19.dist-info/top_level.txt,sha256=Yt-6Y1olZEDqCs2REeqI30WjYx0pLGQSVqzYmDd67N8,12
-ioa_observe_sdk-1.0.19.dist-info/RECORD,,
+ioa_observe_sdk-1.0.21.dist-info/licenses/LICENSE.md,sha256=55VjUfgjWOS4vv3Cf55gfq-RxjPgRIO2vlgYPUuC5lA,11362
+ioa_observe_sdk-1.0.21.dist-info/METADATA,sha256=M6yTeHdBIWvp72-hLTpv8AztbhAxUwA8hVoduyRqxHU,7058
+ioa_observe_sdk-1.0.21.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ioa_observe_sdk-1.0.21.dist-info/top_level.txt,sha256=Yt-6Y1olZEDqCs2REeqI30WjYx0pLGQSVqzYmDd67N8,12
+ioa_observe_sdk-1.0.21.dist-info/RECORD,,