insightfactory-databricks-langgraph-tracer 1.0.0.dev15__tar.gz → 1.0.1.dev18__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (20) hide show
  1. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/CHANGELOG.md +52 -1
  2. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/PKG-INFO +76 -1
  3. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/README.md +75 -0
  4. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/pyproject.toml +1 -1
  5. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/mlflow_utils.py +15 -0
  6. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/tracer.py +250 -0
  7. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/.gitignore +0 -0
  8. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/LICENSE +0 -0
  9. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/THIRD_PARTY_NOTICES +0 -0
  10. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/__init__.py +0 -0
  11. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/_generated_keys.py +0 -0
  12. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/aggregate.py +0 -0
  13. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/autolog_compat.py +0 -0
  14. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/cost.py +0 -0
  15. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/databricks_utils.py +0 -0
  16. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/errors.py +0 -0
  17. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/pricing/model-pricing-overrides.json +0 -0
  18. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/pricing/model-pricing.json +0 -0
  19. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/py.typed +0 -0
  20. {insightfactory_databricks_langgraph_tracer-1.0.0.dev15 → insightfactory_databricks_langgraph_tracer-1.0.1.dev18}/src/databricks_langgraph_tracer/uc_export.py +0 -0
@@ -20,6 +20,56 @@ from this section.
20
20
 
21
21
  ## [Unreleased]
22
22
 
