use parent runtime

libs/langgraph/langgraph/pregel/main.py

@@ -12,6 +12,7 @@ from dataclasses import is_dataclass
 from functools import partial
 from inspect import isclass
 from typing import Any, Callable, Generic, Optional, Union, cast, get_type_hints
+from unittest.mock import DEFAULT
 from uuid import UUID, uuid5
 
 from langchain_core.globals import get_debug
@@ -116,7 +117,7 @@ from langgraph.pregel._validate import validate_graph, validate_keys
 from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry
 from langgraph.pregel.debug import get_bolded_text, get_colored_text, tasks_w_writes
 from langgraph.pregel.protocol import PregelProtocol, StreamChunk, StreamProtocol
-from langgraph.runtime import Runtime
+from langgraph.runtime import DEFAULT_RUNTIME, Runtime
 from langgraph.store.base import BaseStore
 from langgraph.types import (
     All,
@@ -2570,12 +2571,16 @@ class Pregel(
             if durability is not None or deprecated_checkpoint_during is not None:
                 config[CONF][CONFIG_KEY_DURABILITY] = durability_
 
-            config[CONF][CONFIG_KEY_RUNTIME] = Runtime(
+            runtime = Runtime(
                 context=context,
                 store=store,
                 stream_writer=stream_writer,
                 previous=None,
             )
+            parent_runtime = config[CONF].get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME)
+            runtime = parent_runtime.merge(runtime)
+            config[CONF][CONFIG_KEY_RUNTIME] = runtime
+
             with SyncPregelLoop(
                 input,
                 stream=StreamProtocol(stream.put, stream_modes),
@@ -2861,12 +2866,16 @@ class Pregel(
             if durability is not None or deprecated_checkpoint_during is not None:
                 config[CONF][CONFIG_KEY_DURABILITY] = durability_
 
-            config[CONF][CONFIG_KEY_RUNTIME] = Runtime(
+            runtime = Runtime(
                 context=context,
                 store=store,
                 stream_writer=stream_writer,
                 previous=None,
             )
+            parent_runtime = config[CONF].get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME)
+            runtime = parent_runtime.merge(runtime)
+            config[CONF][CONFIG_KEY_RUNTIME] = runtime
+
             async with AsyncPregelLoop(
                 input,
                 stream=StreamProtocol(stream.put_nowait, stream_modes),

libs/langgraph/tests/test_runtime.py

@@ -7,16 +7,14 @@ from langgraph.graph import END, START, StateGraph
 from langgraph.runtime import Runtime, get_runtime
 
 
-@dataclass
-class Context:
-    api_key: str
-
-
-class State(TypedDict):
-    message: str
-
-
 def test_injected_runtime() -> None:
+    @dataclass
+    class Context:
+        api_key: str
+
+    class State(TypedDict):
+        message: str
+
     def injected_runtime(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
         return {"message": f"api key: {runtime.context.api_key}"}
 
@@ -32,6 +30,13 @@ def test_injected_runtime() -> None:
 
 
 def test_context_runtime() -> None:
+    @dataclass
+    class Context:
+        api_key: str
+
+    class State(TypedDict):
+        message: str
+
     def context_runtime(state: State) -> dict[str, Any]:
         runtime = get_runtime(Context)
         return {"message": f"api key: {runtime.context.api_key}"}
@@ -45,3 +50,56 @@ def test_context_runtime() -> None:
         {"message": "hello world"}, context=Context(api_key="sk_123456")
     )
     assert result == {"message": "api key: sk_123456"}
+
+
+def test_override_runtime() -> None:
+    @dataclass
+    class Context:
+        api_key: str
+
+    runtime1 = Runtime(context=Context(api_key="abc"))
+    runtime = runtime1.override(context=Context(api_key="def"))
+    assert runtime.context.api_key == "def"
+
+
+def test_merge_runtime() -> None:
+    @dataclass
+    class Context:
+        api_key: str
+
+    runtime1 = Runtime(context=Context(api_key="abc"))
+    runtime2 = Runtime(context=Context(api_key="def"))
+    runtime = runtime1.merge(runtime2)
+    assert runtime.context.api_key == "def"
+
+
+def test_runtime_propogated_to_subgraph() -> None:
+    @dataclass
+    class Context:
+        username: str
+
+    class State(TypedDict, total=False):
+        foo: str
+        bar: str
+
+    def subgraph_node_1(state: State, runtime: Runtime[Context]):
+        return {"bar": f"hi {runtime.context.username}!"}
+
+    subgraph_builder = StateGraph(State, context_schema=Context)
+    subgraph_builder.add_node(subgraph_node_1)
+    subgraph_builder.set_entry_point("subgraph_node_1")
+    subgraph = subgraph_builder.compile()
+
+    def main_node(state: State, runtime: Runtime[Context]):
+        return {"foo": f"hello {runtime.context.username}!"}
+
+    builder = StateGraph(State, context_schema=Context)
+    builder.add_node(main_node)
+    builder.add_node("node_1", subgraph)
+    builder.set_entry_point("main_node")
+    builder.add_edge("main_node", "node_1")
+    graph = builder.compile()
+
+    context = Context(username="Alice")
+    result = graph.invoke({}, context=context)
+    assert result == {"foo": "hello Alice!", "bar": "hi Alice!"}