Fix Langfuse: clean roots + explicit child parents

src/deepagents_rl/train/swe_env.py

@@ -194,6 +194,7 @@ class DeepAgentSWEEnv:
         self._root_cm: Any = None
         self._root: Any = None
         self._trace_id: str = ""
+        self._root_span_id: str = ""
 
     def _langfuse_enabled(self) -> bool:
         if os.getenv("DARL_LANGFUSE", "1").strip() in ("0", "false", "no"):
@@ -219,6 +220,17 @@ class DeepAgentSWEEnv:
             print(f"[langfuse] init failed: {e}")
             return None
 
+    def _child_trace_context(self) -> dict[str, str] | None:
+        """Attach children to THIS rollout's root — never ambient OTEL context.
+
+        Passing only ``trace_id`` (no parent) makes the Langfuse SDK invent a
+        NonRecordingSpan parent id → orphan parents + mega-trace pollution when
+        parallel GRPO rollouts share get_client().
+        """
+        if not self._trace_id or not self._root_span_id:
+            return None
+        return {"trace_id": self._trace_id, "parent_span_id": self._root_span_id}
+
     def _end_trace(self) -> None:
         if self._root_cm is not None:
             try:
@@ -227,6 +239,8 @@ class DeepAgentSWEEnv:
                 pass
         self._root_cm = None
         self._root = None
+        self._root_span_id = ""
+        # Keep _trace_id until next reset for logging; clear after flush.
         if self._lf is not None:
             try:
                 self._lf.flush()
@@ -249,7 +263,8 @@ class DeepAgentSWEEnv:
             input=payload,
         )
         lf = self._lf
-        if lf is None or self._root is None:
+        ctx = self._child_trace_context()
+        if lf is None or ctx is None:
             try:
                 out = method(self, *args, **kwargs)
             except Exception as e:  # noqa: BLE001
@@ -265,12 +280,14 @@ class DeepAgentSWEEnv:
         with lf.start_as_current_observation(
             name=name,
             as_type="tool",
+            trace_context=ctx,
             input=payload,
             metadata={
                 "rollout_id": self._rollout_id,
                 "instance_id": self._task.get("instance_id"),
                 "harness": "deepagents-swe-grpo",
                 "tool_index": n,
+                "trace_id": self._trace_id,
             },
         ) as span:
             try:
@@ -348,12 +365,13 @@ class DeepAgentSWEEnv:
         _vlog("env.reset.langfuse", enabled=lf is not None)
         if lf is not None:
             try:
-                # Fresh root trace per rollout (avoid nesting everything under one parent).
-                self._trace_id = lf.create_trace_id()
+                # IMPORTANT: do NOT pass trace_context={"trace_id": ...} alone.
+                # Langfuse SDK then fabricates a NonRecordingSpan parent → orphan
+                # parent ids in the UI and children from other rollouts can leak in.
+                # A plain start_as_current_observation creates a clean root trace.
                 self._root_cm = lf.start_as_current_observation(
                     name=f"grpo-swe-rollout:{instance_id}",
                     as_type="agent",
-                    trace_context={"trace_id": self._trace_id},
                     input={
                         "instance_id": instance_id,
                         "repo": repo,
@@ -381,12 +399,27 @@ class DeepAgentSWEEnv:
                     end_on_exit=True,
                 )
                 self._root = self._root_cm.__enter__()
-                _vlog("env.reset.langfuse_trace", trace_id=self._trace_id)
+                self._trace_id = str(
+                    getattr(self._root, "trace_id", None)
+                    or getattr(self._root, "traceId", None)
+                    or ""
+                )
+                self._root_span_id = str(
+                    getattr(self._root, "id", None)
+                    or getattr(self._root, "observation_id", None)
+                    or ""
+                )
+                _vlog(
+                    "env.reset.langfuse_trace",
+                    trace_id=self._trace_id,
+                    root_span_id=self._root_span_id,
+                )
             except Exception as e:  # noqa: BLE001
                 _vlog("env.reset.langfuse_failed", error=str(e))
                 self._root_cm = None
                 self._root = None
                 self._trace_id = ""
+                self._root_span_id = ""
 
         def _materialize() -> None:
             _vlog("env.reset.materialize", repo=repo, base_commit=base_commit[:12])
@@ -402,10 +435,12 @@ class DeepAgentSWEEnv:
             )
             _vlog("env.reset.sandbox_ready", sandbox=str(self._sandbox))
 
-        if lf is not None and self._root is not None:
+        ctx = self._child_trace_context()
+        if lf is not None and ctx is not None:
             with lf.start_as_current_observation(
                 name="repo_materialize",
                 as_type="span",
+                trace_context=ctx,
                 input={"repo": repo, "base_commit": base_commit},
             ) as span:
                 _materialize()
@@ -688,10 +723,12 @@ class DeepAgentSWEEnv:
 
             return float(report.get("reward_total") or 0.0)
 
-        if lf is not None and self._root is not None:
+        ctx = self._child_trace_context()
+        if lf is not None and ctx is not None:
             with lf.start_as_current_observation(
                 name="reward_grader",
                 as_type="evaluator",
+                trace_context=ctx,
                 input={
                     "instance_id": self._task.get("instance_id"),
                     "n_tool_calls": len(self._tool_calls),