feat: Add context coercion for LangGraph runtime (langchain-ai#5736)

Fixes langchain-ai#5735

Implement context coercion functionality for LangGraph runtime to
improve API usability.

Key changes:
- Added `_coerce_context` function in `pregel/main.py`
- Supports coercion for:
  - Pydantic BaseModel
  - Dataclasses
  - TypedDict
- Comprehensive test coverage added in `tests/test_runtime.py`
- Handles edge cases like None context and missing fields

The implementation allows users to pass dictionaries as context, which
will be automatically converted to the expected schema type, making the
API more flexible and user-friendly.

---------

Co-authored-by: open-swe[bot] <[email protected]>
Co-authored-by: Sydney Runkle <[email protected]>
Co-authored-by: Sydney Runkle <[email protected]>
Co-authored-by: Eugene Yurtsev <[email protected]>

Author: 4 people

libs/langgraph/langgraph/pregel/main.py

@@ -2571,7 +2571,7 @@ def stream_writer(c: Any) -> None:
                 config[CONF][CONFIG_KEY_DURABILITY] = durability_
 
             runtime = Runtime(
-                context=context,
+                context=_coerce_context(self.context_schema, context),
                 store=store,
                 stream_writer=stream_writer,
                 previous=None,
@@ -2866,7 +2866,7 @@ def stream_writer(c: Any) -> None:
                 config[CONF][CONFIG_KEY_DURABILITY] = durability_
 
             runtime = Runtime(
-                context=context,
+                context=_coerce_context(self.context_schema, context),
                 store=store,
                 stream_writer=stream_writer,
                 previous=None,
@@ -3224,3 +3224,33 @@ def _output(
                 yield (ns, payload)
             else:
                 yield payload
+
+
+def _coerce_context(
+    context_schema: type[ContextT] | None, context: Any
+) -> ContextT | None:
+    """Coerce context input to the appropriate schema type.
+
+    If context is a dict and context_schema is a dataclass or pydantic model, we coerce.
+    Else, we return the context as-is.
+
+    Args:
+        context_schema: The schema type to coerce to (BaseModel, dataclass, or TypedDict)
+        context: The context value to coerce
+
+    Returns:
+        The coerced context value or None if context is None
+    """
+    if context is None:
+        return None
+
+    if context_schema is None:
+        return context
+
+    schema_is_class = issubclass(context_schema, BaseModel) or is_dataclass(
+        context_schema
+    )
+    if isinstance(context, dict) and schema_is_class:
+        return context_schema(**context)  # type: ignore[misc]
+
+    return cast(ContextT, context)

libs/langgraph/tests/test_runtime.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass
 from typing import Any
 
+import pytest
+from pydantic import BaseModel, ValidationError
 from typing_extensions import TypedDict
 
 from langgraph.graph import END, START, StateGraph
@@ -106,3 +108,284 @@ def main_node(state: State, runtime: Runtime[Context]):
     context = Context(username="Alice")
     result = graph.invoke({}, context=context)
     assert result == {"subgraph": "Alice!", "main": "Alice!"}
+
+
+def test_context_coercion_dataclass() -> None:
+    """Test that dict context is coerced to dataclass."""
+
+    @dataclass
+    class Context:
+        api_key: str
+        timeout: int = 30
+
+    class State(TypedDict):
+        message: str
+
+    def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
+        return {
+            "message": f"api_key: {runtime.context.api_key}, timeout: {runtime.context.timeout}"
+        }
+
+    graph = StateGraph(state_schema=State, context_schema=Context)
+    graph.add_node("node", node_with_context)
+    graph.add_edge(START, "node")
+    graph.add_edge("node", END)
+    compiled = graph.compile()
+
+    # Test dict coercion with all fields
+    result = compiled.invoke(
+        {"message": "test"}, context={"api_key": "sk_test", "timeout": 60}
+    )
+    assert result == {"message": "api_key: sk_test, timeout: 60"}
+
+    # Test dict coercion with default field
+    result = compiled.invoke({"message": "test"}, context={"api_key": "sk_test2"})
+    assert result == {"message": "api_key: sk_test2, timeout: 30"}
+
+    # Test with actual dataclass instance (should still work)
+    result = compiled.invoke(
+        {"message": "test"}, context=Context(api_key="sk_test3", timeout=90)
+    )
+    assert result == {"message": "api_key: sk_test3, timeout: 90"}
+
+
+def test_context_coercion_pydantic() -> None:
+    """Test that dict context is coerced to Pydantic model."""
+
+    class Context(BaseModel):
+        api_key: str
+        timeout: int = 30
+        tags: list[str] = []
+
+    class State(TypedDict):
+        message: str
+
+    def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
+        return {
+            "message": f"api_key: {runtime.context.api_key}, timeout: {runtime.context.timeout}, tags: {runtime.context.tags}"
+        }
+
+    graph = StateGraph(state_schema=State, context_schema=Context)
+    graph.add_node("node", node_with_context)
+    graph.add_edge(START, "node")
+    graph.add_edge("node", END)
+    compiled = graph.compile()
+
+    # Test dict coercion with all fields
+    result = compiled.invoke(
+        {"message": "test"},
+        context={"api_key": "sk_test", "timeout": 60, "tags": ["prod", "v2"]},
+    )
+    assert result == {"message": "api_key: sk_test, timeout: 60, tags: ['prod', 'v2']"}
+
+    # Test dict coercion with defaults
+    result = compiled.invoke({"message": "test"}, context={"api_key": "sk_test2"})
+    assert result == {"message": "api_key: sk_test2, timeout: 30, tags: []"}
+
+    # Test with actual Pydantic instance (should still work)
+    result = compiled.invoke(
+        {"message": "test"},
+        context=Context(api_key="sk_test3", timeout=90, tags=["test"]),
+    )
+    assert result == {"message": "api_key: sk_test3, timeout: 90, tags: ['test']"}
+
+
+def test_context_coercion_typeddict() -> None:
+    """Test that dict context with TypedDict schema passes through as-is."""
+
+    class Context(TypedDict):
+        api_key: str
+        timeout: int
+
+    class State(TypedDict):
+        message: str
+
+    def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
+        # TypedDict context is just a dict at runtime
+        return {
+            "message": f"api_key: {runtime.context['api_key']}, timeout: {runtime.context['timeout']}"
+        }
+
+    graph = StateGraph(state_schema=State, context_schema=Context)
+    graph.add_node("node", node_with_context)
+    graph.add_edge(START, "node")
+    graph.add_edge("node", END)
+    compiled = graph.compile()
+
+    # Test dict passes through for TypedDict
+    result = compiled.invoke(
+        {"message": "test"}, context={"api_key": "sk_test", "timeout": 60}
+    )
+    assert result == {"message": "api_key: sk_test, timeout: 60"}
+
+
+def test_context_coercion_none() -> None:
+    """Test that None context is handled properly."""
+
+    @dataclass
+    class Context:
+        api_key: str
+
+    class State(TypedDict):
+        message: str
+
+    def node_without_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
+        # Should be None when no context provided
+        return {"message": f"context is None: {runtime.context is None}"}
+
+    graph = StateGraph(state_schema=State, context_schema=Context)
+    graph.add_node("node", node_without_context)
+    graph.add_edge(START, "node")
+    graph.add_edge("node", END)
+    compiled = graph.compile()
+
+    # Test with None context
+    result = compiled.invoke({"message": "test"}, context=None)
+    assert result == {"message": "context is None: True"}
+
+    # Test without context parameter (defaults to None)
+    result = compiled.invoke({"message": "test"})
+    assert result == {"message": "context is None: True"}
+
+
+def test_context_coercion_errors() -> None:
+    """Test error handling for invalid context."""
+
+    @dataclass
+    class Context:
+        api_key: str  # Required field
+
+    class State(TypedDict):
+        message: str
+
+    def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
+        return {"message": "should not reach here"}
+
+    graph = StateGraph(state_schema=State, context_schema=Context)
+    graph.add_node("node", node_with_context)
+    graph.add_edge(START, "node")
+    graph.add_edge("node", END)
+    compiled = graph.compile()
+
+    # Test missing required field
+    with pytest.raises(TypeError):
+        compiled.invoke({"message": "test"}, context={"timeout": 60})
+
+    # Test invalid dict keys
+    with pytest.raises(TypeError):
+        compiled.invoke(
+            {"message": "test"}, context={"api_key": "test", "invalid_field": "value"}
+        )
+
+
+@pytest.mark.anyio
+async def test_context_coercion_async() -> None:
+    """Test context coercion with async methods."""
+
+    @dataclass
+    class Context:
+        api_key: str
+        async_mode: bool = True
+
+    class State(TypedDict):
+        message: str
+
+    async def async_node(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
+        return {
+            "message": f"async api_key: {runtime.context.api_key}, async_mode: {runtime.context.async_mode}"
+        }
+
+    graph = StateGraph(state_schema=State, context_schema=Context)
+    graph.add_node("node", async_node)
+    graph.add_edge(START, "node")
+    graph.add_edge("node", END)
+    compiled = graph.compile()
+
+    # Test dict coercion with ainvoke
+    result = await compiled.ainvoke(
+        {"message": "test"}, context={"api_key": "sk_async", "async_mode": False}
+    )
+    assert result == {"message": "async api_key: sk_async, async_mode: False"}
+
+    # Test dict coercion with astream
+    chunks = []
+    async for chunk in compiled.astream(
+        {"message": "test"}, context={"api_key": "sk_stream"}
+    ):
+        chunks.append(chunk)
+
+    # Find the chunk with our node output
+    node_output = None
+    for chunk in chunks:
+        if "node" in chunk:
+            node_output = chunk["node"]
+            break
+
+    assert node_output == {"message": "async api_key: sk_stream, async_mode: True"}
+
+
+def test_context_coercion_stream() -> None:
+    """Test context coercion with sync stream method."""
+
+    @dataclass
+    class Context:
+        api_key: str
+        stream_mode: str = "default"
+
+    class State(TypedDict):
+        message: str
+
+    def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
+        return {
+            "message": f"stream api_key: {runtime.context.api_key}, mode: {runtime.context.stream_mode}"
+        }
+
+    graph = StateGraph(state_schema=State, context_schema=Context)
+    graph.add_node("node", node_with_context)
+    graph.add_edge(START, "node")
+    graph.add_edge("node", END)
+    compiled = graph.compile()
+
+    # Test dict coercion with stream
+    chunks = []
+    for chunk in compiled.stream(
+        {"message": "test"}, context={"api_key": "sk_stream", "stream_mode": "fast"}
+    ):
+        chunks.append(chunk)
+
+    # Find the chunk with our node output
+    node_output = None
+    for chunk in chunks:
+        if "node" in chunk:
+            node_output = chunk["node"]
+            break
+
+    assert node_output == {"message": "stream api_key: sk_stream, mode: fast"}
+
+
+def test_context_coercion_pydantic_validation_errors() -> None:
+    """Test that Pydantic validation errors are raised."""
+
+    class Context(BaseModel):
+        api_key: str
+        timeout: int
+
+    class State(TypedDict):
+        message: str
+
+    def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
+        return {
+            "message": f"api_key: {runtime.context.api_key}, timeout: {runtime.context.timeout}"
+        }
+
+    graph = StateGraph(state_schema=State, context_schema=Context)
+    graph.add_node("node", node_with_context)
+    graph.add_edge(START, "node")
+    graph.add_edge("node", END)
+
+    compiled = graph.compile()
+
+    with pytest.raises(ValidationError):
+        compiled.invoke(
+            {"message": "test"}, context={"api_key": "sk_test", "timeout": "not_an_int"}
+        )