switch generics

libs/langgraph/RESUME_WRITES_FLOW.md

@@ -0,0 +1,67 @@
+# RESUME Writes Stripping: Complete Flow Reference
+
+## Legend
+
+| Column | Meaning |
+|---|---|
+| **Level** | P = Parent, S = Subgraph |
+| **`is_replaying`** | `CONFIG_KEY_CHECKPOINT_ID` key exists in `config[CONF]` (line 249) |
+| **`__enter__` via** | Which branch loads the checkpoint: **ckpt_id** (explicit checkpoint_id in checkpoint_config), **replay_state** (parent's ReplayState), **latest** (fetch most recent) |
+| **`RESUMING`** | Value of `CONFIG_KEY_RESUMING` in configurable (set by parent for subgraphs, absent for outer graph) |
+| **`is_resuming`** | Computed at line 633 — controls whether to "proceed past previous checkpoint" |
+| **`in_map`** | `replaying_from_checkpoint_map` — subgraph's ns found in checkpoint_map |
+| **Strip?** | Are RESUME pending writes stripped? (line 662-671) |
+
+## Setup
+
+```
+Parent:  START → executor (subgraph, checkpointer=True) → END
+Subgraph: START → step_a → ask_1 (interrupt) → ask_2 (interrupt) → END
+```
+
+## The Table
+
+| # | Scenario | Level | User call | `__enter__` via | `is_replaying` | `RESUMING` | `is_resuming` | `in_map` | Strip? | Why correct |
+|---|---|---|---|---|---|---|---|---|---|---|
+| 1 | **Fresh run** | P | `invoke({"v":[]}, cfg)` | latest (None) | False | _(absent)_ | False | — | N/A | No checkpoint yet, no writes to strip |
+| 1 | | S | _(Send from parent)_ | latest (None) | True¹ | False | False | False | N/A | No checkpoint yet |
+| 2 | **Resume single interrupt** | P | `invoke(Cmd(resume="a"), cfg)` | latest | False | _(absent)_ | True | — | No | Resuming — keep RESUME writes for interrupt() to return answer |
+| 2 | | S | _(Send)_ | latest | True¹ | True | True | False | No | `RESUMING=True` → keep. interrupt() returns "a", node completes |
+| 3 | **Resume 1st of 2 interrupts** | P | `invoke(Cmd(resume="a1"), cfg)` | latest | False | _(absent)_ | True | — | No | Keep RESUME writes — ask_1's answer must survive |
+| 3 | | S | _(Send)_ | latest | True¹ | True | True | False | **No** | ask_1 gets "a1" from RESUME write. ask_2 has no RESUME write → interrupt() re-fires. Correct. |
+| 4 | **Replay parent ckpt** (parent was mid-subgraph) | P | `invoke(None, parent_hist_cfg)` | ckpt_id | True | _(absent)_ | True | — | **Yes** | Replaying — strip stale RESUME writes so interrupts re-fire |
+| 4 | | S | _(Send)_ | replay_state² | True¹ | _(popped)_³ | False | False | **Yes** | `is_replaying=T`, `RESUMING` absent → strip. Subgraph replays cleanly |
+| 5 | **Time-travel to subgraph ckpt** (THE BUG) | P | `invoke(None, sub_cfg)` | ckpt_id⁴ | True | _(absent)_ | True | — | **Yes** | Parent replays from historical checkpoint |
+| 5 | | S | _(Send)_ | **ckpt_id**⁵ | True¹ | **True** | **True** | **True** | **Yes** ✨ | `in_map=True` overrides `RESUMING=True` → force strip. THE FIX. |
+| 5 | | S _(without fix)_ | _(Send)_ | ckpt_id⁵ | True¹ | **True** | **True** | _(no check)_ | **No** ❌ | BUG: `RESUMING=True` prevents strip → stale RESUME values → interrupt() doesn't re-fire |
+| 6 | **Fork from subgraph ckpt** | P | `invoke(None, update_state(sub_cfg,...))` | ckpt_id | True | _(absent)_ | True | — | **Yes** | Same as case 5 — fork creates new ckpt, but checkpoint_map still resolves |
+| 6 | | S | _(Send)_ | ckpt_id⁵ | True¹ | True | True | **True** | **Yes** ✨ | Same fix applies |
+| 7 | **Resume after case 5 re-interrupts** | P | `invoke(Cmd(resume="a2"), cfg)` | latest | False | _(absent)_ | True | — | No | Normal resume — keep RESUME writes |
+| 7 | | S | _(Send)_ | latest | True¹ | True | True | False⁶ | **No** | ask_2 gets "a2" from fresh RESUME write. Correct. |
+
+## Footnotes
+
+**¹** `is_replaying` is always `True` for subgraphs on tick 1 because `_algo.py` sets `CONFIG_KEY_CHECKPOINT_ID: None` — the key exists (even with `None` value), so `key in dict` is `True`. After tick 1, line 563 sets `is_replaying = False`.
+
+**²** `replay_state` branch: parent passed `CONFIG_KEY_REPLAY_STATE = ReplayState(parent_ckpt_id)`. The subgraph uses `replay_state.get_checkpoint()` which does `checkpointer.list(before=parent_ckpt_id, limit=1)` to find the subgraph's checkpoint from before the replay point.
+
+**³** The `replay_state` branch in `__enter__` (line 1158) explicitly pops `CONFIG_KEY_RESUMING` from config. This makes `is_resuming = False` in `_first()` because for nested graphs the fallback (`self.input is None or input_is_command`) is False (input is a Send arg).
+
+**⁴** Parent `__init__` clears `checkpoint_ns → ""` and `checkpoint_id → None` (line 273-277), then resolves `""` from checkpoint_map → gets `parent_checkpoint_id` onto `checkpoint_config` (line 278-290).
+
+**⁵** Subgraph `__init__` resolves its namespace (e.g. `"executor:task_id"`) from checkpoint_map → gets `subgraph_checkpoint_id` onto `checkpoint_config`. This is why the new first branch in `__enter__` (line 1141) fires — `checkpoint_config` has a truthy `checkpoint_id`.
+
+**⁶** After case 5 completes/re-interrupts and user resumes, the config is a normal thread config with no checkpoint_map entry for the subgraph. `in_map` is False, so normal resume logic applies.
+
+## The core tension (case 5)
+
+The parent **can't distinguish** these cases when propagating flags to subgraphs:
+
+| Parent sees | What's actually happening | Subgraph should strip RESUME? |
+|---|---|---|
+| `input=None`, has checkpoint | Resume after interrupt | Yes (replaying) |
+| `input=None`, has checkpoint | Resume after interrupt | No (resuming) |
+| `input=Command(resume=...)` | Active resume | No (resuming) |
+| `input=None`, has checkpoint | Time-travel to subgraph | Yes (replaying) |
+
+The **only** distinguishing signal at the subgraph level is whether its namespace appears in `checkpoint_map`.

libs/langgraph/langgraph/pregel/main.py

@@ -2456,7 +2456,7 @@ class Pregel(
         debug: bool | None = None,
         version: Literal["v2"],
         **kwargs: Unpack[DeprecatedKwargs],
-    ) -> Iterator[StreamPart[OutputT, StateT]]: ...
+    ) -> Iterator[StreamPart[StateT, OutputT]]: ...
 
     @overload
     def stream(
@@ -2787,7 +2787,7 @@ class Pregel(
         debug: bool | None = None,
         version: Literal["v2"],
         **kwargs: Unpack[DeprecatedKwargs],
-    ) -> AsyncIterator[StreamPart[OutputT, StateT]]: ...
+    ) -> AsyncIterator[StreamPart[StateT, OutputT]]: ...
 
     @overload
     def astream(
@@ -3194,7 +3194,7 @@ class Pregel(
         durability: Durability | None = None,
         version: Literal["v2"],
         **kwargs: Any,
-    ) -> list[StreamPart[OutputT, StateT]]: ...
+    ) -> list[StreamPart[StateT, OutputT]]: ...
 
     @overload
     def invoke(
@@ -3364,7 +3364,7 @@ class Pregel(
         durability: Durability | None = None,
         version: Literal["v2"],
         **kwargs: Any,
-    ) -> list[StreamPart[OutputT, StateT]]: ...
+    ) -> list[StreamPart[StateT, OutputT]]: ...
 
     @overload
     async def ainvoke(

libs/langgraph/langgraph/pregel/protocol.py

@@ -117,7 +117,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
         interrupt_after: All | Sequence[str] | None = None,
         subgraphs: bool = False,
         version: Literal["v2"],
-    ) -> Iterator[StreamPart[OutputT, StateT]]: ...
+    ) -> Iterator[StreamPart[StateT, OutputT]]: ...
 
     @overload
     @abstractmethod
