Apply patch

libs/langgraph/langgraph/_internal/_runnable.py

@@ -61,6 +61,11 @@ try:
 except ImportError:
     _StreamingCallbackHandler = None  # type: ignore
 
+try:
+    from langgraph.pregel.protocol import PregelProtocol
+except ImportError:
+    PregelProtocol = None  # type: ignore
+
 
 def _set_config_context(
     config: RunnableConfig, run: Any = None
@@ -482,6 +487,56 @@ def is_async_generator(
     )
 
 
+class _PregelWrapper(Runnable):
+    """Wrapper for PregelProtocol instances to handle runtime context propagation.
+
+    When a compiled subgraph (PregelProtocol) is added as a node, this wrapper
+    extracts the runtime context from the config and passes it explicitly to
+    the subgraph's invoke method.
+    """
+
+    def __init__(self, pregel: PregelProtocol, name: str | None = None):
+        self.pregel = pregel
+        self._name = name or getattr(pregel, "name", None) or pregel.__class__.__name__
+
+    def __getattr__(self, name: str) -> Any:
+        """Delegate attribute access to the wrapped pregel instance."""
+        return getattr(self.pregel, name)
+
+    def get_name(self, suffix: str | None = None, *, name: str | None = None) -> str:
+        """Get the name of the runnable."""
+        name = name or self._name
+        return f"{name}{suffix}" if suffix else name
+
+    def invoke(
+        self, input: Any, config: RunnableConfig | None = None, **kwargs: Any
+    ) -> Any:
+        """Invoke the wrapped PregelProtocol with runtime context extracted from config."""
+        if config is None:
+            config = ensure_config()
+
+        # Extract runtime context from config
+        runtime = config.get(CONF, {}).get(CONFIG_KEY_RUNTIME)
+        context = runtime.context if runtime else None
+
+        # Invoke the subgraph with the extracted context
+        return self.pregel.invoke(input, config, context=context, **kwargs)
+
+    async def ainvoke(
+        self, input: Any, config: RunnableConfig | None = None, **kwargs: Any
+    ) -> Any:
+        """Async invoke the wrapped PregelProtocol with runtime context extracted from config."""
+        if config is None:
+            config = ensure_config()
+
+        # Extract runtime context from config
+        runtime = config.get(CONF, {}).get(CONFIG_KEY_RUNTIME)
+        context = runtime.context if runtime else None
+
+        # Invoke the subgraph with the extracted context
+        return await self.pregel.ainvoke(input, config, context=context, **kwargs)
+
+
 def coerce_to_runnable(
     thing: RunnableLike, *, name: str | None, trace: bool
 ) -> Runnable:
@@ -494,6 +549,12 @@ def coerce_to_runnable(
         A Runnable.
     """
     if isinstance(thing, Runnable):
+        # Check if this is a PregelProtocol instance (compiled subgraph)
+        # and wrap it to handle runtime context propagation
+        # Only wrap if the subgraph has a context_schema (needs runtime context)
+        if (PregelProtocol is not None and isinstance(thing, PregelProtocol) 
+            and hasattr(thing, 'context_schema') and thing.context_schema is not None):
+            return _PregelWrapper(thing, name=name)
         return thing
     elif is_async_generator(thing) or inspect.isgeneratorfunction(thing):
         return RunnableLambda(thing, name=name)
@@ -896,3 +957,5 @@ async def _consume_aiter(it: AsyncIterator[Any]) -> Any:
         else:
             output = chunk
     return output
+
+

test_simple_subgraph.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python3
+"""Simple test to check if basic subgraph functionality still works."""
+
+import sys
+import os
+
+# Add the langgraph library to the path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'libs', 'langgraph'))
+
+from dataclasses import dataclass
+from langgraph.graph.state import StateGraph
+from langgraph.runtime import Runtime
+from typing_extensions import TypedDict
+
+
+@dataclass
+class Context:
+    username: str
+
+class State(TypedDict):
+    my_key: str
+
+# Simple subgraph without runtime context
+def simple_node(state: State):
+    return {'my_key': state['my_key'] + ' processed'}
+
+subgraph_builder = StateGraph(State)
+subgraph_builder.add_node('simple_node', simple_node)
+subgraph_builder.set_entry_point('simple_node')
+subgraph_builder.set_finish_point('simple_node')
+subgraph = subgraph_builder.compile()
+
+# Parent graph
+def main_node(state: State):
+    return {'my_key': 'hello'}
+
+builder = StateGraph(State)
+builder.add_node('main_node', main_node)
+builder.add_node('subgraph_node', subgraph)
+builder.set_entry_point('main_node')
+builder.add_edge('main_node', 'subgraph_node')
+builder.set_finish_point('subgraph_node')
+graph = builder.compile()
+
+# Test basic functionality
+try:
+    result = graph.invoke({'my_key': 'start'})
+    print(f"Basic subgraph test SUCCESS: {result}")
+    
+    # Test graph drawing
+    try:
+        mermaid = graph.get_graph().draw_mermaid(with_styles=False)
+        print("Graph drawing SUCCESS")
+        print(f"Mermaid: {mermaid[:100]}...")
+    except Exception as e:
+        print(f"Graph drawing FAILED: {e}")
+    
+    sys.exit(0)
+except Exception as e:
+    print(f"Basic subgraph test FAILED: {e}")
+    import traceback
+    traceback.print_exc()
+    sys.exit(1)

test_subgraph_context_fix.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python3
+"""Test script to verify the runtime context propagation fix for subgraphs."""
+
+import sys
+import os
+
+# Add the langgraph library to the path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'libs', 'langgraph'))
+
+from dataclasses import dataclass
+from langgraph.graph.state import StateGraph
+from langgraph.runtime import Runtime
+from typing_extensions import TypedDict
+
+
+@dataclass
+class Context:
+    username: str
+
+class State(TypedDict):
+    foo: str
+
+# Subgraph
+def subgraph_node_1(state: State, runtime: Runtime[Context]):
+    return {'foo': '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()
+
+# Parent graph
+def main_node(state: State, runtime: Runtime[Context]):
+    return {'foo': '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()
+
+# Test the fix
+try:
+    context = Context(username='Alice')
+    result = graph.invoke({'foo': 'world'}, context=context)
+    print(f"SUCCESS: {result}")
+    print("Runtime context is now properly propagated to subgraphs!")
+    sys.exit(0)
+except Exception as e:
+    print(f"FAILED: {e}")
+    sys.exit(1)