diff --git a/libs/langgraph/RESUME_WRITES_FLOW.md b/libs/langgraph/RESUME_WRITES_FLOW.md new file mode 100644 index 000000000..0b77af868 --- /dev/null +++ b/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`. diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 16ffd86a0..db2860ab6 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/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, @@ -31,7 +32,6 @@ from langgraph_sdk.client import ( ) from langgraph_sdk.schema import ( Checkpoint, - Context, QueryParamTypes, ThreadState, ) @@ -66,6 +66,7 @@ from langgraph.types import ( StreamMode, StreamPart, ) +from langgraph.typing import ContextT, InputT, OutputT, StateT logger = logging.getLogger(__name__) @@ -109,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. @@ -689,10 +693,10 @@ class RemoteGraph(PregelProtocol): @overload def stream( self, - input: dict[str, Any] | Any, + input: InputT | Command | None, config: RunnableConfig | None = None, *, - context: Context | 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, @@ -701,15 +705,15 @@ class RemoteGraph(PregelProtocol): params: QueryParamTypes | None = None, version: Literal["v2"], **kwargs: Any, - ) -> Iterator[StreamPart]: ... + ) -> Iterator[StreamPart[OutputT, StateT]]: ... @overload def stream( self, - input: dict[str, Any] | Any, + input: InputT | Command | None, config: RunnableConfig | None = None, *, - context: Context | 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, @@ -722,10 +726,10 @@ class RemoteGraph(PregelProtocol): def stream( self, - input: dict[str, Any] | Any, + input: InputT | Command | None, config: RunnableConfig | None = None, *, - context: Context | 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, @@ -844,10 +848,10 @@ class RemoteGraph(PregelProtocol): @overload def astream( self, - input: dict[str, Any] | Any, + input: InputT | Command | None, config: RunnableConfig | None = None, *, - context: Context | 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, @@ -856,15 +860,15 @@ class RemoteGraph(PregelProtocol): params: QueryParamTypes | None = None, version: Literal["v2"], **kwargs: Any, - ) -> AsyncIterator[StreamPart]: ... + ) -> AsyncIterator[StreamPart[OutputT, StateT]]: ... @overload def astream( self, - input: dict[str, Any] | Any, + input: InputT | Command | None, config: RunnableConfig | None = None, *, - context: Context | 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, @@ -877,10 +881,10 @@ class RemoteGraph(PregelProtocol): async def astream( self, - input: dict[str, Any] | Any, + input: InputT | Command | None, config: RunnableConfig | None = None, *, - context: Context | 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, @@ -1015,25 +1019,25 @@ class RemoteGraph(PregelProtocol): @overload def invoke( self, - input: dict[str, Any] | Any, + input: InputT | Command | None, config: RunnableConfig | None = None, *, - context: Context | 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: Context | 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, @@ -1044,10 +1048,10 @@ class RemoteGraph(PregelProtocol): def invoke( self, - input: dict[str, Any] | Any, + input: InputT | Command | None, config: RunnableConfig | None = None, *, - context: Context | 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, @@ -1097,25 +1101,25 @@ class RemoteGraph(PregelProtocol): @overload async def ainvoke( self, - input: dict[str, Any] | Any, + input: InputT | Command | None, config: RunnableConfig | None = None, *, - context: Context | 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: Context | 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, @@ -1126,10 +1130,10 @@ class RemoteGraph(PregelProtocol): async def ainvoke( self, - input: dict[str, Any] | Any, + input: InputT | Command | None, config: RunnableConfig | None = None, *, - context: Context | 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, diff --git a/libs/langgraph/tests/test_type_checking.py b/libs/langgraph/tests/test_type_checking.py index 5f807ddaa..69ce54176 100644 --- a/libs/langgraph/tests/test_type_checking.py +++ b/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]