Cdp use fixes

pyproject.toml

@@ -6,7 +6,7 @@
 
 [project]
 name = "lmnr"
-version = "0.7.10"
+version = "0.7.11"
 description = "Python SDK for Laminar"
 authors = [
   { name = "lmnr.ai", email = "founders@lmnr.ai" }

src/lmnr/opentelemetry_lib/tracing/_instrument_initializers.py

@@ -51,9 +51,19 @@ def init_instrumentor(self, *args, **kwargs) -> BaseInstrumentor | None:
 
 
 class BrowserUseInstrumentorInitializer(InstrumentorInitializer):
-    def init_instrumentor(
-        self, client, async_client, *args, **kwargs
-    ) -> BaseInstrumentor | None:
+    """Instruments for different versions of browser-use:
+
+    - browser-use < 0.5: BrowserUseLegacyInstrumentor to track agent_step and
+      other structure spans. Session instrumentation is controlled by
+      Instruments.PLAYWRIGHT (or Instruments.PATCHRIGHT for several versions
+      in 0.4.* that used patchright)
+    - browser-use ~= 0.5: Structure spans live in browser_use package itself.
+      Session instrumentation is controlled by Instruments.PLAYWRIGHT
+    - browser-use >= 0.6.0rc1: BubusInstrumentor to keep spans structure.
+      Session instrumentation is controlled by Instruments.BROWSER_USE_SESSION
+    """
+
+    def init_instrumentor(self, *args, **kwargs) -> BaseInstrumentor | None:
         if not is_package_installed("browser-use"):
             return None
 
@@ -65,6 +75,19 @@ def init_instrumentor(
 
             return BrowserUseLegacyInstrumentor()
 
+        return None
+
+
+class BrowserUseSessionInstrumentorInitializer(InstrumentorInitializer):
+    def init_instrumentor(
+        self, client, async_client, *args, **kwargs
+    ) -> BaseInstrumentor | None:
+        if not is_package_installed("browser-use"):
+            return None
+
+        version = get_package_version("browser-use")
+        from packaging.version import parse
+
         if version and parse(version) >= parse("0.6.0rc1"):
             from lmnr.sdk.browser.browser_use_cdp_otel import BrowserUseInstrumentor
 
@@ -73,6 +96,16 @@ def init_instrumentor(
         return None
 
 
+class BubusInstrumentorInitializer(InstrumentorInitializer):
+    def init_instrumentor(self, *args, **kwargs) -> BaseInstrumentor | None:
+        if not is_package_installed("bubus"):
+            return None
+
+        from lmnr.sdk.browser.bubus_otel import BubusInstrumentor
+
+        return BubusInstrumentor()
+
+
 class ChromaInstrumentorInitializer(InstrumentorInitializer):
     def init_instrumentor(self, *args, **kwargs) -> BaseInstrumentor | None:
         if not is_package_installed("chromadb"):

src/lmnr/opentelemetry_lib/tracing/instruments.py

@@ -17,6 +17,8 @@ class Instruments(Enum):
     ANTHROPIC = "anthropic"
     BEDROCK = "bedrock"
     BROWSER_USE = "browser_use"
+    BROWSER_USE_SESSION = "browser_use_session"
+    BUBUS = "bubus"
     CHROMA = "chroma"
     COHERE = "cohere"
     CREWAI = "crewai"
@@ -60,6 +62,8 @@ class Instruments(Enum):
     Instruments.ANTHROPIC: initializers.AnthropicInstrumentorInitializer(),
     Instruments.BEDROCK: initializers.BedrockInstrumentorInitializer(),
     Instruments.BROWSER_USE: initializers.BrowserUseInstrumentorInitializer(),
+    Instruments.BROWSER_USE_SESSION: initializers.BrowserUseSessionInstrumentorInitializer(),
+    Instruments.BUBUS: initializers.BubusInstrumentorInitializer(),
     Instruments.CHROMA: initializers.ChromaInstrumentorInitializer(),
     Instruments.COHERE: initializers.CohereInstrumentorInitializer(),
     Instruments.CREWAI: initializers.CrewAIInstrumentorInitializer(),

src/lmnr/sdk/browser/browser_use_cdp_otel.py

@@ -1,3 +1,6 @@
+import asyncio
+import uuid
+
 from lmnr.sdk.client.asynchronous.async_client import AsyncLaminarClient
 from lmnr.sdk.browser.utils import with_tracer_and_client_wrapper
 from lmnr.version import __version__
@@ -12,7 +15,6 @@
 from opentelemetry.trace import get_tracer, Tracer
 from typing import Collection
 from wrapt import wrap_function_wrapper
-import uuid
 
 # Stable versions, e.g. 0.6.0, satisfy this condition too
 _instruments = ("browser-use >= 0.6.0rc1",)
@@ -33,12 +35,7 @@
 ]
 
 
-@with_tracer_and_client_wrapper
-async def _wrap(
-    tracer: Tracer, client: AsyncLaminarClient, to_wrap, wrapped, instance, args, kwargs
-):
-    result = await wrapped(*args, **kwargs)
-
+async def process_wrapped_result(result, instance, client, to_wrap):
     if to_wrap.get("action") == "inject_session_recorder":
         is_registered = await is_recorder_present(result)
         if not is_registered:
@@ -50,6 +47,14 @@ async def _wrap(
             cdp_session = await instance.get_or_create_cdp_session(target_id)
             await take_full_snapshot(cdp_session)
 
+
+@with_tracer_and_client_wrapper
+async def _wrap(
+    tracer: Tracer, client: AsyncLaminarClient, to_wrap, wrapped, instance, args, kwargs
+):
+    result = await wrapped(*args, **kwargs)
+    asyncio.create_task(process_wrapped_result(result, instance, client, to_wrap))
+
     return result
 
 

src/lmnr/sdk/browser/bubus_otel.py

@@ -0,0 +1,71 @@
+from typing import Collection
+
+from lmnr import Laminar
+from lmnr.opentelemetry_lib.tracing.context import get_current_context
+from lmnr.sdk.log import get_default_logger
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import unwrap
+from opentelemetry.trace import NonRecordingSpan, get_current_span
+from wrapt import wrap_function_wrapper
+
+
+_instruments = ("bubus >= 1.3.0",)
+event_id_to_span_context = {}
+logger = get_default_logger(__name__)
+
+
+def wrap_dispatch(wrapped, instance, args, kwargs):
+    event = args[0] if args and len(args) > 0 else kwargs.get("event", None)
+    if event and hasattr(event, "event_id"):
+        event_id = event.event_id
+        if event_id:
+            span = get_current_span(get_current_context())
+            event_id_to_span_context[event_id] = span.get_span_context()
+    return wrapped(*args, **kwargs)
+
+
+async def wrap_process_event(wrapped, instance, args, kwargs):
+    event = args[0] if args and len(args) > 0 else kwargs.get("event", None)
+    span_context = None
+    if event and hasattr(event, "event_id"):
+        event_id = event.event_id
+        if event_id:
+            span_context = event_id_to_span_context.get(event_id)
+    if not span_context:
+        return await wrapped(*args, **kwargs)
+    if not Laminar.is_initialized():
+        return await wrapped(*args, **kwargs)
+    with Laminar.use_span(NonRecordingSpan(span_context)):
+        return await wrapped(*args, **kwargs)
+
+
+class BubusInstrumentor(BaseInstrumentor):
+    def __init__(self):
+        super().__init__()
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        try:
+            wrap_function_wrapper("bubus.service", "EventBus.dispatch", wrap_dispatch)
+        except (ModuleNotFoundError, ImportError):
+            pass
+        try:
+            wrap_function_wrapper(
+                "bubus.service", "EventBus.process_event", wrap_process_event
+            )
+        except (ModuleNotFoundError, ImportError):
+            pass
+
+    def _uninstrument(self, **kwargs):
+        try:
+            unwrap("bubus.service", "EventBus.dispatch")
+        except (ModuleNotFoundError, ImportError):
+            pass
+        try:
+            unwrap("bubus.service", "EventBus.process_event")
+        except (ModuleNotFoundError, ImportError):
+            pass
+        event_id_to_span_context.clear()