refac

libs/langgraph/langgraph/pregel/_call.py

@@ -21,6 +21,7 @@ from langgraph._internal._runnable import (
     run_in_executor,
 )
 from langgraph.config import get_config
+from langgraph.pregel._read import _validate_traceable_config
 from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry
 from langgraph.types import CachePolicy, RetryPolicy
 
@@ -221,20 +222,15 @@ def get_runnable_for_task(func: Callable[..., Any]) -> Runnable:
 
         # Check for traceable config early to handle __unwrapped__
         raw_config = getattr(func, "__traceable_config__", None)
-        traceable_config = None
-        if raw_config and isinstance(raw_config, dict):
-            traceable_config = {
-                "process_inputs": raw_config.get("process_inputs"),
-                "process_outputs": raw_config.get("process_outputs"),
-                "enabled": raw_config.get("enabled"),
-                "__unwrapped__": raw_config.get("__unwrapped__"),
-            }
+        traceable_config = _validate_traceable_config(raw_config)
 
         # Use unwrapped function if available to avoid double-tracing
         # when @traceable and @task are used together
-        func_to_run = func
-        if traceable_config and traceable_config.get("__unwrapped__"):
-            func_to_run = traceable_config["__unwrapped__"]
+        func_to_run: Callable[..., Any] = func
+        if traceable_config:
+            unwrapped = traceable_config.get("__unwrapped__")
+            if unwrapped is not None:
+                func_to_run = unwrapped
 
         if is_async_callable(func_to_run):
             run = RunnableCallable(
@@ -262,10 +258,11 @@ def get_runnable_for_task(func: Callable[..., Any]) -> Runnable:
         base_trace_inputs = functools.partial(
             _explode_args_trace_inputs, inspect.signature(func)
         )
-        if traceable_config and traceable_config.get("process_inputs"):
-            trace_inputs = _compose_trace_inputs(
-                base_trace_inputs, traceable_config["process_inputs"]
-            )
+        process_inputs = (
+            traceable_config.get("process_inputs") if traceable_config else None
+        )
+        if process_inputs is not None:
+            trace_inputs = _compose_trace_inputs(base_trace_inputs, process_inputs)
         else:
             trace_inputs = base_trace_inputs
 

libs/langgraph/langgraph/pregel/_read.py

@@ -41,11 +41,13 @@ def _validate_traceable_config(raw: Any) -> TraceableConfig | None:
     """
     if not isinstance(raw, dict):
         return None
+    # Support both "__unwrapped__" (langsmith) and "wrapped" (legacy) keys
+    unwrapped = raw.get("__unwrapped__") or raw.get("wrapped")
     return {
         "process_inputs": raw.get("process_inputs"),
         "process_outputs": raw.get("process_outputs"),
         "enabled": raw.get("enabled"),  # None means use external context
-        "wrapped": raw.get("wrapped"),
+        "__unwrapped__": unwrapped,
     }
 
 

libs/langgraph/tests/test_traceable_integration.py

@@ -52,7 +52,7 @@ def _set_traceable_config(
             "process_inputs": process_inputs,
             "process_outputs": process_outputs,
             "enabled": enabled,
-            "wrapped": func,
+            "__unwrapped__": func,
         },
     )
     return func