From 6b17b12ae473fe3c3755b03363ef8a31a09b934c Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Wed, 7 Jan 2026 13:33:25 -0800 Subject: [PATCH] feat: Honor traceable config --- .../langgraph/_internal/_runnable.py | 471 +++++++++++------- libs/langgraph/langgraph/pregel/_read.py | 86 +++- .../tests/test_traceable_integration.py | 422 ++++++++++++++++ 3 files changed, 788 insertions(+), 191 deletions(-) create mode 100644 libs/langgraph/tests/test_traceable_integration.py diff --git a/libs/langgraph/langgraph/_internal/_runnable.py b/libs/langgraph/langgraph/_internal/_runnable.py index 63e03f544..15938bdae 100644 --- a/libs/langgraph/langgraph/_internal/_runnable.py +++ b/libs/langgraph/langgraph/_internal/_runnable.py @@ -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": ""} + + +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": ""} + 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: diff --git a/libs/langgraph/langgraph/pregel/_read.py b/libs/langgraph/langgraph/pregel/_read.py index 8d4c21135..077de0843 100644 --- a/libs/langgraph/langgraph/pregel/_read.py +++ b/libs/langgraph/langgraph/pregel/_read.py @@ -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: diff --git a/libs/langgraph/tests/test_traceable_integration.py b/libs/langgraph/tests/test_traceable_integration.py new file mode 100644 index 000000000..be056e509 --- /dev/null +++ b/libs/langgraph/tests/test_traceable_integration.py @@ -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": ""} + + +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": ""} + + +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