skip maybe

This commit is contained in:
Sydney Runkle
2026-03-04 13:20:52 -08:00
parent a3823395cf
commit 52b586370d
2 changed files with 105 additions and 0 deletions
+32
View File
@@ -1109,6 +1109,22 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
if saved.pending_writes is not None
else []
)
# When re-invoking from a specific checkpoint (not skip_done_tasks)
# that isn't a subgraph being resumed by its parent (CONFIG_KEY_RESUMING),
# clear cached RESUME writes so that interrupt() re-fires instead of
# returning stale cached values. This mirrors the behavior of fork
# checkpoints created via update_state, which start with no writes.
# Command(resume=...) will add its own fresh RESUME writes in _first().
if (
not self.skip_done_tasks
and CONFIG_KEY_RESUMING
not in self.config.get(CONF, {})
):
self.checkpoint_pending_writes = [
w
for w in self.checkpoint_pending_writes
if w[1] != RESUME
]
self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
self.channels, self.managed = channels_from_checkpoint(
@@ -1288,6 +1304,22 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
if saved.pending_writes is not None
else []
)
# When re-invoking from a specific checkpoint (not skip_done_tasks)
# that isn't a subgraph being resumed by its parent (CONFIG_KEY_RESUMING),
# clear cached RESUME writes so that interrupt() re-fires instead of
# returning stale cached values. This mirrors the behavior of fork
# checkpoints created via update_state, which start with no writes.
# Command(resume=...) will add its own fresh RESUME writes in _first().
if (
not self.skip_done_tasks
and CONFIG_KEY_RESUMING
not in self.config.get(CONF, {})
):
self.checkpoint_pending_writes = [
w
for w in self.checkpoint_pending_writes
if w[1] != RESUME
]
self.submit = await self.stack.enter_async_context(
AsyncBackgroundExecutor(self.config)
+73
View File
@@ -5316,6 +5316,79 @@ def test_multiple_interrupt_state_persistence(
assert state.values["steps"] == ["step1", "step2"]
def test_fork_from_resolved_interrupt_retriggers(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""Replaying from a checkpoint before an interrupt node should re-trigger
the interrupt rather than reusing cached resume values from the original
execution. This covers both the checkpoint that directly resolved the
interrupt and an earlier checkpoint (before the interrupt node)."""
called: list[str] = []
class State(TypedDict):
value: Annotated[list[str], operator.add]
def node_a(state: State) -> State:
called.append("node_a")
return {"value": ["a"]}
def ask_human(state: State) -> State:
called.append("ask_human")
answer = interrupt("What is your input?")
return {"value": [f"human:{answer}"]}
def node_b(state: State) -> State:
called.append("node_b")
return {"value": ["b"]}
graph = (
StateGraph(State)
.add_node("node_a", node_a)
.add_node("ask_human", ask_human)
.add_node("node_b", node_b)
.add_edge(START, "node_a")
.add_edge("node_a", "ask_human")
.add_edge("ask_human", "node_b")
.compile(checkpointer=sync_checkpointer)
)
config = {"configurable": {"thread_id": "1"}}
# 1. Run until interrupt
result = graph.invoke({"value": []}, config)
assert "__interrupt__" in result
# 2. Resume with answer — completes the full graph
result = graph.invoke(Command(resume="hello"), config)
assert result == {"value": ["a", "human:hello", "b"]}
# 3. Find checkpoint before ask_human (after node_a completed)
history = list(graph.get_state_history(config))
before_ask = [s for s in history if s.next == ("ask_human",)][-1]
# 4. Replay from that checkpoint — interrupt should re-fire
called.clear()
replay_result = graph.invoke(None, before_ask.config)
assert "__interrupt__" in replay_result
assert replay_result["value"] == ["a"]
assert replay_result["__interrupt__"][0].value == "What is your input?"
# ask_human was called but hit interrupt before returning
assert "ask_human" in called
# node_a should NOT run (it's before our checkpoint)
assert "node_a" not in called
# node_b should NOT run (interrupt halted execution)
assert "node_b" not in called
# 5. Resume the re-triggered interrupt with a new answer
called.clear()
result = graph.invoke(Command(resume="world"), before_ask.config)
assert result == {"value": ["a", "human:world", "b"]}
assert "ask_human" in called
assert "node_b" in called
def test_concurrent_execution_thread_safety():
"""Test thread safety during concurrent execution."""