@@ -161,7 +161,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
         interrupt_after: All | Sequence[str] | None = None,
         subgraphs: bool = False,
         version: Literal["v2"],
-    ) -> AsyncIterator[StreamPart[OutputT, StateT]]: ...
+    ) -> AsyncIterator[StreamPart[StateT, OutputT]]: ...
 
     @overload
     @abstractmethod

libs/langgraph/langgraph/pregel/remote.py

@@ -5,6 +5,7 @@ from collections.abc import AsyncIterator, Iterator, Sequence
 from dataclasses import asdict
 from typing import (
     Any,
+    Generic,
     Literal,
     cast,
     overload,
@@ -65,6 +66,7 @@ from langgraph.types import (
     StreamMode,
     StreamPart,
 )
+from langgraph.typing import ContextT, InputT, OutputT, StateT
 
 logger = logging.getLogger(__name__)
 
@@ -108,7 +110,10 @@ class RemoteException(Exception):
     pass
 
 
-class RemoteGraph(PregelProtocol):
+class RemoteGraph(
+    PregelProtocol[StateT, ContextT, InputT, OutputT],
+    Generic[StateT, ContextT, InputT, OutputT],
+):
     """The `RemoteGraph` class is a client implementation for calling remote
     APIs that implement the LangGraph Server API specification.
 
@@ -688,9 +693,10 @@ class RemoteGraph(PregelProtocol):
     @overload
     def stream(
         self,
-        input: dict[str, Any] | Any,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
+        context: ContextT | None = None,
         stream_mode: StreamMode | list[StreamMode] | None = None,
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
@@ -699,14 +705,15 @@ class RemoteGraph(PregelProtocol):
         params: QueryParamTypes | None = None,
         version: Literal["v2"],
         **kwargs: Any,
-    ) -> Iterator[StreamPart]: ...
+    ) -> Iterator[StreamPart[StateT, OutputT]]: ...
 
     @overload
     def stream(
         self,
-        input: dict[str, Any] | Any,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
+        context: ContextT | None = None,
         stream_mode: StreamMode | list[StreamMode] | None = None,
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
@@ -719,9 +726,10 @@ class RemoteGraph(PregelProtocol):
 
     def stream(
         self,
-        input: dict[str, Any] | Any,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
+        context: ContextT | None = None,
         stream_mode: StreamMode | list[StreamMode] | None = None,
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
@@ -769,6 +777,7 @@ class RemoteGraph(PregelProtocol):
             input=input,
             command=command,
             config=sanitized_config,
+            context=context,
             stream_mode=stream_modes,
             interrupt_before=interrupt_before,
             interrupt_after=interrupt_after,
@@ -839,9 +848,10 @@ class RemoteGraph(PregelProtocol):
     @overload
     def astream(
         self,
-        input: dict[str, Any] | Any,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
+        context: ContextT | None = None,
         stream_mode: StreamMode | list[StreamMode] | None = None,
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
@@ -850,14 +860,15 @@ class RemoteGraph(PregelProtocol):
         params: QueryParamTypes | None = None,
         version: Literal["v2"],
         **kwargs: Any,
-    ) -> AsyncIterator[StreamPart]: ...
+    ) -> AsyncIterator[StreamPart[StateT, OutputT]]: ...
 
     @overload
     def astream(
         self,
-        input: dict[str, Any] | Any,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
+        context: ContextT | None = None,
         stream_mode: StreamMode | list[StreamMode] | None = None,
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
@@ -870,9 +881,10 @@ class RemoteGraph(PregelProtocol):
 
     async def astream(
         self,
-        input: dict[str, Any] | Any,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
+        context: ContextT | None = None,
         stream_mode: StreamMode | list[StreamMode] | None = None,
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
@@ -920,6 +932,7 @@ class RemoteGraph(PregelProtocol):
             input=input,
             command=command,
             config=sanitized_config,
+            context=context,
             stream_mode=stream_modes,
             interrupt_before=interrupt_before,
             interrupt_after=interrupt_after,
@@ -1006,23 +1019,25 @@ class RemoteGraph(PregelProtocol):
     @overload
     def invoke(
         self,
-        input: dict[str, Any] | Any,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
+        context: ContextT | None = None,
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
         headers: dict[str, str] | None = None,
         params: QueryParamTypes | None = None,
         version: Literal["v2"],
         **kwargs: Any,
-    ) -> GraphOutput[dict[str, Any]]: ...
+    ) -> GraphOutput[OutputT]: ...
 
     @overload
     def invoke(
         self,
-        input: dict[str, Any] | Any,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
+        context: ContextT | None = None,
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
         headers: dict[str, str] | None = None,
@@ -1033,9 +1048,10 @@ class RemoteGraph(PregelProtocol):
 
     def invoke(
         self,
-        input: dict[str, Any] | Any,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
+        context: ContextT | None = None,
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
         headers: dict[str, str] | None = None,
@@ -1061,6 +1077,7 @@ class RemoteGraph(PregelProtocol):
         for chunk in self.stream(  # type: ignore[misc, call-overload]
             input,
             config=config,
+            context=context,
             interrupt_before=interrupt_before,
             interrupt_after=interrupt_after,
             headers=headers,
@@ -1084,23 +1101,25 @@ class RemoteGraph(PregelProtocol):
     @overload
     async def ainvoke(
         self,
-        input: dict[str, Any] | Any,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
+        context: ContextT | None = None,
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
         headers: dict[str, str] | None = None,
         params: QueryParamTypes | None = None,
         version: Literal["v2"],
         **kwargs: Any,
-    ) -> GraphOutput[dict[str, Any]]: ...
+    ) -> GraphOutput[OutputT]: ...
 
     @overload
     async def ainvoke(
         self,
-        input: dict[str, Any] | Any,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
+        context: ContextT | None = None,
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
         headers: dict[str, str] | None = None,
@@ -1111,9 +1130,10 @@ class RemoteGraph(PregelProtocol):
 
     async def ainvoke(
         self,
-        input: dict[str, Any] | Any,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
+        context: ContextT | None = None,
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
         headers: dict[str, str] | None = None,
@@ -1139,6 +1159,7 @@ class RemoteGraph(PregelProtocol):
         async for chunk in self.astream(  # type: ignore[misc, call-overload]
             input,
             config=config,
+            context=context,
             interrupt_before=interrupt_before,
             interrupt_after=interrupt_after,
             headers=headers,

libs/langgraph/langgraph/types.py

@@ -335,7 +335,7 @@ StreamPart = TypeAliasType(
     | CheckpointStreamPart[StateT]
     | TasksStreamPart
     | DebugStreamPart[StateT],
-    type_params=(OutputT, StateT),
+    type_params=(StateT, OutputT),
 )
 """A discriminated union of all v2 stream part types.
 

libs/langgraph/tests/test_remote_graph.py

@@ -1,5 +1,6 @@
 import re
 import sys
+from dataclasses import dataclass
 from typing import Annotated
 from unittest.mock import AsyncMock, MagicMock
 
@@ -10,6 +11,7 @@ from langchain_core.runnables import RunnableConfig
 from langchain_core.runnables.graph import Edge as DrawableEdge
 from langchain_core.runnables.graph import Node as DrawableNode
 from langgraph_sdk.schema import StreamPart
+from pydantic import BaseModel
 from typing_extensions import TypedDict
 
 from langgraph.errors import GraphInterrupt
@@ -908,6 +910,188 @@ async def test_ainvoke():
     assert result == {"messages": [{"type": "human", "content": "world"}]}
 
 
+def test_stream_context():
+    """Test that context is passed through to the SDK client in stream."""
+    mock_sync_client = MagicMock()
+    mock_sync_client.runs.stream.return_value = [
+        StreamPart(event="values", data={"chunk": "data1"}),
+    ]
+
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        sync_client=mock_sync_client,
+    )
+
+    config = {"configurable": {"thread_id": "thread_1"}}
+    context = {"model_name": "anthropic", "user_id": "123"}
+    stream_parts = list(
+        remote_pregel.stream(
+            {"input": "data"},
+            config,
+            context=context,
+            stream_mode="values",
+        )
+    )
+
+    assert stream_parts == [{"chunk": "data1"}]
+    _, kwargs = mock_sync_client.runs.stream.call_args
+    assert kwargs["context"] == {"model_name": "anthropic", "user_id": "123"}
+
+
+def test_stream_context_none():
+    """Test that context defaults to None when not provided."""
+    mock_sync_client = MagicMock()
+    mock_sync_client.runs.stream.return_value = [
+        StreamPart(event="values", data={"chunk": "data1"}),
+    ]
+
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        sync_client=mock_sync_client,
+    )
+
+    config = {"configurable": {"thread_id": "thread_1"}}
+    list(remote_pregel.stream({"input": "data"}, config, stream_mode="values"))
+
+    _, kwargs = mock_sync_client.runs.stream.call_args
+    assert kwargs["context"] is None
+
+
+@pytest.mark.anyio
+async def test_astream_context():
+    """Test that context is passed through to the SDK client in astream."""
+    mock_async_client = MagicMock()
+    async_iter = MagicMock()
+    async_iter.__aiter__.return_value = [
+        StreamPart(event="values", data={"chunk": "data1"}),
+    ]
+    mock_async_client.runs.stream.return_value = async_iter
+
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        client=mock_async_client,
+    )
+
+    config = {"configurable": {"thread_id": "thread_1"}}
+    context = {"model_name": "anthropic"}
+    chunks = []
+    async for chunk in remote_pregel.astream(
+        {"input": "data"},
+        config,
+        context=context,
+        stream_mode="values",
+    ):
+        chunks.append(chunk)
+
+    assert chunks == [{"chunk": "data1"}]
+    _, kwargs = mock_async_client.runs.stream.call_args
+    assert kwargs["context"] == {"model_name": "anthropic"}
+
+
+def test_invoke_context():
+    """Test that context is passed through to the SDK client in invoke."""
+    mock_sync_client = MagicMock()
+    mock_sync_client.runs.stream.return_value = [
+        StreamPart(event="values", data={"result": "done"}),
+    ]
+
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        sync_client=mock_sync_client,
+    )
+
+    config = {"configurable": {"thread_id": "thread_1"}}
+    context = {"model_name": "openai"}
+    result = remote_pregel.invoke({"input": "data"}, config, context=context)
+
+    assert result == {"result": "done"}
+    _, kwargs = mock_sync_client.runs.stream.call_args
+    assert kwargs["context"] == {"model_name": "openai"}
+
+
+@pytest.mark.anyio
+async def test_ainvoke_context():
+    """Test that context is passed through to the SDK client in ainvoke."""
+    mock_async_client = MagicMock()
+    async_iter = MagicMock()
+    async_iter.__aiter__.return_value = [
+        StreamPart(event="values", data={"result": "done"}),
+    ]
+    mock_async_client.runs.stream.return_value = async_iter
+
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        client=mock_async_client,
+    )
+
+    config = {"configurable": {"thread_id": "thread_1"}}
+    context = {"user_id": "456"}
+    result = await remote_pregel.ainvoke({"input": "data"}, config, context=context)
+
+    assert result == {"result": "done"}
+    _, kwargs = mock_async_client.runs.stream.call_args
+    assert kwargs["context"] == {"user_id": "456"}
+
+
+def test_stream_context_dataclass():
+    """Test that a dataclass context is passed through to the SDK client."""
+
+    @dataclass
+    class MyContext:
+        model_name: str
+        user_id: str
+
+    mock_sync_client = MagicMock()
+    mock_sync_client.runs.stream.return_value = [
+        StreamPart(event="values", data={"chunk": "data1"}),
+    ]
+
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        sync_client=mock_sync_client,
+    )
+
+    config = {"configurable": {"thread_id": "thread_1"}}
+    ctx = MyContext(model_name="anthropic", user_id="123")
+    list(
+        remote_pregel.stream(
+            {"input": "data"}, config, context=ctx, stream_mode="values"
+        )
+    )
+
+    _, kwargs = mock_sync_client.runs.stream.call_args
+    assert kwargs["context"] == ctx
+
+
+def test_stream_context_base_model():
+    """Test that a BaseModel context is passed through to the SDK client."""
+
+    class MyContext(BaseModel):
+        model_name: str
+        user_id: str
+
+    mock_sync_client = MagicMock()
+    mock_sync_client.runs.stream.return_value = [
+        StreamPart(event="values", data={"chunk": "data1"}),
+    ]
+
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        sync_client=mock_sync_client,
+    )
+
+    config = {"configurable": {"thread_id": "thread_1"}}
+    ctx = MyContext(model_name="anthropic", user_id="123")
+    list(
+        remote_pregel.stream(
+            {"input": "data"}, config, context=ctx, stream_mode="values"
+        )
+    )
+
+    _, kwargs = mock_sync_client.runs.stream.call_args
+    assert kwargs["context"] == ctx
+
+
 @pytest.mark.skip(
     "Unskip this test to manually test the LangSmith Deployment integration"
 )

