Fix context propagation bugs and remove handler suppression

Bug 1: Replace stale _current_run snapshot with ambient context in
outbound interceptor. Add _get_current_run_for_propagation() helper
that filters _ContextBridgeRunTree from ambient context. Outbound
methods now read get_current_run_tree() for @Traceable nesting instead
of a frozen reference from workflow entry.

Bug 2: Add tracing_context() to Nexus inbound interceptor for both
execute_nexus_operation_start and execute_nexus_operation_cancel,
matching the activity inbound pattern. Ensures @Traceable functions
in Nexus handlers have a LangSmith client even with add_temporal_runs=False.

Remove handler suppression (is_handler check, _workflow_is_active flag)
to align with OTel interceptor which creates spans for all handlers
unconditionally.

Add dump_traces() to test infrastructure for per-root-trace assertions.
Restructure comprehensive tests so user_pipeline only wraps start_workflow,
with polling/signals/queries as independent root traces.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: xumaple

temporalio/contrib/langsmith/_interceptor.py

@@ -123,6 +123,18 @@ def _extract_nexus_context(
     return _ReplaySafeRunTree(run, executor=executor) if run else None
 
 
+def _get_current_run_for_propagation() -> RunTree | None:
+    """Get the current ambient run for context propagation.
+
+    Filters out _ContextBridgeRunTree, which is internal scaffolding
+    that should never be serialized into headers or used as parent runs.
+    """
+    run = get_current_run_tree()
+    if isinstance(run, _ContextBridgeRunTree):
+        return None
+    return run
+
+
 # ---------------------------------------------------------------------------
 # Sandbox safety: patch @traceable's aio_to_thread
 # ---------------------------------------------------------------------------
@@ -742,46 +754,26 @@ class _LangSmithWorkflowInboundInterceptor(
     """Instruments workflow execution with LangSmith runs."""
 
     _config: ClassVar[LangSmithInterceptor]
-    _current_run: _ReplaySafeRunTree | None = None
 
     def init(self, outbound: temporalio.worker.WorkflowOutboundInterceptor) -> None:
-        super().init(
-            _LangSmithWorkflowOutboundInterceptor(outbound, self._config, self)
-        )
+        super().init(_LangSmithWorkflowOutboundInterceptor(outbound, self._config))
 
     @contextmanager
     def _workflow_maybe_run(
         self,
         name: str,
         headers: Mapping[str, Payload] | None = None,
-        *,
-        is_handler: bool = False,
     ) -> Iterator[_ReplaySafeRunTree | None]:
         """Workflow-specific run creation with metadata.
 
-        Extracts parent from headers (if provided) and stores the run (or parent
-        fallback) as ``_current_run`` so the outbound interceptor can propagate
-        context even when ``add_temporal_runs=False``.
-
-        Always sets up ``tracing_context`` so ``@traceable`` functions called
-        from workflow code can discover the parent and LangSmith client,
-        independent of the ``add_temporal_runs`` toggle.
-
-        When ``is_handler`` is True and no LangSmith context is found in
-        headers, skips trace creation if a workflow run is already active
-        (``_current_run`` is set). This suppresses orphan traces from
-        uninstrumented client operations (e.g. query polling) while still
-        allowing handler traces when invoked with propagated context.
+        Extracts parent from headers (if provided) and sets up
+        ``tracing_context`` so ``@traceable`` functions called from workflow
+        code can discover the parent and LangSmith client, independent of the
+        ``add_temporal_runs`` toggle.
         """
         parent = _extract_context(headers, self._config._executor) if headers else None
         if parent is not None:
             parent.ls_client = self._config._client
-        # Handler from an uninstrumented client during workflow execution:
-        # no LangSmith headers but _current_run is set. Skip trace creation
-        # to avoid orphan/duplicate handler traces (e.g. query polling).
-        if is_handler and parent is None and self._current_run is not None:
-            yield None
-            return
         info = temporalio.workflow.info()
         extra_metadata = {
             "temporalWorkflowID": info.workflow_id,
@@ -815,12 +807,7 @@ def _workflow_maybe_run(
                 parent=parent,
                 extra_metadata=extra_metadata,
             ) as run:
-                prev_run = self._current_run
-                self._current_run = run or parent
-                try:
-                    yield run
-                finally:
-                    self._current_run = prev_run
+                yield run
 
     async def execute_workflow(
         self, input: temporalio.worker.ExecuteWorkflowInput
@@ -833,31 +820,23 @@ async def execute_workflow(
             return await super().execute_workflow(input)
 
     async def handle_signal(self, input: temporalio.worker.HandleSignalInput) -> None:
-        with self._workflow_maybe_run(
-            f"HandleSignal:{input.signal}", input.headers, is_handler=True
-        ):
+        with self._workflow_maybe_run(f"HandleSignal:{input.signal}", input.headers):
             return await super().handle_signal(input)
 
     async def handle_query(self, input: temporalio.worker.HandleQueryInput) -> Any:
-        with self._workflow_maybe_run(
-            f"HandleQuery:{input.query}", input.headers, is_handler=True
-        ):
+        with self._workflow_maybe_run(f"HandleQuery:{input.query}", input.headers):
             return await super().handle_query(input)
 
     def handle_update_validator(
         self, input: temporalio.worker.HandleUpdateInput
     ) -> None:
-        with self._workflow_maybe_run(
-            f"ValidateUpdate:{input.update}", input.headers, is_handler=True
-        ):
+        with self._workflow_maybe_run(f"ValidateUpdate:{input.update}", input.headers):
             return super().handle_update_validator(input)
 
     async def handle_update_handler(
         self, input: temporalio.worker.HandleUpdateInput
     ) -> Any:
-        with self._workflow_maybe_run(
-            f"HandleUpdate:{input.update}", input.headers, is_handler=True
-        ):
+        with self._workflow_maybe_run(f"HandleUpdate:{input.update}", input.headers):
             return await super().handle_update_handler(input)
 
 
@@ -875,19 +854,22 @@ def __init__(
         self,
         next: temporalio.worker.WorkflowOutboundInterceptor,
         config: LangSmithInterceptor,
-        inbound: _LangSmithWorkflowInboundInterceptor,
     ) -> None:
         super().__init__(next)
         self._config = config
-        self._inbound = inbound
 
     @contextmanager
     def _traced_outbound(
         self, name: str, input: _InputWithHeaders
     ) -> Iterator[_ReplaySafeRunTree | None]:
-        """Outbound workflow run creation with context injection into input.headers."""
-        with self._config.maybe_run(name, parent=self._inbound._current_run) as run:
-            context_source = run or self._inbound._current_run
+        """Outbound workflow run creation with context injection into input.headers.
+
+        Uses ambient context (``get_current_run_tree()``) instead of a cached
+        snapshot, so ``@traceable`` step functions that wrap outbound calls
+        correctly parent the outbound run under themselves.
+        """
+        with self._config.maybe_run(name) as run:
+            context_source = run or _get_current_run_for_propagation()
             if context_source:
                 input.headers = _inject_context(input.headers, context_source)
             yield run
@@ -923,8 +905,8 @@ async def signal_external_workflow(
             return await super().signal_external_workflow(input)
 
     def continue_as_new(self, input: temporalio.worker.ContinueAsNewInput) -> NoReturn:
-        # No trace created, but inject context from inbound's current run
-        current_run = getattr(self._inbound, "_current_run", None)
+        # No trace created, but inject context from ambient run
+        current_run = _get_current_run_for_propagation()
         if current_run:
             input.headers = _inject_context(input.headers, current_run)
         super().continue_as_new(input)
@@ -934,9 +916,8 @@ async def start_nexus_operation(
     ) -> temporalio.workflow.NexusOperationHandle[Any]:
         with self._config.maybe_run(
             f"StartNexusOperation:{input.service}/{input.operation_name}",
-            parent=self._inbound._current_run,
         ) as run:
-            context_source = run or self._inbound._current_run
+            context_source = run or _get_current_run_for_propagation()
             if context_source:
                 input.headers = _inject_nexus_context(
                     input.headers or {}, context_source
@@ -969,20 +950,42 @@ async def execute_nexus_operation_start(
         | nexusrpc.handler.StartOperationResultAsync
     ):
         parent = _extract_nexus_context(input.ctx.headers, self._config._executor)
-        with self._config.maybe_run(
-            f"RunStartNexusOperationHandler:{input.ctx.service}/{input.ctx.operation}",
-            run_type="tool",
-            parent=parent,
-        ):
-            return await self.next.execute_nexus_operation_start(input)
+        if parent is not None and hasattr(parent, "ls_client"):
+            parent.ls_client = self._config._client
+        ctx_kwargs: dict[str, Any] = {
+            "client": self._config._client,
+            "enabled": True,
+        }
+        if self._config._project_name:
+            ctx_kwargs["project_name"] = self._config._project_name
+        if parent:
+            ctx_kwargs["parent"] = parent
+        with tracing_context(**ctx_kwargs):
+            with self._config.maybe_run(
+                f"RunStartNexusOperationHandler:{input.ctx.service}/{input.ctx.operation}",
+                run_type="tool",
+                parent=parent,
+            ):
+                return await self.next.execute_nexus_operation_start(input)
 
     async def execute_nexus_operation_cancel(
         self, input: temporalio.worker.ExecuteNexusOperationCancelInput
     ) -> None:
         parent = _extract_nexus_context(input.ctx.headers, self._config._executor)
-        with self._config.maybe_run(
-            f"RunCancelNexusOperationHandler:{input.ctx.service}/{input.ctx.operation}",
-            run_type="tool",
-            parent=parent,
-        ):
-            return await self.next.execute_nexus_operation_cancel(input)
+        if parent is not None and hasattr(parent, "ls_client"):
+            parent.ls_client = self._config._client
+        ctx_kwargs: dict[str, Any] = {
+            "client": self._config._client,
+            "enabled": True,
+        }
+        if self._config._project_name:
+            ctx_kwargs["project_name"] = self._config._project_name
+        if parent:
+            ctx_kwargs["parent"] = parent
+        with tracing_context(**ctx_kwargs):
+            with self._config.maybe_run(
+                f"RunCancelNexusOperationHandler:{input.ctx.service}/{input.ctx.operation}",
+                run_type="tool",
+                parent=parent,
+            ):
+                return await self.next.execute_nexus_operation_cancel(input)

tests/contrib/langsmith/conftest.py

@@ -58,24 +58,18 @@ def clear(self) -> None:
         self._by_id.clear()
 
 
-def dump_runs(collector: InMemoryRunCollector) -> list[str]:
-    """Reconstruct parent-child hierarchy from collected runs.
+def dump_traces(collector: InMemoryRunCollector) -> list[list[str]]:
+    """Reconstruct parent-child hierarchy grouped by root trace.
 
-    Returns a list of indented strings, e.g.:
-        ["StartWorkflow:MyWf", "  RunWorkflow:MyWf", "    StartActivity:do_thing"]
+    Returns a list of traces, where each trace is a list of indented
+    strings (same format as dump_runs). Each trace starts from a
+    different root run.
     """
     runs = collector.runs
     children: dict[str | None, list[_RunRecord]] = {}
     for r in runs:
         children.setdefault(r.parent_run_id, []).append(r)
 
-    result: list[str] = []
-
-    def _walk(parent_id: str | None, depth: int) -> None:
-        for child in children.get(parent_id, []):
-            result.append("  " * depth + child.name)
-            _walk(child.id, depth + 1)
-
     # Strict: reject dangling parent references
     known_ids = {r.id for r in runs}
     for r in runs:
@@ -84,10 +78,26 @@ def _walk(parent_id: str | None, depth: int) -> None:
                 f"Run {r.name!r} (id={r.id}) has parent_run_id={r.parent_run_id} "
                 f"which is not in the collected runs — dangling parent reference"
             )
-    # Only walk true roots (parent_run_id is None)
-    _walk(None, 0)
 
-    return result
+    traces: list[list[str]] = []
+    for root in children.get(None, []):
+        trace: list[str] = []
+
+        def _walk(parent_id: str | None, depth: int) -> None:
+            for child in children.get(parent_id, []):
+                trace.append("  " * depth + child.name)
+                _walk(child.id, depth + 1)
+
+        trace.append(root.name)
+        _walk(root.id, 1)
+        traces.append(trace)
+
+    return traces
+
+
+def dump_runs(collector: InMemoryRunCollector) -> list[str]:
+    """Flat list of all runs across all traces."""
+    return [run for trace in dump_traces(collector) for run in trace]
 
 
 def make_mock_ls_client(collector: InMemoryRunCollector) -> MagicMock: