Fix tracing hierarchy for imperative api

libs/langgraph/langgraph/func/__init__.py

@@ -19,7 +19,7 @@ from typing_extensions import ParamSpec
 from langgraph.channels.ephemeral_value import EphemeralValue
 from langgraph.channels.last_value import LastValue
 from langgraph.checkpoint.base import BaseCheckpointSaver
-from langgraph.constants import END, START, TAG_HIDDEN
+from langgraph.constants import CONF, END, START, TAG_HIDDEN
 from langgraph.pregel import Pregel
 from langgraph.pregel.call import get_runnable_for_func
 from langgraph.pregel.read import PregelNode
@@ -39,11 +39,11 @@ def call(
     retry: Optional[RetryPolicy] = None,
 ) -> concurrent.futures.Future[T]:
     from langgraph.constants import CONFIG_KEY_CALL
-    from langgraph.utils.config import get_configurable
+    from langgraph.utils.config import get_config
 
-    conf = get_configurable()
-    impl = conf[CONFIG_KEY_CALL]
-    fut = impl(func, input, retry=retry)
+    config = get_config()
+    impl = config[CONF][CONFIG_KEY_CALL]
+    fut = impl(func, input, retry=retry, callbacks=config["callbacks"])
     return fut
 
 

libs/langgraph/langgraph/pregel/algo.py

@@ -1,5 +1,6 @@
 import sys
 from collections import defaultdict, deque
+from contextvars import Context
 from functools import partial
 from hashlib import sha1
 from typing import (
@@ -19,6 +20,7 @@ from typing import (
 )
 from uuid import UUID
 
+from langchain_core.callbacks import Callbacks
 from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager
 from langchain_core.runnables.config import RunnableConfig
 
@@ -107,18 +109,25 @@ class PregelTaskWrites(NamedTuple):
 
 
 class Call:
-    __slots__ = ("func", "input", "retry")
+    __slots__ = ("func", "input", "retry", "callbacks")
 
     func: Callable
     input: Any
     retry: Optional[RetryPolicy]
+    callbacks: Optional[Callbacks]
 
     def __init__(
-        self, func: Callable, input: Any, *, retry: Optional[RetryPolicy]
+        self,
+        func: Callable,
+        input: Any,
+        *,
+        retry: Optional[RetryPolicy],
+        callbacks: Optional[Context],
     ) -> None:
         self.func = func
         self.input = input
         self.retry = retry
+        self.callbacks = callbacks
 
 
 def should_interrupt(
@@ -465,9 +474,8 @@ def prepare_single_task(
                 patch_config(
                     merge_configs(config, {"metadata": metadata}),
                     run_name=name,
-                    callbacks=(
-                        manager.get_child(f"graph:step:{step}") if manager else None
-                    ),
+                    callbacks=call.callbacks
+                    or (manager.get_child(f"graph:step:{step}") if manager else None),
                     configurable={
                         CONFIG_KEY_TASK_ID: task_id,
                         # deque.extend is thread-safe

libs/langgraph/langgraph/pregel/executor.py

@@ -63,10 +63,11 @@ class BackgroundExecutor(ContextManager):
         __next_tick__: bool = False,
         **kwargs: P.kwargs,
     ) -> concurrent.futures.Future[T]:
+        ctx = copy_context()
         if __next_tick__:
-            task = self.executor.submit(next_tick, fn, *args, **kwargs)
+            task = self.executor.submit(next_tick, ctx.run, fn, *args, **kwargs)
         else:
-            task = self.executor.submit(fn, *args, **kwargs)
+            task = self.executor.submit(ctx.run, fn, *args, **kwargs)
         self.tasks[task] = (__cancel_on_exit__, __reraise_on_exit__)
         task.add_done_callback(self.done)
         return task

libs/langgraph/langgraph/pregel/runner.py

@@ -17,6 +17,8 @@ from typing import (
     cast,
 )
 
+from langchain_core.callbacks import Callbacks
+
 from langgraph.constants import (
     CONF,
     CONFIG_KEY_CALL,
@@ -148,9 +150,12 @@ class PregelRunner:
             input: Any,
             *,
             retry: Optional[RetryPolicy] = None,
+            callbacks: Optional[Callbacks] = None,
         ) -> concurrent.futures.Future[Any]:
             (fut,) = writer(
-                task, [(PUSH, None)], calls=[Call(func, input, retry=retry)]
+                task,
+                [(PUSH, None)],
+                calls=[Call(func, input, retry=retry, callbacks=callbacks)],
             )
             assert fut is not None, "writer did not return a future for call"
             return fut
@@ -337,9 +342,12 @@ class PregelRunner:
             input: Any,
             *,
             retry: Optional[RetryPolicy] = None,
+            callbacks: Optional[Callbacks] = None,
         ) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]:
             (fut,) = writer(
-                task, [(PUSH, None)], calls=[Call(func, input, retry=retry)]
+                task,
+                [(PUSH, None)],
+                calls=[Call(func, input, retry=retry, callbacks=callbacks)],
             )
             assert fut is not None, "writer did not return a future for call"
             if asyncio.iscoroutinefunction(func):

libs/langgraph/langgraph/types.py

@@ -443,6 +443,7 @@ def interrupt(value: Any) -> Any:
         GraphInterrupt: On the first invocation within the node, halts execution and surfaces the provided value to the client.
     """
     from langgraph.constants import (
+        CONF,
         CONFIG_KEY_CHECKPOINT_NS,
         CONFIG_KEY_SCRATCHPAD,
         CONFIG_KEY_SEND,
@@ -453,9 +454,9 @@ def interrupt(value: Any) -> Any:
         RESUME,
     )
     from langgraph.errors import GraphInterrupt
-    from langgraph.utils.config import get_configurable
+    from langgraph.utils.config import get_config
 
-    conf = get_configurable()
+    conf = get_config()[CONF]
     # track interrupt index
     scratchpad: PregelScratchpad = conf[CONFIG_KEY_SCRATCHPAD]
     if "interrupt_counter" not in scratchpad:

libs/langgraph/langgraph/utils/config.py

@@ -297,7 +297,7 @@ def ensure_config(*configs: Optional[RunnableConfig]) -> RunnableConfig:
     return empty
 
 
-def get_configurable() -> dict[str, Any]:
+def get_config() -> RunnableConfig:
     if sys.version_info < (3, 11):
         try:
             if asyncio.current_task():
@@ -307,6 +307,6 @@ def get_configurable() -> dict[str, Any]:
         except RuntimeError:
             pass
     if var_config := var_child_runnable_config.get():
-        return var_config[CONF]
+        return var_config
     else:
         raise RuntimeError("Called get_configurable outside of a runnable context")

libs/langgraph/tests/test_pregel_async.py

@@ -2465,7 +2465,8 @@ async def test_imp_task(checkpointer_name: str) -> None:
             answer = interrupt("question")
             return [m + answer for m in mapped]
 
-        thread1 = {"configurable": {"thread_id": "1"}}
+        tracer = FakeTracer()
+        thread1 = {"configurable": {"thread_id": "1"}, "callbacks": [tracer]}
         assert [c async for c in graph.astream([0, 1], thread1)] == [
             {"mapper": "00"},
             {"mapper": "11"},
@@ -2481,6 +2482,9 @@ async def test_imp_task(checkpointer_name: str) -> None:
             },
         ]
         assert mapper_calls == 2
+        assert len(tracer.runs) == 1
+        assert len(tracer.runs[0].child_runs) == 1
+        assert tracer.runs[0].child_runs[0].name == "graph"
 
         assert await graph.ainvoke(Command(resume="answer"), thread1) == [
             "00answer",