libs/langgraph/tests/test_type_checking.py

@@ -8,7 +8,8 @@ from pydantic import BaseModel
 from typing_extensions import TypedDict
 
 from langgraph.graph import StateGraph
-from langgraph.types import Command
+from langgraph.pregel.remote import RemoteGraph
+from langgraph.types import Command, GraphOutput, StreamPart
 
 
 def test_typed_dict_state() -> None:
@@ -159,3 +160,75 @@ def test_add_node_with_explicit_input_schema() -> None:
     # because it violates the principles of contravariance
     workflow.add_node("a_narrow", a, input_schema=ANarrow)  # type: ignore[arg-type]
     workflow.add_node("b_narrow", b, input_schema=BNarrow)  # type: ignore[arg-type]
+
+
+@pytest.mark.skip("Purely for type checking")
+def test_remote_graph_generics_typed_dict() -> None:
+    """RemoteGraph parameterized with TypedDict should propagate types."""
+
+    class MyState(TypedDict):
+        messages: list[str]
+
+    rg: RemoteGraph[MyState, None, MyState, MyState] = RemoteGraph(
+        "test", url="http://localhost:8123"
+    )
+
+    # v2 invoke should return GraphOutput[MyState]
+    result: GraphOutput[MyState] = rg.invoke({"messages": ["hi"]}, version="v2")
+    _val: MyState = result.value
+
+    # v1 invoke should return dict[str, Any] | Any
+    _v1_result: dict[str, Any] | Any = rg.invoke({"messages": ["hi"]})
+
+    # v2 stream should yield StreamPart[MyState, MyState]
+    for part in rg.stream({"messages": ["hi"]}, version="v2"):
+        _part: StreamPart[MyState, MyState] = part
+
+    # input should accept the state type
+    rg.invoke({"messages": ["hi"]}, version="v2")
+
+    # input should also accept Command
+    rg.invoke(Command(), version="v2")
+
+    # input should also accept None
+    rg.invoke(None, version="v2")
+
+
+@pytest.mark.skip("Purely for type checking")
+def test_remote_graph_generics_pydantic() -> None:
+    """RemoteGraph parameterized with Pydantic model should propagate types."""
+
+    class PydanticState(BaseModel):
+        messages: list[str]
+
+    rg: RemoteGraph[PydanticState, None, PydanticState, PydanticState] = RemoteGraph(
+        "test", url="http://localhost:8123"
+    )
+
+    result: GraphOutput[PydanticState] = rg.invoke(
+        PydanticState(messages=["hi"]), version="v2"
+    )
+    _val: PydanticState = result.value
+
+
+@pytest.mark.skip("Purely for type checking")
+def test_remote_graph_separate_input_output() -> None:
+    """RemoteGraph with different input/output schemas."""
+
+    class InputState(TypedDict):
+        query: str
+
+    class OutputState(TypedDict):
+        answer: str
+
+    class FullState(InputState, OutputState): ...
+
+    rg: RemoteGraph[FullState, None, InputState, OutputState] = RemoteGraph(
+        "test", url="http://localhost:8123"
+    )
+
+    result: GraphOutput[OutputState] = rg.invoke({"query": "hi"}, version="v2")
+    _val: OutputState = result.value
+
+    # wrong input type should fail type checking
+    rg.invoke({"answer": "wrong"}, version="v2")  # type: ignore[call-overload]