Merge branch 'master' into potel-base

Author: antonpirker

sentry_sdk/integrations/langchain.py

@@ -1,3 +1,4 @@
+import itertools
 from collections import OrderedDict
 from functools import wraps
 
@@ -87,12 +88,9 @@ def __init__(self, span):
 class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
     """Base callback handler that can be used to handle callbacks from langchain."""
 
-    span_map = OrderedDict()  # type: OrderedDict[UUID, WatchedSpan]
-
-    max_span_map_size = 0
-
     def __init__(self, max_span_map_size, include_prompts, tiktoken_encoding_name=None):
         # type: (int, bool, Optional[str]) -> None
+        self.span_map = OrderedDict()  # type: OrderedDict[UUID, WatchedSpan]
         self.max_span_map_size = max_span_map_size
         self.include_prompts = include_prompts
 
@@ -458,7 +456,14 @@ def new_configure(
                 **kwargs,
             )
 
-        if not any(isinstance(cb, SentryLangchainCallback) for cb in callbacks_list):
+        inheritable_callbacks_list = (
+            inheritable_callbacks if isinstance(inheritable_callbacks, list) else []
+        )
+
+        if not any(
+            isinstance(cb, SentryLangchainCallback)
+            for cb in itertools.chain(callbacks_list, inheritable_callbacks_list)
+        ):
             # Avoid mutating the existing callbacks list
             callbacks_list = [
                 *callbacks_list,

sentry_sdk/integrations/ray.py

@@ -44,77 +44,85 @@ def _patch_ray_remote():
     old_remote = ray.remote
 
     @functools.wraps(old_remote)
-    def new_remote(f, *args, **kwargs):
-        # type: (Callable[..., Any], *Any, **Any) -> Callable[..., Any]
+    def new_remote(f=None, *args, **kwargs):
+        # type: (Optional[Callable[..., Any]], *Any, **Any) -> Callable[..., Any]
+
         if inspect.isclass(f):
             # Ray Actors
             # (https://docs.ray.io/en/latest/ray-core/actors.html)
             # are not supported
             # (Only Ray Tasks are supported)
-            return old_remote(f, *args, *kwargs)
-
-        def _f(*f_args, _tracing=None, **f_kwargs):
-            # type: (Any, Optional[dict[str, Any]],  Any) -> Any
-            """
-            Ray Worker
-            """
-            _check_sentry_initialized()
-
-            root_span_name = qualname_from_function(f) or DEFAULT_TRANSACTION_NAME
-            sentry_sdk.get_current_scope().set_transaction_name(
-                root_span_name,
-                source=TransactionSource.TASK,
-            )
-            with sentry_sdk.continue_trace(_tracing or {}):
+            return old_remote(f, *args, **kwargs)
+
+        def wrapper(user_f):
+            # type: (Callable[..., Any]) -> Any
+            def new_func(*f_args, _tracing=None, **f_kwargs):
+                # type: (Any, Optional[dict[str, Any]], Any) -> Any
+                _check_sentry_initialized()
+
+                root_span_name = qualname_from_function(f) or DEFAULT_TRANSACTION_NAME
+                sentry_sdk.get_current_scope().set_transaction_name(
+                    root_span_name,
+                    source=TransactionSource.TASK,
+                )
+                with sentry_sdk.continue_trace(_tracing or {}):
+                    with sentry_sdk.start_span(
+                        op=OP.QUEUE_TASK_RAY,
+                        name=qualname_from_function(user_f),
+                        origin=RayIntegration.origin,
+                        source=TransactionSource.TASK,
+                    ) as root_span:
+                        try:
+                            result = user_f(*f_args, **f_kwargs)
+                            root_span.set_status(SPANSTATUS.OK)
+                        except Exception:
+                            root_span.set_status(SPANSTATUS.INTERNAL_ERROR)
+                            exc_info = sys.exc_info()
+                            _capture_exception(exc_info)
+                            reraise(*exc_info)
+
+                        return result
+
+            if f:
+                rv = old_remote(new_func)
+            else:
+                rv = old_remote(*args, **kwargs)(new_func)
+            old_remote_method = rv.remote
+
+            def _remote_method_with_header_propagation(*args, **kwargs):
+                # type: (*Any, **Any) -> Any
+                """
+                Ray Client
+                """
                 with sentry_sdk.start_span(
-                    op=OP.QUEUE_TASK_RAY,
-                    name=root_span_name,
+                    op=OP.QUEUE_SUBMIT_RAY,
+                    name=qualname_from_function(user_f),
                     origin=RayIntegration.origin,
-                    source=TransactionSource.TASK,
-                ) as root_span:
+                    only_if_parent=True,
+                ) as span:
+                    tracing = {
+                        k: v
+                        for k, v in sentry_sdk.get_current_scope().iter_trace_propagation_headers()
+                    }
                     try:
-                        result = f(*f_args, **f_kwargs)
-                        root_span.set_status(SPANSTATUS.OK)
+                        result = old_remote_method(*args, **kwargs, _tracing=tracing)
+                        span.set_status(SPANSTATUS.OK)
                     except Exception:
-                        root_span.set_status(SPANSTATUS.INTERNAL_ERROR)
+                        span.set_status(SPANSTATUS.INTERNAL_ERROR)
                         exc_info = sys.exc_info()
                         _capture_exception(exc_info)
                         reraise(*exc_info)
 
                     return result
 
