mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-07 18:27:52 +02:00
going crazy w/ tests
This commit is contained in:
Sydney Runkle
@@ -735,7 +735,10 @@ class PregelLoop:
|
||||
# (fork scenario), True when resuming from latest checkpoint.
|
||||
self.config = patch_configurable(
|
||||
self.config,
|
||||
{CONFIG_KEY_RESUMING: has_resume or (is_resuming and self.skip_done_tasks)},
|
||||
{
|
||||
CONFIG_KEY_RESUMING: has_resume
|
||||
or (is_resuming and self.skip_done_tasks)
|
||||
},
|
||||
)
|
||||
# set flag
|
||||
self.status = "pending"
|
||||
@@ -1123,9 +1126,9 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
# We must NOT drop them when resuming (subgraph via CONFIG_KEY_RESUMING,
|
||||
# or top graph via Command(resume=...)) because with multiple interrupts
|
||||
# previously resolved RESUME values need to be preserved.
|
||||
is_resuming = (
|
||||
self.config.get(CONF, {}).get(CONFIG_KEY_RESUMING) is True
|
||||
) or (isinstance(self.input, Command) and self.input.resume is not None)
|
||||
is_resuming = (self.config.get(CONF, {}).get(CONFIG_KEY_RESUMING) is True) or (
|
||||
isinstance(self.input, Command) and self.input.resume is not None
|
||||
)
|
||||
if not self.skip_done_tasks and not is_resuming:
|
||||
self.checkpoint_pending_writes = [
|
||||
w for w in self.checkpoint_pending_writes if w[1] != RESUME
|
||||
@@ -1314,9 +1317,9 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
# We must NOT drop them when resuming (subgraph via CONFIG_KEY_RESUMING,
|
||||
# or top graph via Command(resume=...)) because with multiple interrupts
|
||||
# previously resolved RESUME values need to be preserved.
|
||||
is_resuming = (
|
||||
self.config.get(CONF, {}).get(CONFIG_KEY_RESUMING) is True
|
||||
) or (isinstance(self.input, Command) and self.input.resume is not None)
|
||||
is_resuming = (self.config.get(CONF, {}).get(CONFIG_KEY_RESUMING) is True) or (
|
||||
isinstance(self.input, Command) and self.input.resume is not None
|
||||
)
|
||||
if not self.skip_done_tasks and not is_resuming:
|
||||
self.checkpoint_pending_writes = [
|
||||
w for w in self.checkpoint_pending_writes if w[1] != RESUME
|
||||
|
||||
@@ -5572,210 +5572,6 @@ def test_fork_after_all_interrupts(
|
||||
assert "node_b" not in called
|
||||
|
||||
|
||||
def test_fork_subgraph_interrupt_no_checkpointer(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Fork/replay with a subgraph that has no checkpointer (checkpointer=False/None).
|
||||
On fork (input=None), the parent propagates CONFIG_KEY_RESUMING=False so the
|
||||
subgraph strips RESUME writes and re-fires the interrupt — consistent with
|
||||
top-level interrupt behavior on fork/replay."""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
class State(TypedDict):
|
||||
value: Annotated[list[str], operator.add]
|
||||
|
||||
def sub_interrupt(state: State) -> State:
|
||||
called.append("sub_interrupt")
|
||||
answer = interrupt("Sub question?")
|
||||
return {"value": [f"sub:{answer}"]}
|
||||
|
||||
subgraph = (
|
||||
StateGraph(State)
|
||||
.add_node("sub_interrupt", sub_interrupt)
|
||||
.add_edge(START, "sub_interrupt")
|
||||
.compile() # no checkpointer
|
||||
)
|
||||
|
||||
def call_subgraph(state: State) -> State:
|
||||
called.append("call_subgraph")
|
||||
return subgraph.invoke(state)
|
||||
|
||||
def after(state: State) -> State:
|
||||
called.append("after")
|
||||
return {"value": ["after"]}
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("call_subgraph", call_subgraph)
|
||||
.add_node("after", after)
|
||||
.add_edge(START, "call_subgraph")
|
||||
.add_edge("call_subgraph", "after")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# 1. Run until interrupt
|
||||
result = graph.invoke({"value": []}, config)
|
||||
assert "__interrupt__" in result
|
||||
assert result["__interrupt__"][0].value == "Sub question?"
|
||||
|
||||
# 2. Resume — completes
|
||||
result = graph.invoke(Command(resume="answer"), config)
|
||||
assert result == {"value": ["sub:answer", "after"]}
|
||||
|
||||
# 3. Find checkpoint before subgraph node
|
||||
history = list(graph.get_state_history(config))
|
||||
before_sub = [s for s in history if s.next == ("call_subgraph",)][-1]
|
||||
|
||||
# 4. Replay — subgraph re-fires interrupt (consistent with top-level behavior)
|
||||
called.clear()
|
||||
replay_result = graph.invoke(None, before_sub.config)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Sub question?"
|
||||
assert "call_subgraph" in called
|
||||
assert "after" not in called
|
||||
|
||||
|
||||
def test_fork_subgraph_interrupt_checkpointer_true(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Fork/replay with a subgraph that has checkpointer=True.
|
||||
On fork, subgraph re-fires interrupts — consistent with top-level behavior."""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
class State(TypedDict):
|
||||
value: Annotated[list[str], operator.add]
|
||||
|
||||
def sub_node(state: State) -> State:
|
||||
called.append("sub_node")
|
||||
return {"value": ["sub_node"]}
|
||||
|
||||
def sub_interrupt(state: State) -> State:
|
||||
called.append("sub_interrupt")
|
||||
answer = interrupt("Sub question?")
|
||||
return {"value": [f"sub:{answer}"]}
|
||||
|
||||
subgraph = (
|
||||
StateGraph(State)
|
||||
.add_node("sub_node", sub_node)
|
||||
.add_node("sub_interrupt", sub_interrupt)
|
||||
.add_edge(START, "sub_node")
|
||||
.add_edge("sub_node", "sub_interrupt")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
def call_subgraph(state: State) -> State:
|
||||
called.append("call_subgraph")
|
||||
return subgraph.invoke(state)
|
||||
|
||||
def after(state: State) -> State:
|
||||
called.append("after")
|
||||
return {"value": ["after"]}
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("call_subgraph", call_subgraph)
|
||||
.add_node("after", after)
|
||||
.add_edge(START, "call_subgraph")
|
||||
.add_edge("call_subgraph", "after")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# 1. Run until interrupt
|
||||
result = graph.invoke({"value": []}, config)
|
||||
assert "__interrupt__" in result
|
||||
assert result["__interrupt__"][0].value == "Sub question?"
|
||||
|
||||
# 2. Resume — completes
|
||||
result = graph.invoke(Command(resume="answer"), config)
|
||||
assert result == {"value": ["sub_node", "sub:answer", "after"]}
|
||||
|
||||
# 3. Find checkpoint before subgraph node
|
||||
history = list(graph.get_state_history(config))
|
||||
before_sub = [s for s in history if s.next == ("call_subgraph",)][-1]
|
||||
|
||||
# 4. Replay — subgraph re-fires interrupt (consistent with top-level behavior)
|
||||
called.clear()
|
||||
replay_result = graph.invoke(None, before_sub.config)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Sub question?"
|
||||
assert "call_subgraph" in called
|
||||
assert "after" not in called
|
||||
|
||||
|
||||
def test_fork_subgraph_two_interrupts_no_checkpointer(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Fork/replay with a subgraph (no checkpointer) containing two interrupt
|
||||
nodes. On fork, subgraph re-fires the first interrupt — consistent with
|
||||
top-level behavior."""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
class State(TypedDict):
|
||||
value: Annotated[list[str], operator.add]
|
||||
|
||||
def sub_int_1(state: State) -> State:
|
||||
called.append("sub_int_1")
|
||||
answer = interrupt("Sub Q1?")
|
||||
return {"value": [f"s1:{answer}"]}
|
||||
|
||||
def sub_int_2(state: State) -> State:
|
||||
called.append("sub_int_2")
|
||||
answer = interrupt("Sub Q2?")
|
||||
return {"value": [f"s2:{answer}"]}
|
||||
|
||||
subgraph = (
|
||||
StateGraph(State)
|
||||
.add_node("sub_int_1", sub_int_1)
|
||||
.add_node("sub_int_2", sub_int_2)
|
||||
.add_edge(START, "sub_int_1")
|
||||
.add_edge("sub_int_1", "sub_int_2")
|
||||
.compile() # no checkpointer
|
||||
)
|
||||
|
||||
def call_subgraph(state: State) -> State:
|
||||
called.append("call_subgraph")
|
||||
return subgraph.invoke(state)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("call_subgraph", call_subgraph)
|
||||
.add_edge(START, "call_subgraph")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# 1. Run until first sub-interrupt
|
||||
result = graph.invoke({"value": []}, config)
|
||||
assert "__interrupt__" in result
|
||||
assert result["__interrupt__"][0].value == "Sub Q1?"
|
||||
|
||||
# 2. Resume first
|
||||
result = graph.invoke(Command(resume="a1"), config)
|
||||
assert "__interrupt__" in result
|
||||
assert result["__interrupt__"][0].value == "Sub Q2?"
|
||||
|
||||
# 3. Resume second — completes
|
||||
result = graph.invoke(Command(resume="a2"), config)
|
||||
assert result == {"value": ["s1:a1", "s2:a2"]}
|
||||
|
||||
# 4. Replay from before subgraph — re-fires first interrupt
|
||||
history = list(graph.get_state_history(config))
|
||||
before_sub = [s for s in history if s.next == ("call_subgraph",)][-1]
|
||||
|
||||
called.clear()
|
||||
replay_result = graph.invoke(None, before_sub.config)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Sub Q1?"
|
||||
|
||||
|
||||
def test_concurrent_execution_thread_safety():
|
||||
"""Test thread safety during concurrent execution."""
|
||||
|
||||
|
||||
@@ -6776,197 +6776,6 @@ async def test_fork_after_all_interrupts(
|
||||
assert "node_b" not in called
|
||||
|
||||
|
||||
async def test_fork_subgraph_interrupt_no_checkpointer(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Fork/replay with subgraph (no checkpointer). On fork, the parent
|
||||
propagates CONFIG_KEY_RESUMING=False so the subgraph strips RESUME writes
|
||||
and re-fires the interrupt — consistent with top-level behavior."""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
class State(TypedDict):
|
||||
value: Annotated[list[str], operator.add]
|
||||
|
||||
def sub_interrupt(state: State) -> State:
|
||||
called.append("sub_interrupt")
|
||||
answer = interrupt("Sub question?")
|
||||
return {"value": [f"sub:{answer}"]}
|
||||
|
||||
subgraph = (
|
||||
StateGraph(State)
|
||||
.add_node("sub_interrupt", sub_interrupt)
|
||||
.add_edge(START, "sub_interrupt")
|
||||
.compile()
|
||||
)
|
||||
|
||||
def call_subgraph(state: State) -> State:
|
||||
called.append("call_subgraph")
|
||||
return subgraph.invoke(state)
|
||||
|
||||
def after(state: State) -> State:
|
||||
called.append("after")
|
||||
return {"value": ["after"]}
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("call_subgraph", call_subgraph)
|
||||
.add_node("after", after)
|
||||
.add_edge(START, "call_subgraph")
|
||||
.add_edge("call_subgraph", "after")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
result = await graph.ainvoke({"value": []}, config)
|
||||
assert "__interrupt__" in result
|
||||
assert result["__interrupt__"][0].value == "Sub question?"
|
||||
|
||||
result = await graph.ainvoke(Command(resume="answer"), config)
|
||||
assert result == {"value": ["sub:answer", "after"]}
|
||||
|
||||
history = [s async for s in graph.aget_state_history(config)]
|
||||
before_sub = [s for s in history if s.next == ("call_subgraph",)][-1]
|
||||
|
||||
called.clear()
|
||||
replay_result = await graph.ainvoke(None, before_sub.config)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Sub question?"
|
||||
assert "call_subgraph" in called
|
||||
assert "after" not in called
|
||||
|
||||
|
||||
async def test_fork_subgraph_interrupt_checkpointer_true(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Fork/replay with subgraph (checkpointer=True). On fork, subgraph
|
||||
re-fires interrupts — consistent with top-level behavior."""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
class State(TypedDict):
|
||||
value: Annotated[list[str], operator.add]
|
||||
|
||||
def sub_node(state: State) -> State:
|
||||
called.append("sub_node")
|
||||
return {"value": ["sub_node"]}
|
||||
|
||||
def sub_interrupt(state: State) -> State:
|
||||
called.append("sub_interrupt")
|
||||
answer = interrupt("Sub question?")
|
||||
return {"value": [f"sub:{answer}"]}
|
||||
|
||||
subgraph = (
|
||||
StateGraph(State)
|
||||
.add_node("sub_node", sub_node)
|
||||
.add_node("sub_interrupt", sub_interrupt)
|
||||
.add_edge(START, "sub_node")
|
||||
.add_edge("sub_node", "sub_interrupt")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
def call_subgraph(state: State) -> State:
|
||||
called.append("call_subgraph")
|
||||
return subgraph.invoke(state)
|
||||
|
||||
def after(state: State) -> State:
|
||||
called.append("after")
|
||||
return {"value": ["after"]}
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("call_subgraph", call_subgraph)
|
||||
.add_node("after", after)
|
||||
.add_edge(START, "call_subgraph")
|
||||
.add_edge("call_subgraph", "after")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
result = await graph.ainvoke({"value": []}, config)
|
||||
assert "__interrupt__" in result
|
||||
assert result["__interrupt__"][0].value == "Sub question?"
|
||||
|
||||
result = await graph.ainvoke(Command(resume="answer"), config)
|
||||
assert result == {"value": ["sub_node", "sub:answer", "after"]}
|
||||
|
||||
history = [s async for s in graph.aget_state_history(config)]
|
||||
before_sub = [s for s in history if s.next == ("call_subgraph",)][-1]
|
||||
|
||||
called.clear()
|
||||
replay_result = await graph.ainvoke(None, before_sub.config)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Sub question?"
|
||||
assert "call_subgraph" in called
|
||||
assert "after" not in called
|
||||
|
||||
|
||||
async def test_fork_subgraph_two_interrupts_no_checkpointer(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Fork/replay with subgraph (no checkpointer) with two interrupt nodes.
|
||||
On fork, subgraph re-fires the first interrupt — consistent with
|
||||
top-level behavior."""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
class State(TypedDict):
|
||||
value: Annotated[list[str], operator.add]
|
||||
|
||||
def sub_int_1(state: State) -> State:
|
||||
called.append("sub_int_1")
|
||||
answer = interrupt("Sub Q1?")
|
||||
return {"value": [f"s1:{answer}"]}
|
||||
|
||||
def sub_int_2(state: State) -> State:
|
||||
called.append("sub_int_2")
|
||||
answer = interrupt("Sub Q2?")
|
||||
return {"value": [f"s2:{answer}"]}
|
||||
|
||||
subgraph = (
|
||||
StateGraph(State)
|
||||
.add_node("sub_int_1", sub_int_1)
|
||||
.add_node("sub_int_2", sub_int_2)
|
||||
.add_edge(START, "sub_int_1")
|
||||
.add_edge("sub_int_1", "sub_int_2")
|
||||
.compile()
|
||||
)
|
||||
|
||||
def call_subgraph(state: State) -> State:
|
||||
called.append("call_subgraph")
|
||||
return subgraph.invoke(state)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("call_subgraph", call_subgraph)
|
||||
.add_edge(START, "call_subgraph")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
result = await graph.ainvoke({"value": []}, config)
|
||||
assert "__interrupt__" in result
|
||||
assert result["__interrupt__"][0].value == "Sub Q1?"
|
||||
|
||||
result = await graph.ainvoke(Command(resume="a1"), config)
|
||||
assert "__interrupt__" in result
|
||||
assert result["__interrupt__"][0].value == "Sub Q2?"
|
||||
|
||||
result = await graph.ainvoke(Command(resume="a2"), config)
|
||||
assert result == {"value": ["s1:a1", "s2:a2"]}
|
||||
|
||||
history = [s async for s in graph.aget_state_history(config)]
|
||||
before_sub = [s for s in history if s.next == ("call_subgraph",)][-1]
|
||||
|
||||
called.clear()
|
||||
replay_result = await graph.ainvoke(None, before_sub.config)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Sub Q1?"
|
||||
|
||||
|
||||
async def test_concurrent_execution():
|
||||
"""Test concurrent execution with async nodes."""
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user