23
+ ### Added
24
+
25
+ - `context_propagation` on `configure_databricks_tracing(...)` and on
26
+ `DatabricksLangGraphTracer(...)`, with the `DATABRICKS_TRACING_CONTEXT_PROPAGATION`
27
+ environment variable behind it. On by default; set it to `false`, `0`, `no` or
28
+ `off` to keep the previous behaviour, which also restores LangChain's default
29
+ callback dispatch.
30
+
31
+ ### Changed
32
+
33
+ - The tracer's callbacks now run inline on the caller's thread while context
34
+ propagation is on, which on an async graph means the event loop. A
35
+ `content_ref_resolver` that blocks now blocks the loop rather than a worker
36
+ thread, so keep it fast.
37
+
38
+ ### Fixed
39
+
40
+ - Spans created outside the tracer now nest under the running graph instead of
41
+ opening their own root trace (#52). A helper decorated with `@mlflow.trace`, or an
42
+ MLflow autolog integration, starts its span from MLflow's context, which the tracer
43
+ previously left empty; a consumer tracing its database helpers that way got one
44
+ extra root trace per query, none of them carrying the root tags, session metadata
45
+ or cost rollup. The tracer now publishes each live span into that context, and runs
46
+ its callbacks inline so the context reaches the code the node calls. Own-span
47
+ parenting is unchanged and still comes from LangChain's run tree, so concurrent
48
+ graph branches cannot cross-parent, and only the first tracer attached to a run
49
+ takes the context. Runs that overlap rather than nest, two `graph.stream()`
50
+ iterators consumed in start order, release their context in the order
51
+ OpenTelemetry requires: a run that ends while another is still stacked above it
52
+ holds its attach until that one ends, so a finished span is never restored over a
53
+ live one. Two limits come with that. A run whose end callback never fires pins the
54
+ attaches below it, which past a depth no real graph reaches is now reported at
55
+ warning. And another handler on the same run that also publishes to MLflow's
56
+ context, in practice `MlflowLangchainTracer` from `mlflow.langchain.autolog()`,
57
+ cannot be coordinated with: LangChain calls the two handlers' end hooks in
58
+ registration order, so whichever releases last restores an ended span. Two of
59
+ MLflow's own tracers on one run leak the same way with this tracer absent, so no
60
+ release order fixes it from inside one handler. The tracer warns once when it
61
+ detects the conflict and names both settings to turn off. Issue #57 tracks it.
62
+
63
+ ## [1.0.1] - 2026-09-03
64
+
65
+ ### Changed
66
+
67
+ - Rewrote the README that PyPI renders in plainer prose. Commands, code samples
68
+ and configuration tables are unchanged, so there is no behaviour or packaging
69
+ change.
70
+ - Corrected the truncation placeholder example, which showed a trailing ellipsis
71
+ in `preview`. The tracer stores a plain slice of the original string.
72
+
23
73
  ## [1.0.0] - 2026-09-01
24
74
  ### Changed
25
75
  - Renamed the PyPI distribution from `databricks-langgraph-tracer` to
@@ -163,7 +213,8 @@ reserved by the repo's immutable-releases policy and cannot be created. Restore
163
213
  python-vX.Y.Z tag links if the names are ever reclaimed. 0.5.0's tag is created
164
214
  out-of-band (its bump predates the release workflow); 0.5.1+ tags are minted by
165
215
  release.yml on push to main. -->
166
- [Unreleased]: https://github.com/insightfactory-ai/if_s_langraph_mlflow_tracer/compare/python-v1.0.0...HEAD
216
+ [Unreleased]: https://github.com/insightfactory-ai/if_s_langraph_mlflow_tracer/compare/python-v1.0.1...HEAD
217
+ [1.0.1]: https://github.com/insightfactory-ai/if_s_langraph_mlflow_tracer/compare/python-v1.0.0...python-v1.0.1
167
218
  [1.0.0]: https://github.com/insightfactory-ai/if_s_langraph_mlflow_tracer/compare/python-v0.6.1...python-v1.0.0
168
219
  [0.6.1]: https://github.com/insightfactory-ai/if_s_langraph_mlflow_tracer/compare/python-v0.6.0...python-v0.6.1
169
220
  [0.6.0]: https://github.com/insightfactory-ai/if_s_langraph_mlflow_tracer/compare/python-v0.5.0...python-v0.6.0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: insightfactory-databricks-langgraph-tracer
3
- Version: 1.0.0.dev15
3
+ Version: 1.0.1.dev18
4
4
  Summary: LangGraph tracing for Databricks MLflow
5
5
  Project-URL: Homepage, https://insightfactory.ai
6
6
  Author-email: "insightfactory.ai" <support@insightfactory.ai>
@@ -71,6 +71,7 @@ overrides the matching variable.
71
71
  | Source tag | none | `source=` kwarg, default `langgraph` |
72
72
  | Multimodal refs | none | Inline image, PDF and file bytes become a reference. `content_ref_resolver=` chooses it. See below. |
73
73
  | Text cap | `DATABRICKS_TRACING_MAX_STRING_CHARS` | Opt-in. `max_string_chars=` truncates long plain-text span content. See below. |
74
+ | Context propagation | `DATABRICKS_TRACING_CONTEXT_PROPAGATION` | On by default. Set to `false`, `0`, `no` or `off`, or pass `context_propagation=False`, to stop publishing spans into MLflow's context. See below. |
74
75
  | Disable | `TESTING` or `BUILDING` set to `true`, or `enabled=False` | Turns the tracer into a no-op. This is the only path that does not raise. |
75
76
  | UC-backed tracing | `MLFLOW_TRACING_UC_BACKED=true` or `uc_tracing=True` | See below. |
76
77
  | SQL warehouse | `MLFLOW_TRACING_SQL_WAREHOUSE_ID`, falling back to `DATABRICKS_WAREHOUSE_ID` | Required when UC-backed. |
@@ -192,6 +193,80 @@ bytes, and multibyte text can be up to four times larger in bytes than in charac
192
193
  CJK or emoji-heavy content size the cap below a quarter of the limit. The placeholder's
193
194
  `bytes` field always reports the exact UTF-8 size of the original.
194
195
 
196
+ ### Spans from outside the tracer
197
+
198
+ A helper decorated with `@mlflow.trace`, or an MLflow autolog integration, starts its
199
+ span through MLflow's context rather than through LangChain's callbacks. The tracer
200
+ publishes each live span into that context, so those spans nest under whichever node is
201
+ running instead of opening a root trace of their own:
202
+
203
+ ```python
204
+ import mlflow
205
+
206
+ @mlflow.trace(span_type="TOOL", name="lakebase.execute_query")
207
+ def execute_query(sql: str): ...
208
+ ```
209
+
210
+ Calling that from a graph node produces one trace, with the tool span parented to the
211
+ node. Without it you get a second root trace per call, carrying none of the root tags,
212
+ session metadata or cost rollup the tracer writes.
213
+
214
+ Attaching the context requires the tracer's callbacks to run in the caller's own
215
+ context, so `run_inline` is set on the handler whenever propagation is on. Only
216
+ LangChain's async callback manager reads that flag: without it an async run sends the
217
+ handler through an executor with a copied context, which the node body never sees. The
218
+ sync manager calls handlers in the caller's thread either way.
219
+
220
+ The tracer still parents its own spans from LangChain's run tree, never from the
221
+ context, so concurrent branches of a graph cannot cross-parent. Runs that overlap
222
+ rather than nest, two `stream(...)` generators created before either is consumed,
223
+ release their attaches in the order OpenTelemetry requires: a run that ends while
224
+ another is still stacked above it holds its attach until that one ends, so a finished
225
+ span is never restored over a live one.
226
+
227
+ The price of never releasing out of order is that a run whose end callback never
228
+ fires pins the attaches beneath it. A node that abandons an inner `stream(...)`
229
+ part-consumed leaves its span active for the life of that context, and spans started
230
+ there parent to a run that has already finished. Sweeping the stack when the
231
+ outermost run ends would clear it and bring back the overlapping-run bug above, so
232
+ the tracer reports it instead: past a depth no real graph reaches, it logs a warning
233
+ naming the count.
234
+
235
+ Only one handler attaches per run. Attach a second context-propagating tracer to the
236
+ same graph and it records its own spans as before, but leaves the context to the
237
+ first.
238
+
239
+ Three limits are worth knowing:
240
+
241
+ - It only works in one direction. A root run is always created without a parent and
242
+ without consulting the context, so a graph invoked from *inside* an `@mlflow.trace`
243
+ function still opens its own root trace rather than nesting into the caller's. This
244
+ is the one place the TypeScript package behaves differently: MLflow's JavaScript
245
+ `startSpan` falls back to the active context, so a graph there does nest into an
246
+ enclosing span.
247
+ - Foreign spans are recorded by MLflow, not by this tracer, so neither the text cap
248
+ nor the multimodal reference substitution applies to them. A traced query helper
249
+ returning a large rowset can push a trace back over the inline read limit that
250
+ `max_string_chars` was set to stay under.
251
+ - The tracer's callbacks run inline on the caller's thread, which on an async graph
252
+ means the event loop. Keep `content_ref_resolver` fast and non-blocking.
253
+ - Another handler on the same run that also publishes to MLflow's context, which in
254
+ practice means `MlflowLangchainTracer` from `mlflow.langchain.autolog()`, cannot be
255
+ coordinated with. LangChain calls the two handlers' end hooks in registration order,
256
+ so the attaches nest but the releases run front to back, and whichever releases last
257
+ restores an ended span. Two of MLflow's own tracers on one run leak the same way
258
+ with this tracer absent, so no release order fixes it here. The tracer warns once
259
+ when it detects the conflict and names both settings to turn off; turn one of them
260
+ off. Issue #57 tracks it.
261
+
262
+ To turn the whole behaviour off, pass `context_propagation=False` or set
263
+ `DATABRICKS_TRACING_CONTEXT_PROPAGATION=false` (`0`, `no` and `off` also work). That
264
+ also restores LangChain's default callback dispatch.
265
+
266
+ The TypeScript package reaches the same outcome through `getActiveTracingSpan()` and
267
+ `withTracingContext()`, named at the call site, because Node has no ambient
268
+ equivalent. See [`../typescript/README.md`](../typescript/README.md).
269
+
195
270
  ### Autolog fallback
196
271
 
197
272
  `mode="autolog"` wires MLflow's built-in LangChain autologging plus compatibility patches. It
@@ -47,6 +47,7 @@ overrides the matching variable.
47
47
  | Source tag | none | `source=` kwarg, default `langgraph` |
48
48
  | Multimodal refs | none | Inline image, PDF and file bytes become a reference. `content_ref_resolver=` chooses it. See below. |
49
49
  | Text cap | `DATABRICKS_TRACING_MAX_STRING_CHARS` | Opt-in. `max_string_chars=` truncates long plain-text span content. See below. |
50
+ | Context propagation | `DATABRICKS_TRACING_CONTEXT_PROPAGATION` | On by default. Set to `false`, `0`, `no` or `off`, or pass `context_propagation=False`, to stop publishing spans into MLflow's context. See below. |
50
51
  | Disable | `TESTING` or `BUILDING` set to `true`, or `enabled=False` | Turns the tracer into a no-op. This is the only path that does not raise. |
51
52
  | UC-backed tracing | `MLFLOW_TRACING_UC_BACKED=true` or `uc_tracing=True` | See below. |
52
53
  | SQL warehouse | `MLFLOW_TRACING_SQL_WAREHOUSE_ID`, falling back to `DATABRICKS_WAREHOUSE_ID` | Required when UC-backed. |
@@ -168,6 +169,80 @@ bytes, and multibyte text can be up to four times larger in bytes than in charac
168
169
  CJK or emoji-heavy content size the cap below a quarter of the limit. The placeholder's
169
170
  `bytes` field always reports the exact UTF-8 size of the original.
170
171
 
172
+ ### Spans from outside the tracer
173
+
174
+ A helper decorated with `@mlflow.trace`, or an MLflow autolog integration, starts its
175
+ span through MLflow's context rather than through LangChain's callbacks. The tracer
176
+ publishes each live span into that context, so those spans nest under whichever node is
177
+ running instead of opening a root trace of their own:
178
+
179
+ ```python
180
+ import mlflow
181
+
182
+ @mlflow.trace(span_type="TOOL", name="lakebase.execute_query")
183
+ def execute_query(sql: str): ...
184
+ ```
185
+
186
+ Calling that from a graph node produces one trace, with the tool span parented to the
187
+ node. Without it you get a second root trace per call, carrying none of the root tags,
188
+ session metadata or cost rollup the tracer writes.
189
+
190
+ Attaching the context requires the tracer's callbacks to run in the caller's own
191
+ context, so `run_inline` is set on the handler whenever propagation is on. Only
192
+ LangChain's async callback manager reads that flag: without it an async run sends the
193
+ handler through an executor with a copied context, which the node body never sees. The
194
+ sync manager calls handlers in the caller's thread either way.
195
+
196
+ The tracer still parents its own spans from LangChain's run tree, never from the
197
+ context, so concurrent branches of a graph cannot cross-parent. Runs that overlap
198
+ rather than nest, two `stream(...)` generators created before either is consumed,
199
+ release their attaches in the order OpenTelemetry requires: a run that ends while
200
+ another is still stacked above it holds its attach until that one ends, so a finished
201
+ span is never restored over a live one.
202
+
203
+ The price of never releasing out of order is that a run whose end callback never
204
+ fires pins the attaches beneath it. A node that abandons an inner `stream(...)`
205
+ part-consumed leaves its span active for the life of that context, and spans started
206
+ there parent to a run that has already finished. Sweeping the stack when the
207
+ outermost run ends would clear it and bring back the overlapping-run bug above, so
208
+ the tracer reports it instead: past a depth no real graph reaches, it logs a warning
209
+ naming the count.
210
+
211
+ Only one handler attaches per run. Attach a second context-propagating tracer to the
212
+ same graph and it records its own spans as before, but leaves the context to the
213
+ first.
214
+
215
+ Three limits are worth knowing:
216
+
217
+ - It only works in one direction. A root run is always created without a parent and
218
+ without consulting the context, so a graph invoked from *inside* an `@mlflow.trace`
219
+ function still opens its own root trace rather than nesting into the caller's. This
220
+ is the one place the TypeScript package behaves differently: MLflow's JavaScript
221
+ `startSpan` falls back to the active context, so a graph there does nest into an
222
+ enclosing span.
223
+ - Foreign spans are recorded by MLflow, not by this tracer, so neither the text cap
224
+ nor the multimodal reference substitution applies to them. A traced query helper
225
+ returning a large rowset can push a trace back over the inline read limit that
226
+ `max_string_chars` was set to stay under.
227
+ - The tracer's callbacks run inline on the caller's thread, which on an async graph
228
+ means the event loop. Keep `content_ref_resolver` fast and non-blocking.
229
+ - Another handler on the same run that also publishes to MLflow's context, which in
230
+ practice means `MlflowLangchainTracer` from `mlflow.langchain.autolog()`, cannot be
231
+ coordinated with. LangChain calls the two handlers' end hooks in registration order,
232
+ so the attaches nest but the releases run front to back, and whichever releases last
233
+ restores an ended span. Two of MLflow's own tracers on one run leak the same way
234
+ with this tracer absent, so no release order fixes it here. The tracer warns once
235
+ when it detects the conflict and names both settings to turn off; turn one of them
236
+ off. Issue #57 tracks it.
237
+
238
+ To turn the whole behaviour off, pass `context_propagation=False` or set
239
+ `DATABRICKS_TRACING_CONTEXT_PROPAGATION=false` (`0`, `no` and `off` also work). That
240
+ also restores LangChain's default callback dispatch.
241
+
242
+ The TypeScript package reaches the same outcome through `getActiveTracingSpan()` and
243
+ `withTracingContext()`, named at the call site, because Node has no ambient
244
+ equivalent. See [`../typescript/README.md`](../typescript/README.md).
245
+
171
246
  ### Autolog fallback
172
247
 
173
248
  `mode="autolog"` wires MLflow's built-in LangChain autologging plus compatibility patches. It
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "insightfactory-databricks-langgraph-tracer"
3
- version = "1.0.0.dev15"
3
+ version = "1.0.1.dev18"
4
4
  description = "LangGraph tracing for Databricks MLflow"
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.12,<3.13"
@@ -161,6 +161,7 @@ def configure_databricks_tracing(
161
161
  sql_warehouse_id: str | None = None,
162
162
  content_ref_resolver: "ContentRefResolver | None" = None,
163
163
  max_string_chars: int | None = None,
164
+ context_propagation: bool | None = None,
164
165
  ) -> TracingConfiguration:
165
166
  """Bootstrap Databricks MLflow tracing and return the active configuration.
166
167
 
@@ -178,6 +179,19 @@ def configure_databricks_tracing(
178
179
  not just multimodal bytes, can push a trace past the SQL inline read limit);
179
180
  falls back to ``DATABRICKS_TRACING_MAX_STRING_CHARS``, else off. Like the
180
181
  resolver, it applies only under ``mode="tracer"``.
182
+
183
+ ``context_propagation`` (default on) publishes each live span into MLflow's
184
+ context, so spans started outside this tracer — an ``@mlflow.trace``-decorated
185
+ helper, an MLflow autolog integration — nest under the running graph instead of
186
+ opening their own root trace (issue #52); falls back to
187
+ ``DATABRICKS_TRACING_CONTEXT_PROPAGATION``, which turns it off for ``false`` /
188
+ ``0`` / ``no`` / ``off`` (case-insensitive) and is otherwise ignored.
189
+
190
+ ``mode`` picks between two alternatives, never both: ``"tracer"`` (the default)
191
+ attaches this package's ``BaseTracer`` for the full shared schema, while
192
+ ``"autolog"`` installs MLflow's LangChain autologging and attaches no callback.
193
+ The three options above shape only the ``"tracer"`` path. Autolog already traces
194
+ through MLflow's own context, so it needs no propagation.
181
195
  """
182
196
  global _bootstrap_done, _active_config
183
197
  # explicit disable path (design Decision 3): enabled=False OR skip-guard -> no-op
@@ -300,6 +314,7 @@ def configure_databricks_tracing(
300
314
  source=source,
301
315
  content_ref_resolver=content_ref_resolver,
302
316
  max_string_chars=max_string_chars,
317
+ context_propagation=context_propagation,
303
318
  )
304
319
  ]
