lmnr 0.6.21__py3-none-any.whl → 0.7.0__py3-none-any.whl

lmnr/opentelemetry_lib/tracing/__init__.py

@@ -10,9 +10,22 @@ from lmnr.opentelemetry_lib.tracing.instruments import (
     Instruments,
     init_instrumentations,
 )
+from lmnr.opentelemetry_lib.tracing.context import (
+    attach_context,
+    detach_context,
+    get_current_context,
+    get_token_stack,
+    _isolated_token_stack,
+    _isolated_token_stack_storage,
+    set_token_stack,
+)
 
 from opentelemetry import trace
-from opentelemetry.instrumentation.threading import ThreadingInstrumentor
+from opentelemetry.context import Context
+
+# instead of importing from opentelemetry.instrumentation.threading,
+# we import from our modified copy to use Laminar's isolated context.
+from ..opentelemetry.instrumentation.threading import ThreadingInstrumentor
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider, SpanProcessor
 from opentelemetry.sdk.trace.export import SpanExporter
@@ -32,6 +45,7 @@ class TracerWrapper(object):
     _async_client: AsyncLaminarClient
     _resource: Resource
     _span_processor: SpanProcessor
+    _original_thread_init = None
 
     def __new__(
         cls,
@@ -91,6 +105,9 @@ class TracerWrapper(object):
 
                 obj._tracer_provider.add_span_processor(obj._span_processor)
 
+                # Setup threading context inheritance
+                obj._setup_threading_inheritance()
+
                 # This is not a real instrumentation and does not generate telemetry
                 # data, but it is required to ensure that OpenTelemetry context
                 # propagation is enabled.
@@ -113,6 +130,43 @@ class TracerWrapper(object):
 
             return cls.instance
 
+    def _setup_threading_inheritance(self):
+        """Setup threading inheritance for isolated context."""
+        if TracerWrapper._original_thread_init is None:
+            # Monkey patch Thread.__init__ to capture context inheritance
+            TracerWrapper._original_thread_init = threading.Thread.__init__
+
+            def patched_thread_init(thread_self, *args, **kwargs):
+                # Capture current isolated context and token stack for inheritance
+                current_context = get_current_context()
+                current_token_stack = get_token_stack().copy()
+
+                # Get the original target function
+                original_target = kwargs.get("target")
+                if not original_target and args:
+                    original_target = args[0]
+
+                # Only inherit if we have a target function
+                if original_target:
+                    # Create a wrapper function that sets up context
+                    def thread_wrapper(*target_args, **target_kwargs):
+                        # Set inherited context and token stack in the new thread
+                        attach_context(current_context)
+                        set_token_stack(current_token_stack)
+                        # Run original target
+                        return original_target(*target_args, **target_kwargs)
+
+                    # Replace the target with our wrapper
+                    if "target" in kwargs:
+                        kwargs["target"] = thread_wrapper
+                    elif args:
+                        args = (thread_wrapper,) + args[1:]
+
+                # Call original init
+                TracerWrapper._original_thread_init(thread_self, *args, **kwargs)
+
+            threading.Thread.__init__ = patched_thread_init
+
     def exit_handler(self):
         if isinstance(self._span_processor, LaminarSpanProcessor):
             self._span_processor.clear()
@@ -124,6 +178,31 @@ class TracerWrapper(object):
         console_log_handler.setFormatter(VerboseColorfulFormatter())
         self._logger.addHandler(console_log_handler)
 
+    def get_isolated_context(self) -> Context:
+        """Get the current isolated context."""
+        return get_current_context()
+
+    def push_span_context(self, span: trace.Span) -> Context:
+        """Push a new context with the given span onto the stack."""
+        current_ctx = get_current_context()
+        new_context = trace.set_span_in_context(span, current_ctx)
+        token = attach_context(new_context)
+
+        # Store the token for later detachment - tokens are much lighter than contexts
+        current_stack = get_token_stack().copy()
+        current_stack.append(token)
+        set_token_stack(current_stack)
+
+        return new_context
+
+    def pop_span_context(self) -> None:
+        """Pop the current span context from the stack."""
+        current_stack = get_token_stack().copy()
+        if current_stack:
+            token = current_stack.pop()
+            set_token_stack(current_stack)
+            detach_context(token)
+
     @staticmethod
     def set_static_params(
         resource_attributes: dict,
@@ -144,6 +223,15 @@ class TracerWrapper(object):
         # Any state cleanup. Now used in between tests
         if isinstance(cls.instance._span_processor, LaminarSpanProcessor):
             cls.instance._span_processor.clear()
+        # Clear the isolated context state for clean test state
+        try:
+            _isolated_token_stack.set([])
+        except LookupError:
+            pass
+        if hasattr(_isolated_token_stack_storage, "token_stack"):
+            _isolated_token_stack_storage.token_stack = []
+        # Reset the isolated context to a fresh state
+        attach_context(Context())
 
     def shutdown(self):
         if self._tracer_provider is None:

lmnr/opentelemetry_lib/tracing/context.py

@@ -0,0 +1,109 @@
+import threading
+
+from abc import ABC, abstractmethod
+from contextvars import ContextVar
+from opentelemetry.context import Context, Token
+
+
+class _IsolatedRuntimeContext(ABC):
+    """The isolated RuntimeContext interface, identical to OpenTelemetry's _RuntimeContext
+    but isolated from the global context.
+    """
+
+    @abstractmethod
+    def attach(self, context: Context) -> Token[Context]:
+        """Sets the current `Context` object. Returns a
+        token that can be used to reset to the previous `Context`.
+
+        Args:
+            context: The Context to set.
+        """
+
+    @abstractmethod
+    def get_current(self) -> Context:
+        """Returns the current `Context` object."""
+
+    @abstractmethod
+    def detach(self, token: Token[Context]) -> None:
+        """Resets Context to a previous value
+
+        Args:
+            token: A reference to a previous Context.
+        """
+
+
+class IsolatedContextVarsRuntimeContext(_IsolatedRuntimeContext):
+    """An isolated implementation of the RuntimeContext interface which wraps ContextVar
+    but uses its own ContextVar instead of the global one.
+    """
+
+    def __init__(self) -> None:
+        self._current_context = ContextVar(
+            "isolated_current_context", default=Context()
+        )
+
+    def attach(self, context: Context) -> Token[Context]:
+        """Sets the current `Context` object. Returns a
+        token that can be used to reset to the previous `Context`.
+
+        Args:
+            context: The Context to set.
+        """
+        return self._current_context.set(context)
+
+    def get_current(self) -> Context:
+        """Returns the current `Context` object."""
+        return self._current_context.get()
+
+    def detach(self, token: Token[Context]) -> None:
+        """Resets Context to a previous value
+
+        Args:
+            token: A reference to a previous Context.
+        """
+        self._current_context.reset(token)
+
+
+# Create the isolated runtime context
+_ISOLATED_RUNTIME_CONTEXT = IsolatedContextVarsRuntimeContext()
+
+# Token stack for push/pop API compatibility - much lighter than copying contexts
+_isolated_token_stack: ContextVar[list[Token[Context]]] = ContextVar(
+    "isolated_token_stack", default=[]
+)
+
+# Thread-local storage for threading support
+_isolated_token_stack_storage = threading.local()
+
+
+def get_token_stack() -> list[Token[Context]]:
+    """Get the token stack, supporting both asyncio and threading."""
+    try:
+        return _isolated_token_stack.get()
+    except LookupError:
+        if not hasattr(_isolated_token_stack_storage, "token_stack"):
+            _isolated_token_stack_storage.token_stack = []
+        return _isolated_token_stack_storage.token_stack
+
+
+def set_token_stack(stack: list[Token[Context]]) -> None:
+    """Set the token stack, supporting both asyncio and threading."""
+    try:
+        _isolated_token_stack.set(stack)
+    except LookupError:
+        _isolated_token_stack_storage.token_stack = stack
+
+
+def get_current_context() -> Context:
+    """Get the current isolated context."""
+    return _ISOLATED_RUNTIME_CONTEXT.get_current()
+
+
+def attach_context(context: Context) -> Token[Context]:
+    """Attach a context to the isolated runtime context."""
+    return _ISOLATED_RUNTIME_CONTEXT.attach(context)
+
+
+def detach_context(token: Token[Context]) -> None:
+    """Detach a context from the isolated runtime context."""
+    _ISOLATED_RUNTIME_CONTEXT.detach(token)

lmnr/opentelemetry_lib/tracing/processor.py

@@ -17,9 +17,6 @@ from lmnr.opentelemetry_lib.tracing.attributes import (
     SPAN_SDK_VERSION,
 )
 from lmnr.opentelemetry_lib.tracing.exporter import LaminarSpanExporter
-from lmnr.opentelemetry_lib.tracing.context_properties import (
-    _set_association_properties_attributes,
-)
 from lmnr.version import PYTHON_VERSION, __version__
 
 
@@ -76,9 +73,11 @@ class LaminarSpanProcessor(SpanProcessor):
         span.set_attribute(SPAN_SDK_VERSION, __version__)
         span.set_attribute(SPAN_LANGUAGE_VERSION, f"python@{PYTHON_VERSION}")
 
-        association_properties = get_value("association_properties")
-        if association_properties is not None:
-            _set_association_properties_attributes(span, association_properties)
+        if span.name == "LangGraph.workflow":
+            graph_context = get_value("lmnr.langgraph.graph") or {}
+            for key, value in graph_context.items():
+                span.set_attribute(f"lmnr.association.properties.{key}", value)
+
         self.instance.on_start(span, parent_context)
 
     def on_end(self, span: Span):

lmnr/opentelemetry_lib/tracing/tracer.py

@@ -1,6 +1,8 @@
 from contextlib import contextmanager
+from typing import Generator, Tuple
 
 from opentelemetry import trace
+from opentelemetry.context import Context
 from lmnr.opentelemetry_lib.tracing import TracerWrapper
 
 
@@ -16,3 +18,30 @@ def get_tracer(flush_on_exit: bool = False):
     finally:
         if flush_on_exit:
             wrapper.flush()
+
+
+@contextmanager
+def get_tracer_with_context(
+    flush_on_exit: bool = False,
+) -> Generator[Tuple[trace.Tracer, Context], None, None]:
+    """Get tracer with isolated context. Returns (tracer, context) tuple."""
+    wrapper = TracerWrapper()
+    try:
+        tracer = wrapper.get_tracer()
+        context = wrapper.get_isolated_context()
+        yield tracer, context
+    finally:
+        if flush_on_exit:
+            wrapper.flush()
+
+
+def copy_current_context() -> Context:
+    """Copy the current isolated context for use in threads/tasks."""
+    wrapper = TracerWrapper()
+    return wrapper.get_isolated_context()
+
+
+def set_context_for_thread(context: Context) -> None:
+    """Set the isolated context for the current thread."""
+    wrapper = TracerWrapper()
+    wrapper.set_isolated_context(context)

lmnr/sdk/browser/browser_use_otel.py

@@ -28,8 +28,8 @@ WRAPPED_METHODS = [
         "object": "Agent",
         "method": "run",
         "span_name": "agent.run",
-        "ignore_input": False,
-        "ignore_output": False,
+        "ignore_input": True,
+        "ignore_output": True,
         "span_type": "DEFAULT",
     },
     {
@@ -47,15 +47,15 @@ WRAPPED_METHODS = [
         "method": "act",
         "span_name": "controller.act",
         "ignore_input": True,
-        "ignore_output": False,
+        "ignore_output": True,
         "span_type": "DEFAULT",
     },
     {
         "package": "browser_use.controller.registry.service",
         "object": "Registry",
         "method": "execute_action",
-        "ignore_input": False,
-        "ignore_output": False,
+        "ignore_input": True,
+        "ignore_output": True,
         "span_type": "TOOL",
     },
 ]

lmnr/sdk/browser/patchright_otel.py

@@ -1,4 +1,6 @@
 from lmnr.sdk.browser.playwright_otel import (
+    _wrap_bring_to_front_async,
+    _wrap_bring_to_front_sync,
     _wrap_new_browser_sync,
     _wrap_new_browser_async,
     _wrap_new_context_sync,
@@ -46,6 +48,12 @@ WRAPPED_METHODS = [
         "method": "launch_persistent_context",
         "wrapper": _wrap_new_context_sync,
     },
+    {
+        "package": "patchright.sync_api",
+        "object": "Page",
+        "method": "bring_to_front",
+        "wrapper": _wrap_bring_to_front_sync,
+    },
 ]
 
 WRAPPED_METHODS_ASYNC = [
@@ -79,6 +87,12 @@ WRAPPED_METHODS_ASYNC = [
         "method": "launch_persistent_context",
         "wrapper": _wrap_new_context_async,
     },
+    {
+        "package": "patchright.async_api",
+        "object": "Page",
+        "method": "bring_to_front",
+        "wrapper": _wrap_bring_to_front_async,
+    },
 ]
 
 

lmnr/sdk/browser/playwright_otel.py

@@ -60,8 +60,15 @@ def _wrap_new_browser_sync(
     browser: SyncBrowser = wrapped(*args, **kwargs)
     session_id = str(uuid.uuid4().hex)
 
+    def create_page_handler(session_id, client):
+        def page_handler(page):
+            start_recording_events_sync(page, session_id, client)
+
+        return page_handler
+
     for context in browser.contexts:
-        context.on("page", lambda p: start_recording_events_sync(p, session_id, client))
+        page_handler = create_page_handler(session_id, client)
+        context.on("page", page_handler)
         for page in context.pages:
             start_recording_events_sync(page, session_id, client)
 
@@ -75,10 +82,15 @@ async def _wrap_new_browser_async(
     browser: Browser = await wrapped(*args, **kwargs)
     session_id = str(uuid.uuid4().hex)
 
+    def create_page_handler(session_id, client):
+        async def page_handler(page):
+            await start_recording_events_async(page, session_id, client)
+
+        return page_handler
+
     for context in browser.contexts:
-        context.on(
-            "page", lambda p: start_recording_events_async(p, session_id, client)
-        )
+        page_handler = create_page_handler(session_id, client)
+        context.on("page", page_handler)
         for page in context.pages:
             await start_recording_events_async(page, session_id, client)
     return browser
@@ -91,7 +103,14 @@ def _wrap_new_context_sync(
     context: SyncBrowserContext = wrapped(*args, **kwargs)
     session_id = str(uuid.uuid4().hex)
 
-    context.on("page", lambda p: start_recording_events_sync(p, session_id, client))
+    def create_page_handler(session_id, client):
+        def page_handler(page):
+            start_recording_events_sync(page, session_id, client)
+
+        return page_handler
+
+    page_handler = create_page_handler(session_id, client)
+    context.on("page", page_handler)
     for page in context.pages:
         start_recording_events_sync(page, session_id, client)
 
@@ -105,7 +124,14 @@ async def _wrap_new_context_async(
     context: BrowserContext = await wrapped(*args, **kwargs)
     session_id = str(uuid.uuid4().hex)
 
-    context.on("page", lambda p: start_recording_events_async(p, session_id, client))
+    def create_page_handler(session_id, client):
+        async def page_handler(page):
+            await start_recording_events_async(page, session_id, client)
+
+        return page_handler
+
+    page_handler = create_page_handler(session_id, client)
+    context.on("page", page_handler)
     for page in context.pages:
         await start_recording_events_async(page, session_id, client)
 

lmnr/sdk/browser/pw_utils.py

@@ -8,6 +8,7 @@ from lmnr.sdk.decorators import observe
 from lmnr.sdk.browser.utils import retry_sync, retry_async
 from lmnr.sdk.client.synchronous.sync_client import LaminarClient
 from lmnr.sdk.client.asynchronous.async_client import AsyncLaminarClient
+from lmnr.opentelemetry_lib.tracing.context import get_current_context
 
 try:
     if is_package_installed("playwright"):
@@ -38,7 +39,9 @@ with open(os.path.join(current_dir, "rrweb", "rrweb.umd.min.cjs"), "r") as f:
 INJECT_PLACEHOLDER = """
 () => {
     const BATCH_TIMEOUT = 2000; // Send events after 2 seconds
-
+    const MAX_WORKER_PROMISES = 50; // Max concurrent worker promises
+    const HEARTBEAT_INTERVAL = 1000;
+    
     window.lmnrRrwebEventsBatch = [];
 
     // Create a Web Worker for heavy JSON processing with chunked processing
@@ -97,6 +100,29 @@ INJECT_PLACEHOLDER = """
     let workerPromises = new Map();
     let workerId = 0;
 
+    // Cleanup function for worker
+    const cleanupWorker = () => {
+        if (compressionWorker) {
+            compressionWorker.terminate();
+            compressionWorker = null;
+        }
+        workerPromises.clear();
+        workerId = 0;
+    };
+
+    // Clean up stale promises to prevent memory leaks
+    const cleanupStalePromises = () => {
+        if (workerPromises.size > MAX_WORKER_PROMISES) {
+            const toDelete = [];
+            for (const [id, promise] of workerPromises) {
+                if (toDelete.length >= workerPromises.size - MAX_WORKER_PROMISES) break;
+                toDelete.push(id);
+                promise.reject(new Error('Promise cleaned up due to memory pressure'));
+            }
+            toDelete.forEach(id => workerPromises.delete(id));
+        }
+    };
+
     // Non-blocking JSON.stringify using chunked processing
     function stringifyNonBlocking(obj, chunkSize = 10000) {
         return new Promise((resolve, reject) => {
@@ -196,6 +222,9 @@ INJECT_PLACEHOLDER = """
     // Alternative: Use transferable objects for maximum efficiency
     async function compressLargeObjectTransferable(data) {
         try {
+            // Clean up stale promises first
+            cleanupStalePromises();
+            
             // Stringify on main thread but non-blocking
             const jsonString = await stringifyNonBlocking(data);
 
@@ -219,11 +248,24 @@ INJECT_PLACEHOLDER = """
                             }
                         }
                     };
+                    
+                    compressionWorker.onerror = (error) => {
+                        console.error('Compression worker error:', error);
+                        cleanupWorker();
+                    };
                 }
 
                 const id = ++workerId;
                 workerPromises.set(id, { resolve, reject });
 
+                // Set timeout to prevent hanging promises
+                setTimeout(() => {
+                    if (workerPromises.has(id)) {
+                        workerPromises.delete(id);
+                        reject(new Error('Compression timeout'));
+                    }
+                }, 10000);
+
                 // Transfer the ArrayBuffer (no copying!)
                 compressionWorker.postMessage({
                     buffer,
@@ -262,15 +304,32 @@ INJECT_PLACEHOLDER = """
                             }
                         }
                     };
+                    
+                    compressionWorker.onerror = (error) => {
+                        console.error('Compression worker error:', error);
+                        cleanupWorker();
+                    };
                 }
 
                 const id = ++workerId;
                 workerPromises.set(id, { resolve, reject });
+                
+                // Set timeout to prevent hanging promises
+                setTimeout(() => {
+                    if (workerPromises.has(id)) {
+                        workerPromises.delete(id);
+                        reject(new Error('Compression timeout'));
+                    }
+                }, 10000);
+                
                 compressionWorker.postMessage({ jsonString, id });
             });
         }
     }
 
+    
+    setInterval(cleanupWorker, 5000);
+    
     function isLargeEvent(type) {
         const LARGE_EVENT_TYPES = [
             2, // FullSnapshot
@@ -300,13 +359,22 @@ INJECT_PLACEHOLDER = """
     setInterval(sendBatchIfReady, BATCH_TIMEOUT);
 
     // Add heartbeat events
-    setInterval(async () => {
+    setInterval(() => {
         window.lmnrRrweb.record.addCustomEvent('heartbeat', {
             title: document.title,
             url: document.URL,
         })
-    }, 1000);
+    }, HEARTBEAT_INTERVAL);
 
+    async function bufferToBase64(buffer) {
+        const base64url = await new Promise(r => {
+            const reader = new FileReader()
+            reader.onload = () => r(reader.result)
+            reader.readAsDataURL(new Blob([buffer]))
+        });
+        return base64url.slice(base64url.indexOf(',') + 1);
+    }
+ 
     window.lmnrRrweb.record({
         async emit(event) {
             try {
@@ -315,9 +383,10 @@ INJECT_PLACEHOLDER = """
                     await compressLargeObject(event.data, true) :
                     await compressSmallObject(event.data);
 
+                const base64Data = await bufferToBase64(compressedResult);
                 const eventToSend = {
                     ...event,
-                    data: compressedResult,
+                    data: base64Data,
                 };
                 window.lmnrRrwebEventsBatch.push(eventToSend);
             } catch (error) {
@@ -461,7 +530,9 @@ async def inject_session_recorder_async(page: Page):
 
 @observe(name="playwright.page", ignore_input=True, ignore_output=True)
 def start_recording_events_sync(page: SyncPage, session_id: str, client: LaminarClient):
-    span = trace.get_current_span()
+
+    ctx = get_current_context()
+    span = trace.get_current_span(ctx)
     trace_id = format(span.get_span_context().trace_id, "032x")
     span.set_attribute("lmnr.internal.has_browser_session", True)
 
@@ -506,7 +577,8 @@ def start_recording_events_sync(page: SyncPage, session_id: str, client: Laminar
 async def start_recording_events_async(
     page: Page, session_id: str, client: AsyncLaminarClient
 ):
-    span = trace.get_current_span()
+    ctx = get_current_context()
+    span = trace.get_current_span(ctx)
     trace_id = format(span.get_span_context().trace_id, "032x")
     span.set_attribute("lmnr.internal.has_browser_session", True)
 

lmnr/sdk/client/asynchronous/resources/browser_events.py

@@ -25,6 +25,7 @@ class AsyncBrowserEvents(BaseAsyncResource):
             "source": f"python@{PYTHON_VERSION}",
             "sdkVersion": __version__,
         }
+
         compressed_payload = gzip.compress(json.dumps(payload).encode("utf-8"))
         response = await self._client.post(
             url,