WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-12 04:37:51 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat: Honor traceable config

Browse Source
This commit is contained in:
William Fu-Hinthorn
2026-01-07 14:55:53 -08:00
parent 8ccead9560
commit 6b17b12ae4
3 changed files with 788 additions and 191 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/_internal/_runnable.py
+286 -185
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import asyncio
import enum
import inspect
import logging
import sys
import warnings
from collections.abc import (
@@ -61,6 +62,38 @@ try:
except ImportError:
_StreamingCallbackHandler = None # type: ignore
logger = logging.getLogger(__name__)
def _safe_process_inputs(processor: Callable[[Any], Any] | None, inputs: Any) -> Any:
"""Safely process trace inputs, returning error placeholder on failure.
This prevents PII leakage if a filter function crashes - we return an error
indicator instead of the raw (potentially sensitive) inputs.
"""
if processor is None:
return inputs
try:
return processor(inputs)
except Exception:
logger.warning("trace_inputs filter failed", exc_info=True)
return {"error": "<trace_inputs processing failed>"}
def _safe_process_outputs(processor: Callable[[Any], Any] | None, outputs: Any) -> Any:
"""Safely process trace outputs, returning error placeholder on failure.
This prevents PII leakage if a filter function crashes - we return an error
indicator instead of the raw (potentially sensitive) outputs.
"""
if processor is None:
return outputs
try:
return processor(outputs)
except Exception:
logger.warning("trace_outputs filter failed", exc_info=True)
return {"error": "<trace_outputs processing failed>"}
def _set_config_context(
config: RunnableConfig, run: Any = None
@@ -544,12 +577,17 @@ class RunnableSeq(Runnable):
*steps: RunnableLike,
name: str | None = None,
trace_inputs: Callable[[Any], Any] | None = None,
trace_outputs: Callable[[Any], Any] | None = None,
trace: bool = True,
) -> None:
"""Create a new RunnableSeq.
Args:
steps: The steps to include in the sequence.
name: The name of the `Runnable`.
trace_inputs: Optional function to transform inputs before tracing.
trace_outputs: Optional function to transform outputs before tracing.
trace: Whether to trace this sequence. Defaults to True.
Raises:
ValueError: If the sequence has less than 2 steps.
@@ -569,6 +607,8 @@ class RunnableSeq(Runnable):
self.steps = steps_flat
self.name = name
self.trace_inputs = trace_inputs
self.trace_outputs = trace_outputs
self.trace = trace
def __or__(
self,
@@ -625,43 +665,56 @@ class RunnableSeq(Runnable):
) -> Any:
if config is None:
config = ensure_config()
# setup callbacks and context
callback_manager = get_callback_manager_for_config(config)
# start the root run
run_manager = callback_manager.on_chain_start(
None,
self.trace_inputs(input) if self.trace_inputs is not None else input,
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
# invoke all steps in sequence
try:
for i, step in enumerate(self.steps):
# mark each step as a child run
config = patch_config(
config, callbacks=run_manager.get_child(f"seq:step:{i + 1}")
)
# 1st step is the actual node,
# others are writers which don't need to be run in context
if i == 0:
# get the run object
for h in run_manager.handlers:
if isinstance(h, LangChainTracer):
run = h.run_map.get(str(run_manager.run_id))
break
if self.trace:
# setup callbacks and context
callback_manager = get_callback_manager_for_config(config)
# start the root run
run_manager = callback_manager.on_chain_start(
None,
_safe_process_inputs(self.trace_inputs, input),
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
# invoke all steps in sequence
try:
for i, step in enumerate(self.steps):
# mark each step as a child run
config = patch_config(
config, callbacks=run_manager.get_child(f"seq:step:{i + 1}")
)
# 1st step is the actual node,
# others are writers which don't need to be run in context
if i == 0:
# get the run object
for h in run_manager.handlers:
if isinstance(h, LangChainTracer):
run = h.run_map.get(str(run_manager.run_id))
break
else:
run = None
# run in context
with set_config_context(config, run) as context:
input = context.run(step.invoke, input, config, **kwargs)
else:
run = None
# run in context
with set_config_context(config, run) as context:
input = context.run(step.invoke, input, config, **kwargs)
else:
input = step.invoke(input, config)
# finish the root run
except BaseException as e:
run_manager.on_chain_error(e)
raise
input = step.invoke(input, config)
# finish the root run
except BaseException as e:
run_manager.on_chain_error(e)
raise
else:
run_manager.on_chain_end(
_safe_process_outputs(self.trace_outputs, input)
)
return input
else:
run_manager.on_chain_end(input)
# No tracing - just execute the steps directly
for i, step in enumerate(self.steps):
input = (
step.invoke(input, config, **kwargs)
if i == 0
else step.invoke(input, config)
)
return input
async def ainvoke(
@@ -672,49 +725,62 @@ class RunnableSeq(Runnable):
) -> Any:
if config is None:
config = ensure_config()
# setup callbacks
callback_manager = get_async_callback_manager_for_config(config)
# start the root run
run_manager = await callback_manager.on_chain_start(
None,
self.trace_inputs(input) if self.trace_inputs is not None else input,
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
# invoke all steps in sequence
try:
for i, step in enumerate(self.steps):
# mark each step as a child run
config = patch_config(
config, callbacks=run_manager.get_child(f"seq:step:{i + 1}")
)
# 1st step is the actual node,
# others are writers which don't need to be run in context
if i == 0:
if ASYNCIO_ACCEPTS_CONTEXT:
# get the run object
for h in run_manager.handlers:
if isinstance(h, LangChainTracer):
run = h.run_map.get(str(run_manager.run_id))
break
if self.trace:
# setup callbacks
callback_manager = get_async_callback_manager_for_config(config)
# start the root run
run_manager = await callback_manager.on_chain_start(
None,
_safe_process_inputs(self.trace_inputs, input),
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
# invoke all steps in sequence
try:
for i, step in enumerate(self.steps):
# mark each step as a child run
config = patch_config(
config, callbacks=run_manager.get_child(f"seq:step:{i + 1}")
)
# 1st step is the actual node,
# others are writers which don't need to be run in context
if i == 0:
if ASYNCIO_ACCEPTS_CONTEXT:
# get the run object
for h in run_manager.handlers:
if isinstance(h, LangChainTracer):
run = h.run_map.get(str(run_manager.run_id))
break
else:
run = None
# run in context
with set_config_context(config, run) as context:
input = await asyncio.create_task(
step.ainvoke(input, config, **kwargs),
context=context,
)
else:
run = None
# run in context
with set_config_context(config, run) as context:
input = await asyncio.create_task(
step.ainvoke(input, config, **kwargs), context=context
)
input = await step.ainvoke(input, config, **kwargs)
else:
input = await step.ainvoke(input, config, **kwargs)
input = await step.ainvoke(input, config)
# finish the root run
except BaseException as e:
await run_manager.on_chain_error(e)
raise
else:
await run_manager.on_chain_end(
_safe_process_outputs(self.trace_outputs, input)
)
return input
else:
# No tracing - just execute the steps directly
for i, step in enumerate(self.steps):
if i == 0:
input = await step.ainvoke(input, config, **kwargs)
else:
input = await step.ainvoke(input, config)
# finish the root run
except BaseException as e:
await run_manager.on_chain_error(e)
raise
else:
await run_manager.on_chain_end(input)
return input
def stream(
@@ -725,80 +791,17 @@ class RunnableSeq(Runnable):
) -> Iterator[Any]:
if config is None:
config = ensure_config()
# setup callbacks
callback_manager = get_callback_manager_for_config(config)
# start the root run
run_manager = callback_manager.on_chain_start(
None,
self.trace_inputs(input) if self.trace_inputs is not None else input,
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
# get the run object
for h in run_manager.handlers:
if isinstance(h, LangChainTracer):
run = h.run_map.get(str(run_manager.run_id))
break
else:
run = None
# create first step config
config = patch_config(
config,
callbacks=run_manager.get_child(f"seq:step:{1}"),
)
# run all in context
with set_config_context(config, run) as context:
try:
# stream the last steps
# transform the input stream of each step with the next
# steps that don't natively support transforming an input stream will
# buffer input in memory until all available, and then start emitting output
for idx, step in enumerate(self.steps):
if idx == 0:
iterator = step.stream(input, config, **kwargs)
else:
config = patch_config(
config,
callbacks=run_manager.get_child(f"seq:step:{idx + 1}"),
)
iterator = step.transform(iterator, config)
# populates streamed_output in astream_log() output if needed
if _StreamingCallbackHandler is not None:
for h in run_manager.handlers:
if isinstance(h, _StreamingCallbackHandler):
iterator = h.tap_output_iter(run_manager.run_id, iterator)
# consume into final output
output = context.run(_consume_iter, iterator)
# sequence doesn't emit output, yield to mark as generator
yield
except BaseException as e:
run_manager.on_chain_error(e)
raise
else:
run_manager.on_chain_end(output)
async def astream(
self,
input: Input,
config: RunnableConfig | None = None,
**kwargs: Any | None,
) -> AsyncIterator[Any]:
if config is None:
config = ensure_config()
# setup callbacks
callback_manager = get_async_callback_manager_for_config(config)
# start the root run
run_manager = await callback_manager.on_chain_start(
None,
self.trace_inputs(input) if self.trace_inputs is not None else input,
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
# stream the last steps
# transform the input stream of each step with the next
# steps that don't natively support transforming an input stream will
# buffer input in memory until all available, and then start emitting output
if ASYNCIO_ACCEPTS_CONTEXT:
if self.trace:
# setup callbacks
callback_manager = get_callback_manager_for_config(config)
# start the root run
run_manager = callback_manager.on_chain_start(
None,
_safe_process_inputs(self.trace_inputs, input),
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
# get the run object
for h in run_manager.handlers:
if isinstance(h, LangChainTracer):
@@ -813,18 +816,133 @@ class RunnableSeq(Runnable):
)
# run all in context
with set_config_context(config, run) as context:
try:
# stream the last steps
# transform the input stream of each step with the next
# steps that don't natively support transforming an input stream will
# buffer input in memory until all available, and then start emitting output
for idx, step in enumerate(self.steps):
if idx == 0:
iterator = step.stream(input, config, **kwargs)
else:
config = patch_config(
config,
callbacks=run_manager.get_child(f"seq:step:{idx + 1}"),
)
iterator = step.transform(iterator, config)
# populates streamed_output in astream_log() output if needed
if _StreamingCallbackHandler is not None:
for h in run_manager.handlers:
if isinstance(h, _StreamingCallbackHandler):
iterator = h.tap_output_iter(
run_manager.run_id, iterator
)
# consume into final output
output = context.run(_consume_iter, iterator)
# sequence doesn't emit output, yield to mark as generator
yield
except BaseException as e:
run_manager.on_chain_error(e)
raise
else:
run_manager.on_chain_end(
_safe_process_outputs(self.trace_outputs, output)
)
else:
# No tracing - just execute the steps directly
for idx, step in enumerate(self.steps):
if idx == 0:
iterator = step.stream(input, config, **kwargs)
else:
iterator = step.transform(iterator, config)
# consume into final output
_consume_iter(iterator)
yield
async def astream(
self,
input: Input,
config: RunnableConfig | None = None,
**kwargs: Any | None,
) -> AsyncIterator[Any]:
if config is None:
config = ensure_config()
if self.trace:
# setup callbacks
callback_manager = get_async_callback_manager_for_config(config)
# start the root run
run_manager = await callback_manager.on_chain_start(
None,
_safe_process_inputs(self.trace_inputs, input),
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
# stream the last steps
# transform the input stream of each step with the next
# steps that don't natively support transforming an input stream will
# buffer input in memory until all available, and then start emitting output
if ASYNCIO_ACCEPTS_CONTEXT:
# get the run object
for h in run_manager.handlers:
if isinstance(h, LangChainTracer):
run = h.run_map.get(str(run_manager.run_id))
break
else:
run = None
# create first step config
config = patch_config(
config,
callbacks=run_manager.get_child(f"seq:step:{1}"),
)
# run all in context
with set_config_context(config, run) as context:
try:
async with AsyncExitStack() as stack:
for idx, step in enumerate(self.steps):
if idx == 0:
aiterator = step.astream(input, config, **kwargs)
else:
config = patch_config(
config,
callbacks=run_manager.get_child(
f"seq:step:{idx + 1}"
),
)
aiterator = step.atransform(aiterator, config)
if hasattr(aiterator, "aclose"):
stack.push_async_callback(aiterator.aclose)
# populates streamed_output in astream_log() output if needed
if _StreamingCallbackHandler is not None:
for h in run_manager.handlers:
if isinstance(h, _StreamingCallbackHandler):
aiterator = h.tap_output_aiter(
run_manager.run_id, aiterator
)
# consume into final output
output = await asyncio.create_task(
_consume_aiter(aiterator), context=context
)
# sequence doesn't emit output, yield to mark as generator
yield
except BaseException as e:
await run_manager.on_chain_error(e)
raise
else:
await run_manager.on_chain_end(
_safe_process_outputs(self.trace_outputs, output)
)
else:
try:
async with AsyncExitStack() as stack:
for idx, step in enumerate(self.steps):
config = patch_config(
config,
callbacks=run_manager.get_child(f"seq:step:{idx + 1}"),
)
if idx == 0:
aiterator = step.astream(input, config, **kwargs)
else:
config = patch_config(
config,
callbacks=run_manager.get_child(
f"seq:step:{idx + 1}"
),
)
aiterator = step.atransform(aiterator, config)
if hasattr(aiterator, "aclose"):
stack.push_async_callback(aiterator.aclose)
@@ -836,46 +954,29 @@ class RunnableSeq(Runnable):
run_manager.run_id, aiterator
)
# consume into final output
output = await asyncio.create_task(
_consume_aiter(aiterator), context=context
)
output = await _consume_aiter(aiterator)
# sequence doesn't emit output, yield to mark as generator
yield
except BaseException as e:
await run_manager.on_chain_error(e)
raise
else:
await run_manager.on_chain_end(output)
await run_manager.on_chain_end(
_safe_process_outputs(self.trace_outputs, output)
)
else:
try:
async with AsyncExitStack() as stack:
for idx, step in enumerate(self.steps):
config = patch_config(
config,
callbacks=run_manager.get_child(f"seq:step:{idx + 1}"),
)
if idx == 0:
aiterator = step.astream(input, config, **kwargs)
else:
aiterator = step.atransform(aiterator, config)
if hasattr(aiterator, "aclose"):
stack.push_async_callback(aiterator.aclose)
# populates streamed_output in astream_log() output if needed
if _StreamingCallbackHandler is not None:
for h in run_manager.handlers:
if isinstance(h, _StreamingCallbackHandler):
aiterator = h.tap_output_aiter(
run_manager.run_id, aiterator
)
# consume into final output
output = await _consume_aiter(aiterator)
# sequence doesn't emit output, yield to mark as generator
yield
except BaseException as e:
await run_manager.on_chain_error(e)
raise
else:
await run_manager.on_chain_end(output)
# No tracing - just execute the steps directly
async with AsyncExitStack() as stack:
for idx, step in enumerate(self.steps):
if idx == 0:
aiterator = step.astream(input, config, **kwargs)
else:
aiterator = step.atransform(aiterator, config)
if hasattr(aiterator, "aclose"):
stack.push_async_callback(aiterator.aclose)
# consume into final output
await _consume_aiter(aiterator)
yield
def _consume_iter(it: Iterator[Any]) -> Any:
libs/langgraph/langgraph/pregel/_read.py
+80 -6
View File
@@ -7,6 +7,7 @@ from typing import (
)
from langchain_core.runnables import Runnable, RunnableConfig
from typing_extensions import TypedDict
from langgraph._internal._config import merge_configs
from langgraph._internal._constants import CONF, CONFIG_KEY_READ
@@ -20,6 +21,28 @@ READ_TYPE = Callable[[str | Sequence[str], bool], Any | dict[str, Any]]
INPUT_CACHE_KEY_TYPE = tuple[Callable[..., Any], tuple[str, ...]]
class TraceableConfig(TypedDict, total=False):
"""Configuration extracted from @traceable decorated functions."""
process_inputs: Callable[[Any], Any] | None
process_outputs: Callable[[Any], Any] | None
enabled: bool | None # None = honor external tracing context
def _validate_traceable_config(raw: Any) -> TraceableConfig | None:
"""Validate __traceable_config__ has expected structure.
Returns validated config dict or None if invalid.
"""
if not isinstance(raw, dict):
return None
return {
"process_inputs": raw.get("process_inputs"),
"process_outputs": raw.get("process_outputs"),
"enabled": raw.get("enabled"), # None means use external context
}
class ChannelRead(RunnableCallable):
"""Implements the logic for reading state from CONFIG_KEY_READ.
Usable both as a runnable as well as a static method to call imperatively."""
@@ -176,6 +199,7 @@ class PregelNode:
attrs = {**self.__dict__, **update}
# Drop the cached properties
attrs.pop("flat_writers", None)
attrs.pop("_traceable_config", None)
attrs.pop("node", None)
attrs.pop("input_cache_key", None)
return PregelNode(**attrs)
@@ -197,20 +221,70 @@ class PregelNode:
writers.pop()
return writers
@cached_property
def _traceable_config(self) -> TraceableConfig | None:
"""Extract @traceable config if bound wraps a traceable function.
Returns validated TraceableConfig with:
- process_inputs: Optional input processing function
- process_outputs: Optional output processing function
- enabled: Whether tracing is enabled (None = honor external context)
- __unwrapped__: The original unwrapped function (if available)
"""
if not isinstance(self.bound, RunnableCallable):
return None
func = self.bound.func or self.bound.afunc
if func is None:
return None
raw_config = getattr(func, "__traceable_config__", None)
config = _validate_traceable_config(raw_config)
if config is None:
return None
# Get the unwrapped original function (set by functools.wraps)
unwrapped = getattr(func, "__wrapped__", None)
if unwrapped:
config["__unwrapped__"] = unwrapped
return config
@cached_property
def node(self) -> Runnable[Any, Any] | None:
"""Get a runnable that combines `bound` and `writers`."""
from langgraph._internal._runnable import coerce_to_runnable
writers = self.flat_writers
if self.bound is DEFAULT_BOUND and not writers:
tc = self._traceable_config
# Build RunnableSeq kwargs from traceable config
seq_kwargs: dict[str, Any] = {}
if tc:
# Only disable tracing if explicitly set to False
# None means honor external context (keep default trace=True)
if tc.get("enabled") is False:
seq_kwargs["trace"] = False
# trace_inputs/trace_outputs map to @traceable's process_inputs/process_outputs
seq_kwargs["trace_inputs"] = tc.get("process_inputs")
seq_kwargs["trace_outputs"] = tc.get("process_outputs")
# Get the bound to use (unwrapped if traceable)
bound = self.bound
if tc and tc.get("__unwrapped__"):
# Replace bound with unwrapped function to avoid double-tracing
bound = coerce_to_runnable(tc["__unwrapped__"], name=None, trace=False)
if bound is DEFAULT_BOUND and not writers:
return None
elif self.bound is DEFAULT_BOUND and len(writers) == 1:
elif bound is DEFAULT_BOUND and len(writers) == 1:
return writers[0]
elif self.bound is DEFAULT_BOUND:
return RunnableSeq(*writers)
elif bound is DEFAULT_BOUND:
return RunnableSeq(*writers, **seq_kwargs)
elif writers:
return RunnableSeq(self.bound, *writers)
return RunnableSeq(bound, *writers, **seq_kwargs)
else:
return self.bound
return bound
@cached_property
def input_cache_key(self) -> INPUT_CACHE_KEY_TYPE:
libs/langgraph/tests/test_traceable_integration.py
+422
View File
@@ -0,0 +1,422 @@
"""Tests for @traceable integration with LangGraph nodes.
These tests verify that LangGraph properly reads `__traceable_config__` from
decorated node functions and applies the configuration to tracing.
Note: Some tests manually set `__traceable_config__` to work independently of
the langsmith package version installed. When langsmith exposes this attribute
via @traceable decorator, the manual setting can be removed.
"""
from __future__ import annotations
from typing import Any
import pytest
from typing_extensions import TypedDict
from langgraph.graph import StateGraph
from tests.fake_tracer import FakeTracer
pytestmark = pytest.mark.anyio
class SimpleState(TypedDict):
value: str
def _set_traceable_config(
func: Any,
*,
process_inputs: Any = None,
process_outputs: Any = None,
enabled: bool | None = None,
) -> Any:
"""Helper to set __traceable_config__ on a function.
This simulates what @traceable decorator should do when it exposes
the __traceable_config__ attribute.
Args:
func: The function to decorate.
process_inputs: Optional function to filter traced inputs.
process_outputs: Optional function to filter traced outputs.
enabled: Whether tracing is enabled. None means honor external context.
"""
setattr(
func,
"__traceable_config__",
{
"process_inputs": process_inputs,
"process_outputs": process_outputs,
"enabled": enabled,
},
)
return func
def test_traceable_config_process_inputs():
"""Test that __traceable_config__ with process_inputs filters inputs in trace."""
tracer = FakeTracer()
def filter_inputs(inputs: dict[str, Any]) -> dict[str, Any]:
"""Filter sensitive data from inputs."""
return {"value": "[REDACTED]"}
def my_node(state: SimpleState) -> SimpleState:
return {"value": f"processed_{state['value']}"}
# Manually set traceable config (simulating @traceable decorator)
_set_traceable_config(my_node, process_inputs=filter_inputs)
builder = StateGraph(SimpleState)
builder.add_node("my_node", my_node)
builder.add_edge("__start__", "my_node")
graph = builder.compile()
result = graph.invoke({"value": "secret_data"}, {"callbacks": [tracer]})
assert result == {"value": "processed_secret_data"}
# Find the node run (should be the RunnableSeq for my_node)
runs = tracer.flattened_runs()
node_runs = [r for r in runs if r.name == "my_node"]
assert len(node_runs) == 1, f"Expected 1 my_node run, got {len(node_runs)}"
# Verify inputs were filtered through process_inputs
node_run = node_runs[0]
assert node_run.inputs == {"value": "[REDACTED]"}, (
f"Expected inputs to be filtered, got {node_run.inputs}"
)
def test_traceable_config_process_outputs():
"""Test that __traceable_config__ with process_outputs filters outputs in trace."""
tracer = FakeTracer()
def filter_outputs(outputs: Any) -> Any:
"""Filter sensitive data from outputs."""
if isinstance(outputs, dict) and "value" in outputs:
return {"value": "[OUTPUT_REDACTED]"}
return outputs
def my_node(state: SimpleState) -> SimpleState:
return {"value": "secret_result"}
# Manually set traceable config
_set_traceable_config(my_node, process_outputs=filter_outputs)
builder = StateGraph(SimpleState)
builder.add_node("my_node", my_node)
builder.add_edge("__start__", "my_node")
graph = builder.compile()
result = graph.invoke({"value": "input"}, {"callbacks": [tracer]})
# Result should be unfiltered (actual execution)
assert result == {"value": "secret_result"}
# Find the node run
runs = tracer.flattened_runs()
node_runs = [r for r in runs if r.name == "my_node"]
assert len(node_runs) == 1, f"Expected 1 my_node run, got {len(node_runs)}"
# Verify outputs were filtered through process_outputs
node_run = node_runs[0]
assert node_run.outputs == {"value": "[OUTPUT_REDACTED]"}, (
f"Expected outputs to be filtered, got {node_run.outputs}"
)
def test_traceable_config_enabled_false():
"""Test that __traceable_config__ with enabled=False skips trace creation."""
tracer = FakeTracer()
def hidden_node(state: SimpleState) -> SimpleState:
return {"value": f"hidden_{state['value']}"}
# Manually set traceable config with enabled=False
_set_traceable_config(hidden_node, enabled=False)
def visible_node(state: SimpleState) -> SimpleState:
return {"value": f"visible_{state['value']}"}
builder = StateGraph(SimpleState)
builder.add_node("hidden_node", hidden_node)
builder.add_node("visible_node", visible_node)
builder.add_edge("__start__", "hidden_node")
builder.add_edge("hidden_node", "visible_node")
graph = builder.compile()
result = graph.invoke({"value": "test"}, {"callbacks": [tracer]})
assert result == {"value": "visible_hidden_test"}
# Verify runs
runs = tracer.flattened_runs()
run_names = [r.name for r in runs]
# visible_node should be traced
assert "visible_node" in run_names, f"Expected visible_node in {run_names}"
# hidden_node should NOT be traced (enabled=False)
assert "hidden_node" not in run_names, (
f"hidden_node should not be traced but found in {run_names}"
)
def test_regular_node_tracing_unchanged():
"""Test that regular nodes (without __traceable_config__) trace normally."""
tracer = FakeTracer()
def regular_node(state: SimpleState) -> SimpleState:
return {"value": f"regular_{state['value']}"}
builder = StateGraph(SimpleState)
builder.add_node("regular_node", regular_node)
builder.add_edge("__start__", "regular_node")
graph = builder.compile()
result = graph.invoke({"value": "test"}, {"callbacks": [tracer]})
assert result == {"value": "regular_test"}
# Verify the node is traced
runs = tracer.flattened_runs()
node_runs = [r for r in runs if r.name == "regular_node"]
assert len(node_runs) == 1, f"Expected 1 regular_node run, got {len(node_runs)}"
# Verify inputs and outputs are captured normally (not filtered)
node_run = node_runs[0]
assert "value" in node_run.inputs
assert node_run.inputs["value"] == "test"
def test_traceable_config_single_run_no_duplication():
"""Test that node with __traceable_config__ creates exactly one run."""
tracer = FakeTracer()
def my_node(state: SimpleState) -> SimpleState:
return {"value": f"result_{state['value']}"}
# Set traceable config (simulating @traceable decorator)
_set_traceable_config(my_node)
builder = StateGraph(SimpleState)
builder.add_node("my_node", my_node)
builder.add_edge("__start__", "my_node")
graph = builder.compile()
result = graph.invoke({"value": "test"}, {"callbacks": [tracer]})
assert result == {"value": "result_test"}
# Count runs with name my_node - should be exactly 1
runs = tracer.flattened_runs()
my_node_runs = [r for r in runs if r.name == "my_node"]
assert len(my_node_runs) == 1, (
f"Expected exactly 1 my_node run (no duplication), got {len(my_node_runs)}. "
f"All run names: {[r.name for r in runs]}"
)
async def test_traceable_config_async():
"""Test that async nodes with __traceable_config__ work correctly."""
tracer = FakeTracer()
def filter_inputs(inputs: dict[str, Any]) -> dict[str, Any]:
return {"value": "[ASYNC_REDACTED]"}
async def async_node(state: SimpleState) -> SimpleState:
return {"value": f"async_{state['value']}"}
# Set traceable config (simulating @traceable decorator)
_set_traceable_config(async_node, process_inputs=filter_inputs)
builder = StateGraph(SimpleState)
builder.add_node("async_node", async_node)
builder.add_edge("__start__", "async_node")
graph = builder.compile()
result = await graph.ainvoke({"value": "secret"}, {"callbacks": [tracer]})
assert result == {"value": "async_secret"}
# Verify inputs were filtered
runs = tracer.flattened_runs()
node_runs = [r for r in runs if r.name == "async_node"]
assert len(node_runs) == 1
assert node_runs[0].inputs == {"value": "[ASYNC_REDACTED]"}
def test_traceable_config_both_inputs_and_outputs():
"""Test __traceable_config__ with both process_inputs and process_outputs."""
tracer = FakeTracer()
def filter_inputs(inputs: dict[str, Any]) -> dict[str, Any]:
return {"value": "[INPUT_FILTERED]"}
def filter_outputs(outputs: Any) -> Any:
if isinstance(outputs, dict):
return {"value": "[OUTPUT_FILTERED]"}
return outputs
def my_node(state: SimpleState) -> SimpleState:
return {"value": f"result_{state['value']}"}
# Set traceable config with both filters
_set_traceable_config(
my_node, process_inputs=filter_inputs, process_outputs=filter_outputs
)
builder = StateGraph(SimpleState)
builder.add_node("my_node", my_node)
builder.add_edge("__start__", "my_node")
graph = builder.compile()
result = graph.invoke({"value": "secret"}, {"callbacks": [tracer]})
# Actual result should be unfiltered
assert result == {"value": "result_secret"}
# Trace should have filtered inputs and outputs
runs = tracer.flattened_runs()
node_runs = [r for r in runs if r.name == "my_node"]
assert len(node_runs) == 1
assert node_runs[0].inputs == {"value": "[INPUT_FILTERED]"}
assert node_runs[0].outputs == {"value": "[OUTPUT_FILTERED]"}
def test_traceable_config_process_inputs_error_handling():
"""Test that errors in process_inputs don't leak PII."""
tracer = FakeTracer()
def bad_filter(inputs: dict[str, Any]) -> dict[str, Any]:
raise ValueError("filter crashed!")
def my_node(state: SimpleState) -> SimpleState:
return {"value": f"processed_{state['value']}"}
_set_traceable_config(my_node, process_inputs=bad_filter)
builder = StateGraph(SimpleState)
builder.add_node("my_node", my_node)
builder.add_edge("__start__", "my_node")
graph = builder.compile()
# Graph should still execute successfully
result = graph.invoke({"value": "secret_pii_data"}, {"callbacks": [tracer]})
assert result == {"value": "processed_secret_pii_data"}
# Trace should show error placeholder, NOT the raw PII data
runs = tracer.flattened_runs()
node_runs = [r for r in runs if r.name == "my_node"]
assert len(node_runs) == 1
# Should NOT contain "secret_pii_data"
assert node_runs[0].inputs == {"error": "<trace_inputs processing failed>"}
def test_traceable_config_process_outputs_error_handling():
"""Test that errors in process_outputs don't leak PII."""
tracer = FakeTracer()
def bad_filter(outputs: Any) -> Any:
raise RuntimeError("output filter crashed!")
def my_node(state: SimpleState) -> SimpleState:
return {"value": "secret_output_pii"}
_set_traceable_config(my_node, process_outputs=bad_filter)
builder = StateGraph(SimpleState)
builder.add_node("my_node", my_node)
builder.add_edge("__start__", "my_node")
graph = builder.compile()
# Graph should still execute successfully
result = graph.invoke({"value": "input"}, {"callbacks": [tracer]})
assert result == {"value": "secret_output_pii"}
# Trace should show error placeholder, NOT the raw PII data
runs = tracer.flattened_runs()
node_runs = [r for r in runs if r.name == "my_node"]
assert len(node_runs) == 1
# Should NOT contain "secret_output_pii"
assert node_runs[0].outputs == {"error": "<trace_outputs processing failed>"}
def test_traceable_config_enabled_none_honors_external_context():
"""Test that enabled=None (default) honors external tracing context."""
tracer = FakeTracer()
def my_node(state: SimpleState) -> SimpleState:
return {"value": f"result_{state['value']}"}
# Set traceable config with enabled=None (the default)
_set_traceable_config(my_node, enabled=None)
builder = StateGraph(SimpleState)
builder.add_node("my_node", my_node)
builder.add_edge("__start__", "my_node")
graph = builder.compile()
result = graph.invoke({"value": "test"}, {"callbacks": [tracer]})
assert result == {"value": "result_test"}
# With enabled=None, tracing should proceed normally (honor external context)
runs = tracer.flattened_runs()
node_runs = [r for r in runs if r.name == "my_node"]
# Node should be traced since we provided a tracer callback
assert len(node_runs) == 1
async def test_traceable_context_propagates_to_children():
"""Traceable node should become the parent context for downstream nodes."""
tracer = FakeTracer()
async def traceable_node(state: SimpleState) -> dict:
return await subgraph.ainvoke({"value": f"traceable_{state['value']}"})
def child_node(state: SimpleState) -> SimpleState:
return {"value": f"child_{state['value']}"}
_set_traceable_config(traceable_node, process_inputs=lambda _: {"value": "[MASK]"})
subgraph_builder = StateGraph(SimpleState)
subgraph_builder.add_node("child_node", child_node)
subgraph_builder.add_edge("__start__", "child_node")
subgraph = subgraph_builder.compile(name="Subgraph")
builder = StateGraph(SimpleState)
builder.add_node("traceable_node", traceable_node)
builder.add_edge("__start__", "traceable_node")
graph = builder.compile()
result = await graph.ainvoke({"value": "secret"}, {"callbacks": [tracer]})
assert result == {"value": "child_traceable_secret"}
runs = tracer.flattened_runs()
traceable_runs = [r for r in runs if r.name == "traceable_node"]
child_runs = [r for r in runs if r.name == "child_node"]
subgraph_runs = [r for r in runs if r.name == "Subgraph"]
assert len(traceable_runs) == 1
assert len(child_runs) == 1
assert len(subgraph_runs) == 1
traceable_run = traceable_runs[0]
child_run = child_runs[0]
subgraph_run = subgraph_runs[0]
assert child_run.trace_id == traceable_run.trace_id
assert child_run.dotted_order.startswith(traceable_run.dotted_order)
assert child_run.parent_run_id == subgraph_run.id
assert subgraph_run.parent_run_id == traceable_run.id