305
320
  _active_config = TracingConfiguration(
@@ -55,7 +55,9 @@ import logging
55
55
  import math
56
56
  import os
57
57
  from collections.abc import Callable
58
+ from contextvars import ContextVar
58
59
  from dataclasses import dataclass
60
+ from typing import Any
59
61
 
60
62
  import mlflow
61
63
  from langchain_core.tracers.base import BaseTracer
@@ -84,6 +86,23 @@ from .cost import (
84
86
  qualify_deployment_name,
85
87
  )
86
88
 
89
+ # OpenTelemetry context propagation (issue #52). These live on a private MLflow
90
+ # module path, so the import is guarded: a future relocation degrades to "no context
91
+ # propagation" rather than breaking every consumer's import of this package. The
92
+ # names are annotated up front so both branches share one optional type. The attach
93
+ # token is an MLflow-internal union (a contextvars Token, or a pair of them), so it
94
+ # is carried as Any rather than restated here.
95
+ _set_span_in_context: "Callable[[LiveSpan], Any] | None"
96
+ _detach_span_from_context: "Callable[[Any], Any] | None"
97
+ try:
98
+ from mlflow.tracing.provider import detach_span_from_context, set_span_in_context
99
+
100
+ _set_span_in_context = set_span_in_context
101
+ _detach_span_from_context = detach_span_from_context
102
+ except ImportError: # pragma: no cover - depends on the installed mlflow layout
103
+ _set_span_in_context = None
104
+ _detach_span_from_context = None
105
+
87
106
  logger = logging.getLogger(__name__)
88
107
 
89
108
  # Pydantic BaseModel classes whose models MLflow expands when serializing. Both v2
@@ -160,6 +179,88 @@ def _resolve_max_string_chars(explicit: int | None) -> int | None:
160
179
  return value if value > 0 else None
161
180
 
162
181
 
182
+ # ---------------------------------------------------------------------------
183
+ # OpenTelemetry context propagation (issue #52).
184
+ #
185
+ # Spans are created with ``mlflow.start_span_no_context`` and parented explicitly
186
+ # from ``_span_map``, which keeps LangChain's callback tree authoritative and makes
187
+ # the tracer immune to context leaking between concurrent graph runs. The cost is
188
+ # that MLflow's *own* context is never populated, so a span started by anything
189
+ # other than this tracer — an ``@mlflow.trace`` tool, an MLflow autolog integration
190
+ # — finds no active parent and opens a second root trace. A consumer tracing its
191
+ # database helpers that way gets one trace per query alongside the graph trace,
192
+ # none of them carrying the root tags or the cost rollup.
193
+ #
194
+ # Publishing each live span into MLflow's context closes that gap: the tracer still
195
+ # parents its own spans from ``_span_map``, and everything else nests under whichever
196
+ # span is active. The attach token is detached in ``_end_trace``'s finally block, so
197
+ # a failure anywhere in span finalization cannot leak it.
198
+ #
199
+ # The attaches live in this context, innermost last. Two things make a plain
200
+ # attach-here / release-there pair wrong, and both were observed:
201
+ #
202
+ # 1. Two context-propagating tracers on one run. LangChain invokes every handler in
203
+ # registration order, so both attach and both release front-to-back, and the
204
+ # second release restores the context from before the *first* attach, leaving the
205
+ # first tracer's already-ended span active.
206
+ # 2. Runs that overlap rather than nest. Two ``graph.stream()`` iterators started
207
+ # before either is consumed end in consumption order, so the releases again run
208
+ # front-to-back and the finished run's span is restored over the live one.
209
+ #
210
+ # Either way a later ``@mlflow.trace`` call joins a finalized trace. This stack fixes
211
+ # both: a run already on it never attaches twice, and a run that ends out of order is
212
+ # only marked, with its token released once every attach above it has ended. Releases
213
+ # therefore always run innermost-first, which is the order OpenTelemetry requires.
214
+ # Entries are replaced rather than reset, so a token never has to be valid in the
215
+ # context that releases it.
216
+ #
217
+ # The price of never releasing out of order is head-of-line blocking: a run whose
218
+ # ``_end_trace`` never fires pins every attach beneath it, because the drain loop
219
+ # stops at the first entry that has not ended. A node that abandons an inner
220
+ # ``stream(...)`` part-consumed, or a generator finalized by the garbage collector in
221
+ # some other context, leaves its span active for the life of that context. Sweeping
222
+ # the stack when the bottom entry ends would fix that and reintroduce case 2, so the
223
+ # blocking is the deliberate trade. ``_MAX_ATTACHMENT_DEPTH`` makes it diagnosable.
224
+ @dataclass(frozen=True)
225
+ class _Attachment:
226
+ """One live context attach: the run that owns it and the token that releases it."""
227
+
228
+ run_id: str
229
+ span_id: str
230
+ token: Any
231
+ ended: bool = False
232
+
233
+
234
+ _attachments: ContextVar[tuple[_Attachment, ...]] = ContextVar(
235
+ "databricks_tracer_attachments", default=()
236
+ )
237
+
238
+ # Depth at which the stack is reported as stuck rather than deep. A graph nests a few
239
+ # runs, not dozens, so anything past this is a run that never ended pinning the ones
240
+ # below it (see the head-of-line note above) rather than genuine depth.
241
+ _MAX_ATTACHMENT_DEPTH = 64
242
+
243
+ _CONTEXT_PROPAGATION_ENV = "DATABRICKS_TRACING_CONTEXT_PROPAGATION"
244
+ # Opt-*out* polarity, so this cannot reuse the ``== "true"`` idiom the other boolean
245
+ # env vars in this package use (see mlflow_utils).
246
+ _CONTEXT_PROPAGATION_FALSE_VALUES = frozenset({"false", "0", "no", "off"})
247
+
248
+
249
+ def _resolve_context_propagation(explicit: bool | None) -> bool:
250
+ """Resolve context propagation: explicit arg wins, else the env var, else on.
251
+
252
+ Only an explicit false-y value ("false" / "0" / "no" / "off", case-insensitive)
253
+ turns it off, so an unset, empty or unrecognized variable leaves the default in
254
+ place. Matches ``_resolve_max_string_chars``, which also ignores an empty value.
255
+ """
256
+ if explicit is not None:
257
+ return explicit
258
+ raw = os.getenv(_CONTEXT_PROPAGATION_ENV)
259
+ if not raw:
260
+ return True
261
+ return raw.strip().lower() not in _CONTEXT_PROPAGATION_FALSE_VALUES
262
+
263
+
163
264
  def _truncate_long_string(value: str, max_string_chars: int | None):
164
265
  """Replace an over-cap string leaf with a compact placeholder; else return it as-is.
165
266
 
@@ -703,6 +804,7 @@ class DatabricksLangGraphTracer(BaseTracer):
703
804
  source: str = "langgraph",
704
805
  content_ref_resolver: ContentRefResolver | None = None,
705
806
  max_string_chars: int | None = None,
807
+ context_propagation: bool | None = None,
706
808
  **kw,
707
809
  ):
708
810
  # ``content_ref_resolver`` maps each multimodal content part to the reference
@@ -729,7 +831,39 @@ class DatabricksLangGraphTracer(BaseTracer):
729
831
  self._source = source
730
832
  self._content_ref_resolver = content_ref_resolver
731
833
  self._max_string_chars = _resolve_max_string_chars(max_string_chars)
834
+ # ``context_propagation`` publishes each live span into MLflow's context so
835
+ # spans started outside this tracer (an ``@mlflow.trace`` tool, an autolog
836
+ # integration) nest under the graph instead of opening their own root trace
837
+ # (issue #52); falls back to DATABRICKS_TRACING_CONTEXT_PROPAGATION, else on.
838
+ # Forced off when the MLflow context helpers are unavailable.
839
+ self._context_propagation = _resolve_context_propagation(
840
+ context_propagation
841
+ ) and (
842
+ _set_span_in_context is not None and _detach_span_from_context is not None
843
+ )
844
+ # Only the *async* callback manager consults run_inline: a sync handler that
845
+ # does not set it is dispatched through run_in_executor with a copied context,
846
+ # so an attach inside it never reaches the node body. langchain-core 1.x
847
+ # defaults BaseTracer to False, which is why the async path needs this. The
848
+ # sync manager calls handlers in the caller's thread either way (verified
849
+ # against langchain-core 1.6). Set per instance, matching autolog_compat, so
850
+ # opting out of context propagation also opts out of the dispatch change.
851
+ #
852
+ # Inline means this handler runs on the event loop rather than a worker
853
+ # thread. It only manipulates in-memory span state, but it also walks the
854
+ # payload for _externalize_content_refs and calls the consumer's
855
+ # content_ref_resolver — which must therefore be fast and non-blocking.
856
+ self.run_inline = self._context_propagation
732
857
  self._span_map: dict[str, LiveSpan] = {}
858
+ # Runs this tracer attached to the context. The tokens themselves live on the
859
+ # _attachments stack; this only records which entries are ours to end, so a
860
+ # second tracer on the same run does not end the first one's attach. Named
861
+ # apart from the stack on purpose: this one is instance-local, the stack is
862
+ # context-local and shared across tracers.
863
+ self._runs_i_attached: set[str] = set()
864
+ # One warning per tracer for each of the two context hazards.
865
+ self._warned_shared_context = False
866
+ self._warned_stack_depth = False
733
867
  # Per-root cost accumulators, keyed by root run id.
734
868
  self._aggregates: dict[str, TraceAggregate] = {}
735
869
  # Run id to root run id, so child LLM spans find their root.
@@ -766,6 +900,7 @@ class DatabricksLangGraphTracer(BaseTracer):
766
900
  inputs=inputs,
767
901
  )
768
902
  self._span_map[str(run.id)] = span
903
+ self._attach_context(run, span)
769
904
  if run.parent_run_id is None: # ROOT
770
905
  self._aggregates[str(run.id)] = TraceAggregate(
771
906
  root_trace_id=span.trace_id
@@ -836,12 +971,127 @@ class DatabricksLangGraphTracer(BaseTracer):
836
971
  except Exception:
837
972
  logger.debug("span end failed", exc_info=True)
838
973
  finally:
974
+ # Detach after every span write, so the context unwinds in the reverse
975
+ # order it was attached (child before parent) and a failure above still
976
+ # releases it.
977
+ self._detach_context(run)
839
978
  # Unconditionally clear this run's root mapping (self-map for a root,
840
979
  # own->root lookup for a child). Never delete the *value*: sibling
841
980
  # children of a still-running root rely on its self-mapping.
842
981
  self._root_by_run.pop(str(run.id), None)
843
982
  super()._end_trace(run)
844
983
 
984
+ def _attach_context(self, run: Run, span: LiveSpan) -> None:
985
+ """Publish ``span`` into MLflow's context as the active span (issue #52).
986
+
987
+ Pushes the release token onto the context-local stack and records the run id on
988
+ this tracer, which is what marks the entry as this tracer's to end. A failure is
989
+ logged and swallowed: losing context propagation degrades nesting for foreign
990
+ spans, it must never break the trace this tracer is recording.
991
+ """
992
+ if not self._context_propagation or _set_span_in_context is None:
993
+ return
994
+ run_id = str(run.id)
995
+ stack = _attachments.get()
996
+ if any(entry.run_id == run_id for entry in stack):
997
+ return # another handler on this run already owns the attach
998
+ try:
999
+ token = _set_span_in_context(span)
1000
+ except Exception:
1001
+ logger.debug("attach span context failed", exc_info=True)
1002
+ return
1003
+ # Pushed only after the attach succeeds, so a failure here leaves the next
1004
+ # handler free to try rather than suppressing it for a token that never was.
1005
+ entries = (*stack, _Attachment(run_id, span.span_id, token))
1006
+ _attachments.set(entries)
1007
+ self._runs_i_attached.add(run_id)
1008
+ if len(entries) > _MAX_ATTACHMENT_DEPTH and not self._warned_stack_depth:
1009
+ self._warned_stack_depth = True
1010
+ logger.warning(
1011
+ "%d spans are attached to MLflow's context in this context, which "
1012
+ "means a run ended without its end callback firing and is holding "
1013
+ "the attaches below it. Spans started here will parent to a run that "
1014
+ "has already finished.",
1015
+ len(entries),
1016
+ )
1017
+
1018
+ def _detach_context(self, run: Run) -> None:
1019
+ """Mark this run's attach as ended and release every attach now on top.
1020
+
1021
+ An attach can only be released once nothing is stacked above it, so a run
1022
+ that ends out of order is marked and left in place; its token goes when the
1023
+ run above it ends. That keeps the releases innermost-first, which is what
1024
+ OpenTelemetry's token contract requires, and it never restores an ended span
1025
+ over a live one.
1026
+
1027
+ A failed release is logged at warning, unlike a failed attach: a failed attach
1028
+ costs nesting for one run, while a failed release leaves an *ended* span active
1029
+ in a context that outlives the run (a reused server thread, a long-lived event
1030
+ loop), so later spans in that context would parent to a finalized span.
1031
+ OpenTelemetry's own ``detach`` swallows most mismatches and logs them itself,
1032
+ so this is a backstop rather than the only signal.
1033
+ """
1034
+ run_id = str(run.id)
1035
+ if run_id not in self._runs_i_attached:
1036
+ return
1037
+ # Discarded first, so a failure below cannot leave the run looking attached.
1038
+ self._runs_i_attached.discard(run_id)
1039
+ if _detach_span_from_context is None:
1040
+ return
1041
+ entries = list(_attachments.get())
1042
+ index = next(
1043
+ (i for i, entry in enumerate(entries) if entry.run_id == run_id), None
1044
+ )
1045
+ if index is None:
1046
+ # Attached in a different context than the one ending the run, so this
1047
+ # context has nothing to release. The attach's own context still holds the
1048
+ # token and unwinds normally when it is discarded.
1049
+ return
1050
+ self._warn_if_context_shared(entries)
1051
+ entries[index] = dataclasses.replace(entries[index], ended=True)
1052
+ while entries and entries[-1].ended:
1053
+ entry = entries.pop()
1054
+ try:
1055
+ _detach_span_from_context(entry.token)
1056
+ except Exception:
1057
+ logger.warning("detach span context failed", exc_info=True)
1058
+ _attachments.set(tuple(entries))
1059
+
1060
+ def _warn_if_context_shared(self, entries: list[_Attachment]) -> None:
1061
+ """Warn once if another context publisher has attached above this tracer.
1062
+
1063
+ The stack this tracer keeps only covers its own attaches. Another handler on
1064
+ the same run that also publishes to MLflow's context, ``MlflowLangchainTracer``
1065
+ from ``mlflow.langchain.autolog()`` being the one in practice, attaches its own
1066
+ token, and LangChain calls the two handlers' end hooks in registration order.
1067
+ Whichever releases last restores the span it recorded at attach time, which is
1068
+ by then an ended span, so a later ``@mlflow.trace`` call joins a finalized
1069
+ trace.
1070
+
1071
+ No release order fixes this from inside one handler, and it is not specific to
1072
+ this tracer: two ``MlflowLangchainTracer`` instances on one run leak the same
1073
+ way with this tracer absent. MLflow's own answer is the global
1074
+ ``_should_attach_span_to_context`` opt-out rather than a coordination protocol.
1075
+ So this warns rather than repairs, and names the setting to turn off. Issue #57.
1076
+
1077
+ An enclosing span the caller opened around the graph is not this case: it sits
1078
+ *below* this tracer's attaches, so the stack top still matches.
1079
+ """
1080
+ if self._warned_shared_context or not entries:
1081
+ return
1082
+ active = mlflow.get_current_active_span()
1083
+ if active is not None and active.span_id == entries[-1].span_id:
1084
+ return
1085
+ self._warned_shared_context = True
1086
+ logger.warning(
1087
+ "another handler has attached a span to MLflow's context inside this run, "
1088
+ "so the context cannot unwind in order and a span started after the run "
1089
+ "may join a finalized trace. Disable context propagation on one of them: "
1090
+ "%s=false for this tracer, or mlflow.langchain.autolog(log_traces=False) "
1091
+ "for MLflow's own LangChain tracer.",
1092
+ _CONTEXT_PROPAGATION_ENV,
1093
+ )
1094
+
845
1095
  def _write_root_trace_fields(self, run: Run, span: LiveSpan) -> None:
846
1096
  """Write the root trace tags and session/user metadata.
847
1097