-        rv = old_remote(_f, *args, *kwargs)
-        old_remote_method = rv.remote
-
-        def _remote_method_with_header_propagation(*args, **kwargs):
-            # type: (*Any, **Any) -> Any
-            """
-            Ray Client
-            """
-            with sentry_sdk.start_span(
-                op=OP.QUEUE_SUBMIT_RAY,
-                name=qualname_from_function(f),
-                origin=RayIntegration.origin,
-                only_if_parent=True,
-            ) as span:
-                tracing = {
-                    k: v
-                    for k, v in sentry_sdk.get_current_scope().iter_trace_propagation_headers()
-                }
-                try:
-                    result = old_remote_method(*args, **kwargs, _tracing=tracing)
-                    span.set_status(SPANSTATUS.OK)
-                except Exception:
-                    span.set_status(SPANSTATUS.INTERNAL_ERROR)
-                    exc_info = sys.exc_info()
-                    _capture_exception(exc_info)
-                    reraise(*exc_info)
-
-                return result
-
-        rv.remote = _remote_method_with_header_propagation
-
-        return rv
+            rv.remote = _remote_method_with_header_propagation
+
+            return rv
+
+        if f is not None:
+            return wrapper(f)
+        else:
+            return wrapper
 
     ray.remote = new_remote
 

tests/integrations/langchain/test_langchain.py

@@ -14,10 +14,15 @@
 
 from langchain_core.callbacks import CallbackManagerForLLMRun
 from langchain_core.messages import BaseMessage, AIMessageChunk
-from langchain_core.outputs import ChatGenerationChunk
+from langchain_core.outputs import ChatGenerationChunk, ChatResult
+from langchain_core.runnables import RunnableConfig
+from langchain_core.language_models.chat_models import BaseChatModel
 
 from sentry_sdk import start_span
-from sentry_sdk.integrations.langchain import LangchainIntegration
+from sentry_sdk.integrations.langchain import (
+    LangchainIntegration,
+    SentryLangchainCallback,
+)
 from langchain.agents import tool, AgentExecutor, create_openai_tools_agent
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 
@@ -341,3 +346,84 @@ def test_span_origin(sentry_init, capture_events):
     assert event["contexts"]["trace"]["origin"] == "manual"
     for span in event["spans"]:
         assert span["origin"] == "auto.ai.langchain"
+
+
+def test_manual_callback_no_duplication(sentry_init):
+    """
+    Test that when a user manually provides a SentryLangchainCallback,
+    the integration doesn't create a duplicate callback.
+    """
+
+    # Track callback instances
+    tracked_callback_instances = set()
+
+    class CallbackTrackingModel(BaseChatModel):
+        """Mock model that tracks callback instances for testing."""
+
+        def _generate(
+            self,
+            messages,
+            stop=None,
+            run_manager=None,
+            **kwargs,
+        ):
+            # Track all SentryLangchainCallback instances
+            if run_manager:
+                for handler in run_manager.handlers:
+                    if isinstance(handler, SentryLangchainCallback):
+                        tracked_callback_instances.add(id(handler))
+
+                for handler in run_manager.inheritable_handlers:
+                    if isinstance(handler, SentryLangchainCallback):
+                        tracked_callback_instances.add(id(handler))
+
+            return ChatResult(
+                generations=[
+                    ChatGenerationChunk(message=AIMessageChunk(content="Hello!"))
+                ],
+                llm_output={},
+            )
+
+        @property
+        def _llm_type(self):
+            return "test_model"
+
+        @property
+        def _identifying_params(self):
+            return {}
+
+    sentry_init(integrations=[LangchainIntegration()])
+
+    # Create a manual SentryLangchainCallback
+    manual_callback = SentryLangchainCallback(
+        max_span_map_size=100, include_prompts=False
+    )
+
+    # Create RunnableConfig with the manual callback
+    config = RunnableConfig(callbacks=[manual_callback])
+
+    # Invoke the model with the config
+    llm = CallbackTrackingModel()
+    llm.invoke("Hello", config)
+
+    # Verify that only ONE SentryLangchainCallback instance was used
+    assert len(tracked_callback_instances) == 1, (
+        f"Expected exactly 1 SentryLangchainCallback instance, "
+        f"but found {len(tracked_callback_instances)}. "
+        f"This indicates callback duplication occurred."
+    )
+
+    # Verify the callback ID matches our manual callback
+    assert id(manual_callback) in tracked_callback_instances
+
+
+def test_span_map_is_instance_variable():
+    """Test that each SentryLangchainCallback instance has its own span_map."""
+    # Create two separate callback instances
+    callback1 = SentryLangchainCallback(max_span_map_size=100, include_prompts=True)
+    callback2 = SentryLangchainCallback(max_span_map_size=100, include_prompts=True)
+
+    # Verify they have different span_map instances
+    assert (
+        callback1.span_map is not callback2.span_map
+    ), "span_map should be an instance variable, not shared between instances"

tests/integrations/ray/test_ray.py

@@ -59,7 +59,10 @@ def read_error_from_log(job_id):
 
 
 @pytest.mark.forked
-def test_tracing_in_ray_tasks():
+@pytest.mark.parametrize(
+    "task_options", [{}, {"num_cpus": 0, "memory": 1024 * 1024 * 10}]
+)
+def test_tracing_in_ray_tasks(task_options):
     setup_sentry()
 
     ray.init(
@@ -69,14 +72,19 @@ def test_tracing_in_ray_tasks():
         }
     )
 
-    # Setup ray task
-    @ray.remote
     def example_task():
         with sentry_sdk.start_span(op="task", name="example task step"):
             ...
 
         return sentry_sdk.get_client().transport.envelopes
 
+    # Setup ray task, calling decorator directly instead of @,
+    # to accommodate for test parametrization
+    if task_options:
+        example_task = ray.remote(**task_options)(example_task)
+    else:
+        example_task = ray.remote(example_task)
+
     with sentry_sdk.start_span(op="test", name="ray client root span"):
         worker_envelopes = ray.get(example_task.remote())