update docstring

libs/langgraph/langgraph/pregel/main.py

@@ -22,6 +22,7 @@ from inspect import isclass
 from typing import (
     Any,
     Generic,
+    Literal,
     cast,
     get_type_hints,
 )
@@ -38,6 +39,7 @@ from langchain_core.runnables.config import (
     get_callback_manager_for_config,
 )
 from langchain_core.runnables.graph import Graph
+from langchain_core.runnables.schema import CustomStreamEvent, StandardStreamEvent
 from langgraph.cache.base import BaseCache
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
@@ -3021,6 +3023,61 @@ class Pregel(
             await asyncio.shield(run_manager.on_chain_error(e))
             raise
 
+    async def astream_events(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        version: Literal["v2"],
+        include_names: Sequence[str] | None = None,
+        include_types: Sequence[str] | None = None,
+        include_tags: Sequence[str] | None = None,
+        exclude_names: Sequence[str] | None = None,
+        exclude_types: Sequence[str] | None = None,
+        exclude_tags: Sequence[str] | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[StandardStreamEvent | CustomStreamEvent]:
+        """Stream events from the graph execution.
+
+        This method extends the base Runnable.astream_events with support for
+        the LangGraph `context` parameter.
+
+        Args:
+            input: The input to the graph.
+            config: The configuration to use for the run.
+            context: The static context to use for the run.
+                !!! version-added "Added in version 1.0.6"
+            version: The version of the event stream schema to use ("v2").
+            include_names: Only include events from runnables with matching names.
+            include_types: Only include events from runnables with matching types.
+            include_tags: Only include events from runnables with matching tags.
+            exclude_names: Exclude events from runnables with matching names.
+            exclude_types: Exclude events from runnables with matching types.
+            exclude_tags: Exclude events from runnables with matching tags.
+            **kwargs: Additional arguments passed to the underlying stream.
+
+        Yields:
+            Events from the graph execution.
+        """
+        async with contextlib.aclosing(
+            super().astream_events(
+                input,
+                config,
+                version=version,
+                include_names=include_names,
+                include_types=include_types,
+                include_tags=include_tags,
+                exclude_names=exclude_names,
+                exclude_types=exclude_types,
+                exclude_tags=exclude_tags,
+                context=context,
+                **kwargs,
+            )
+        ) as stream:  # type: ignore[type-var]
+            async for event in stream:
+                yield event
+
     def invoke(
         self,
         input: InputT | Command | None,

libs/langgraph/langgraph/pregel/remote.py

@@ -907,6 +907,7 @@ class RemoteGraph(PregelProtocol):
         input: Any,
         config: RunnableConfig | None = None,
         *,
+        context: Any = None,
         version: Literal["v1", "v2"],
         include_names: Sequence[All] | None = None,
         include_types: Sequence[All] | None = None,

libs/langgraph/tests/test_runtime.py

@@ -389,3 +389,45 @@ def test_context_coercion_pydantic_validation_errors() -> None:
         compiled.invoke(
             {"message": "test"}, context={"api_key": "sk_test", "timeout": "not_an_int"}
         )
+
+
+@pytest.mark.anyio
+async def test_context_with_astream_events() -> None:
+    """Test that context is properly passed through astream_events."""
+
+    @dataclass
+    class Context:
+        api_key: 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}"}
+
+    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()
+
+    events = []
+    async for event in compiled.astream_events(
+        {"message": "test"},
+        version="v2",
+        context=Context(api_key="sk_events_123"),
+    ):
+        events.append(event)
+
+    # Verify we got events
+    assert len(events) > 0
+
+    # Find the final on_chain_end event (no parent_ids means it's the root)
+    end_events = [
+        e for e in events if e["event"] == "on_chain_end" and not e.get("parent_ids")
+    ]
+    assert len(end_events) == 1
+
+    # Verify the output contains our context value
+    output = end_events[0]["data"]["output"]
+    assert output["message"] == "api_key: sk_events_123"