fix(langgraph): stop at an interrupt_before breakpoint when a drain is requested before it

At a step boundary, tick() checked for a drain before the interrupt_before
breakpoints. A drain requested in the step before a breakpoint stopped the
run with GraphDrained, and the resume after it (invoke(None, config)) marks
the next step's breakpoints as passed, so the breakpoint never fired. The
drain check now runs after the breakpoint check, so the run stops at the
breakpoint as it would have without the drain. Subgraph breakpoints share
the fix, since subgraphs share the run's control.

libs/langgraph/langgraph/pregel/_loop.py

@@ -685,10 +685,6 @@ class PregelLoop:
             self.status = "done"
             return False
 
-        if self.control is not None and self.control.drain_requested:
-            self.status = "draining"
-            return False
-
         # if there are pending writes from a previous loop, apply them
         if self._reapplies_pending_writes and self.checkpoint_pending_writes:
             self._reapply_writes_to_succeeded_nodes(self.tasks)
@@ -701,6 +697,13 @@ class PregelLoop:
             self.status = "interrupt_before"
             raise GraphInterrupt()
 
+        # stop before running the next tasks if a drain was requested. after
+        # the breakpoint check: a resume passes the next step's breakpoints,
+        # so draining in front of one would skip it
+        if self.control is not None and self.control.drain_requested:
+            self.status = "draining"
+            return False
+
         # produce debug output
         self._emit("tasks", map_debug_tasks, self.tasks.values())
 

libs/langgraph/tests/test_runtime.py

@@ -248,6 +248,95 @@ def test_drain_from_subgraph_can_resume_parent() -> None:
     }
 
 
+class _BreakpointState(TypedDict, total=False):
+    first: str
+    second: str
+
+
+def _drain_before_breakpoint_builder(
+    second_runs: list[str],
+) -> StateGraph[_BreakpointState]:
+    def first(state: _BreakpointState, runtime: Runtime) -> _BreakpointState:
+        runtime.control.request_drain("rollout")
+        return {"first": "done"}
+
+    def second(state: _BreakpointState) -> _BreakpointState:
+        second_runs.append("second")
+        return {"second": "done"}
+
+    builder = StateGraph(_BreakpointState)
+    builder.add_node("first", first)
+    builder.add_node("second", second)
+    builder.add_edge(START, "first")
+    builder.add_edge("first", "second")
+    return builder
+
+
+@pytest.mark.parametrize("at_compile", [True, False])
+def test_drain_before_interrupt_before_stops_at_breakpoint(at_compile: bool) -> None:
+    # A resume passes the next step's breakpoints, so a drain right before a
+    # breakpoint has to stop there instead, or the resumed run skips it.
+    second_runs: list[str] = []
+    breakpoints = {"interrupt_before": ["second"]}
+    compiled = _drain_before_breakpoint_builder(second_runs).compile(
+        checkpointer=MemorySaver(), **(breakpoints if at_compile else {})
+    )
+    run_kwargs = {} if at_compile else breakpoints
+    config = {"configurable": {"thread_id": "drain-breakpoint"}}
+
+    assert compiled.invoke({}, config, control=RunControl(), **run_kwargs) == {
+        "first": "done"
+    }
+    assert compiled.get_state(config).next == ("second",)
+    assert second_runs == []
+
+    assert compiled.invoke(None, config, **run_kwargs) == {
+        "first": "done",
+        "second": "done",
+    }
+    assert second_runs == ["second"]
+
+
+def test_drain_before_subgraph_interrupt_before_stops_at_breakpoint() -> None:
+    second_runs: list[str] = []
+    child = _drain_before_breakpoint_builder(second_runs).compile(
+        interrupt_before=["second"]
+    )
+    parent = StateGraph(_BreakpointState)
+    parent.add_node("child", child)
+    parent.add_edge(START, "child")
+    compiled = parent.compile(checkpointer=MemorySaver())
+    config = {"configurable": {"thread_id": "drain-subgraph-breakpoint"}}
+
+    compiled.invoke({}, config, control=RunControl())
+    state = compiled.get_state(config, subgraphs=True)
+    assert state.next == ("child",)
+    assert state.tasks[0].state.next == ("second",)
+    assert second_runs == []
+
+    assert compiled.invoke(None, config) == {"first": "done", "second": "done"}
+    assert second_runs == ["second"]
+
+
+@pytest.mark.anyio
+async def test_drain_before_interrupt_before_stops_at_breakpoint_async() -> None:
+    second_runs: list[str] = []
+    compiled = _drain_before_breakpoint_builder(second_runs).compile(
+        checkpointer=MemorySaver(), interrupt_before=["second"]
+    )
+    config = {"configurable": {"thread_id": "drain-breakpoint-async"}}
+
+    assert await compiled.ainvoke({}, config, control=RunControl()) == {"first": "done"}
+    assert (await compiled.aget_state(config)).next == ("second",)
+    assert second_runs == []
+
+    assert await compiled.ainvoke(None, config) == {
+        "first": "done",
+        "second": "done",
+    }
+    assert second_runs == ["second"]
+
+
 @pytest.mark.anyio
 async def test_drain_requested_in_terminal_step_finishes_normally_async() -> None:
     class State(TypedDict, total=False):