github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-08-01 19:06:01 +00:00
Code Issues Packages Projects Releases Wiki Activity
deer-flow/backend/tests/test_run_worker_rollback.py
Aari 244ce7739f
fix(runtime): linearize delta-mode checkpoint resume (#4460) 
* fix(runtime): linearize delta-mode checkpoint resume

Resuming a run from an older checkpoint forks the lineage, and in delta
mode that fork's state cannot be materialized correctly: the delta
history walk collects every pending_writes entry stored on each on-path
ancestor, but a shared parent also carries the writes of the sibling
child that was abandoned. Those writes replay into the fork, so the run
starts from a message list that still contains the answer it was meant to
replace — regenerating in a branched thread surfaced this as the
superseded assistant message reappearing beside the new one after a
reload. All three saver implementations are affected, so write-to-child
ownership is a gap in the upstream delta contract rather than one
saver's slip.

Rather than reimplement that walk, express the fork as what it means:
materialize the requested checkpoint's state, write it as an Overwrite on
the current head (which has no siblings), and run linearly. The abandoned
turn stays in history as the rewritten head's ancestry.

This runs after the rollback point is captured, so cancel-with-rollback
still restores the real pre-run head, and fails closed — an unreadable
resume checkpoint raises instead of falling back to the corrupt fork.
Full mode keeps forking: its checkpoints carry complete channel_values
and need no replay.

* fix(runtime): restore complete delta resume state

* fix(runtime): linearize delta rollback restoration

* Apply suggestions from code review

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

* fix(runtime): serialize delta resume preparation

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
2026-07-26 21:59:19 +08:00

2853 lines
104 KiB
Python
Raw Blame History

import asyncio
import copy
from contextlib import suppress
from types import SimpleNamespace
from typing import Annotated, Any, NotRequired, TypedDict
from unittest.mock import AsyncMock, call, patch
import pytest
from langchain_core.messages import AIMessage, AIMessageChunk, AnyMessage, HumanMessage
from langgraph.channels.delta import DeltaChannel
from langgraph.checkpoint.base import empty_checkpoint
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages
from langgraph.types import Overwrite
from deerflow.agents.thread_state import merge_artifacts, merge_message_writes
from deerflow.runtime.checkpoint_state import CheckpointStateAccessor
from deerflow.runtime.context_keys import CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY
from deerflow.runtime.runs.manager import CancelOutcome, ConflictError, RunManager
from deerflow.runtime.runs.schemas import RunStatus
from deerflow.runtime.runs.worker import (
RollbackPoint,
RunContext,
_agent_factory_supports_app_config,
_build_runtime_context,
_bump_channel_version,
_capture_rollback_point,
_collect_pre_existing_message_ids,
_ensure_interrupted_title,
_extract_llm_error_fallback_message,
_install_runtime_context,
_LargeFileToolChunkBatcher,
_rollback_to_pre_run_checkpoint,
_try_extract_from_message,
run_agent,
)
class FakeCheckpointer:
def __init__(self):
self.adelete_thread = AsyncMock()
self.aget_tuple = AsyncMock(return_value=None)
self.aput_writes = AsyncMock()
@pytest.mark.anyio
async def test_pending_cancel_stops_waiting_for_prior_finalization():
run_manager = RunManager()
prior = await run_manager.create("thread-cancel-while-waiting")
prior.status = RunStatus.interrupted
prior.finalizing = True
record = await run_manager.create("thread-cancel-while-waiting")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
factory_called = False
def agent_factory(**_kwargs):
nonlocal factory_called
factory_called = True
raise AssertionError("cancelled pending run must not build an agent")
record.task = asyncio.create_task(
run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=None),
agent_factory=agent_factory,
graph_input={"messages": []},
config={},
)
)
await asyncio.sleep(0)
outcome = await run_manager.cancel(record.run_id)
await asyncio.wait_for(record.task, timeout=0.2)
assert outcome == CancelOutcome.cancelled
assert prior.finalizing is True
assert factory_called is False
assert record.status == RunStatus.interrupted
bridge.publish_end.assert_awaited_once_with(record.run_id)
def _make_rollback_point(*, checkpoint_id="ckpt-1", messages=("before",), pending_writes=()):
materialized_messages = tuple(messages)
return RollbackPoint(
config={
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
"checkpoint_id": checkpoint_id,
}
},
state_values={},
messages=materialized_messages,
metadata={"source": "input"},
pending_writes=tuple(pending_writes),
)
def _stub_mutation_graph(monkeypatch, *, restored_config):
"""Replace the rollback mutation graph with a stub returning ``restored_config``."""
mock_graph = SimpleNamespace()
mock_graph.aupdate_state = AsyncMock(return_value=restored_config)
monkeypatch.setattr(
"deerflow.runtime.runs.worker.build_state_mutation_graph",
lambda *args, **kwargs: mock_graph,
)
return mock_graph
def _rollback_accessor(checkpointer):
return CheckpointStateAccessor(graph=SimpleNamespace(), checkpointer=checkpointer, mode="full")
class _DeltaChannelState(TypedDict):
messages: Annotated[list[AnyMessage], DeltaChannel(merge_message_writes, snapshot_frequency=1000)]
class _FullChannelState(TypedDict):
messages: Annotated[list[AnyMessage], add_messages]
def _build_message_append_graph(state_schema: type, checkpointer: Any):
async def _append_message(state: dict[str, Any]) -> dict[str, Any]:
return {"messages": [HumanMessage(content=f"turn-{len(state.get('messages') or [])}")]}
builder = StateGraph(state_schema)
builder.add_node("append_message", _append_message)
builder.set_entry_point("append_message")
builder.set_finish_point("append_message")
return builder.compile(checkpointer=checkpointer)
@pytest.mark.parametrize("tool_name", ["write_file", "str_replace"])
def test_large_file_tool_chunk_batcher_streams_bounded_batches(tool_name: str):
batcher = _LargeFileToolChunkBatcher(batch_size=2)
first = AIMessageChunk(
content="",
id="ai-1",
tool_call_chunks=[
{
"id": "call-1",
"index": 0,
"name": tool_name,
"args": '{"path":"/mnt/user-data/outputs/report.md","content":"Hel',
}
],
)
continuation = AIMessageChunk(
content="",
id="ai-1",
tool_call_chunks=[{"index": 0, "name": None, "args": 'lo"}'}],
)
assert batcher.push((first, {})) == []
published = batcher.push((continuation, {}))
assert len(published) == 1
message, metadata = published[0]
assert metadata == {}
assert message.tool_calls[0]["args"]["content"] == "Hello"
assert batcher.flush() == []
def test_large_file_tool_chunk_batcher_preserves_visible_and_non_file_chunks():
batcher = _LargeFileToolChunkBatcher()
visible_text = AIMessageChunk(content="Writing the report now.", id="ai-1")
search_tool = AIMessageChunk(
content="",
id="ai-2",
tool_call_chunks=[
{
"id": "call-2",
"index": 0,
"name": "web_search",
"args": '{"query":"vector databases"}',
}
],
)
write_with_reasoning = AIMessageChunk(
content="",
id="ai-3",
additional_kwargs={"reasoning_content": "Choosing a filename."},
tool_call_chunks=[
{
"id": "call-3",
"index": 0,
"name": "write_file",
"args": '{"path":"/mnt/user-data/outputs/report.md"}',
}
],
)
assert batcher.push((visible_text, {})) == [(visible_text, {})]
assert batcher.push((search_tool, {})) == [(search_tool, {})]
visible_reasoning = batcher.push((write_with_reasoning, {}))
assert len(visible_reasoning) == 1
filtered_message, filtered_metadata = visible_reasoning[0]
assert filtered_metadata == {}
assert filtered_message.additional_kwargs == {"reasoning_content": "Choosing a filename."}
assert filtered_message.tool_call_chunks == []
pending_file_chunks = batcher.flush()
assert len(pending_file_chunks) == 1
assert pending_file_chunks[0][0].tool_call_chunks[0]["name"] == "write_file"
def test_large_file_tool_chunk_batcher_separates_subgraph_namespaces():
batcher = _LargeFileToolChunkBatcher()
first = AIMessageChunk(
content="",
id="shared-ai-id",
tool_call_chunks=[{"id": "call-a", "index": 0, "name": "write_file", "args": '{"path":"a.md","content":"A'}],
)
second = AIMessageChunk(
content="",
id="shared-ai-id",
tool_call_chunks=[{"id": "call-b", "index": 0, "name": "write_file", "args": '{"path":"b.md","content":"B'}],
)
assert batcher.push((first, {"langgraph_checkpoint_ns": "task-a"})) == []
published = batcher.push((second, {"langgraph_checkpoint_ns": "task-b"}))
assert len(published) == 1
assert published[0][1]["langgraph_checkpoint_ns"] == "task-a"
assert batcher.flush()[0][1]["langgraph_checkpoint_ns"] == "task-b"
@pytest.mark.parametrize("metadata", [None, "not-a-dict"])
def test_large_file_tool_chunk_batcher_accepts_non_dict_metadata(metadata: Any):
batcher = _LargeFileToolChunkBatcher(batch_size=1)
message = AIMessageChunk(
content="",
id="ai-file",
tool_call_chunks=[
{
"id": "call-file",
"index": 0,
"name": "write_file",
"args": '{"path":"report.md","content":"draft"}',
}
],
)
published = batcher.push((message, metadata))
assert len(published) == 1
assert published[0][0].tool_call_chunks[0]["args"] == '{"path":"report.md","content":"draft"}'
assert published[0][1] == {}
def test_large_file_tool_chunk_batcher_does_not_retain_non_file_names():
batcher = _LargeFileToolChunkBatcher()
for index in range(100):
message = AIMessageChunk(
content="",
id=f"ai-{index}",
tool_call_chunks=[
{
"id": f"call-{index}",
"index": 0,
"name": "web_search",
"args": '{"query":"deerflow"}',
}
],
)
assert batcher.push((message, {})) == [(message, {})]
assert batcher.tool_names == {}
def test_large_file_tool_chunk_batcher_starts_batching_after_split_name_matches():
batcher = _LargeFileToolChunkBatcher(batch_size=1)
name_prefix = AIMessageChunk(
content="",
id="ai-file",
tool_call_chunks=[{"id": "call-file", "index": 0, "name": "write_", "args": ""}],
)
name_suffix = AIMessageChunk(
content="",
id="ai-file",
tool_call_chunks=[
{
"index": 0,
"name": "file",
"args": '{"path":"report.md","content":"draft"}',
}
],
)
assert batcher.push((name_prefix, {})) == [(name_prefix, {})]
assert set(batcher.tool_names.values()) == {"write_"}
assert len(batcher.push((name_suffix, {}))) == 1
assert set(batcher.tool_names.values()) == {"write_file"}
def test_large_file_tool_chunk_batcher_keeps_identity_across_batches_then_releases_it():
batcher = _LargeFileToolChunkBatcher(batch_size=1)
first = AIMessageChunk(
content="",
id="ai-file",
tool_call_chunks=[
{
"id": "call-file",
"index": 0,
"name": "write_file",
"args": '{"path":"report.md","content":"Hel',
}
],
)
continuation = AIMessageChunk(
content="",
id="ai-file",
tool_call_chunks=[{"index": 0, "name": None, "args": 'lo"}'}],
)
assert len(batcher.push((first, {}))) == 1
assert len(batcher.push((continuation, {}))) == 1
assert set(batcher.tool_names.values()) == {"write_file"}
assert batcher.finish() == []
assert batcher.tool_names == {}
def test_build_runtime_context_includes_app_config_when_present():
app_config = object()
context = _build_runtime_context("thread-1", "run-1", None, app_config)
assert context["thread_id"] == "thread-1"
assert context["run_id"] == "run-1"
assert context["app_config"] is app_config
def test_install_runtime_context_preserves_existing_thread_id_and_threads_app_config():
app_config = object()
config = {"context": {"thread_id": "caller-thread"}}
_install_runtime_context(
config,
{
"thread_id": "record-thread",
"run_id": "run-1",
"app_config": app_config,
},
)
assert config["context"]["thread_id"] == "caller-thread"
assert config["context"]["run_id"] == "run-1"
assert config["context"]["app_config"] is app_config
def test_install_runtime_context_overrides_internal_pre_existing_message_ids():
config = {"context": {CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY: {"spoofed"}}}
_install_runtime_context(
config,
{
"thread_id": "record-thread",
"run_id": "run-1",
CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY: frozenset({"old-ai"}),
},
)
assert config["context"][CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] == frozenset({"old-ai"})
@pytest.mark.anyio
async def test_run_agent_batches_incremental_file_args_and_keeps_complete_values():
run_manager = RunManager()
record = await run_manager.create("thread-file-stream")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
complete_message = AIMessage(
content="",
id="ai-file",
tool_calls=[
{
"id": "call-file",
"name": "write_file",
"args": {
"path": "/mnt/user-data/outputs/report.md",
"content": "Hello world",
},
"type": "tool_call",
}
],
)
class DummyAgent:
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
del graph_input, config, stream_mode, subgraphs
yield (
"messages",
(
AIMessageChunk(
content="",
id="ai-file",
tool_call_chunks=[
{
"id": "call-file",
"index": 0,
"name": "write_file",
"args": '{"path":"/mnt/user-data/outputs/report.md","content":"Hello',
}
],
),
{},
),
)
yield (
"messages",
(
AIMessageChunk(
content="",
id="ai-file",
tool_call_chunks=[{"index": 0, "name": None, "args": ' world"}'}],
),
{},
),
)
yield ("values", {"messages": [complete_message]})
await run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=None),
agent_factory=lambda **_kwargs: DummyAgent(),
graph_input={},
config={},
stream_modes=["messages-tuple", "values"],
)
message_events = [call.args for call in bridge.publish.await_args_list if call.args[1] == "messages"]
assert len(message_events) == 1
assert message_events[0][2][0]["tool_calls"][0]["args"]["content"] == "Hello world"
values_events = [call.args[2] for call in bridge.publish.await_args_list if call.args[1] == "values"]
assert any(event["messages"][0]["tool_calls"][0]["args"]["content"] == "Hello world" for event in values_events)
@pytest.mark.parametrize(
("stream_error", "flush_publish_error", "expected_error"),
[
(True, False, "stream failed"),
(True, True, "stream failed"),
(False, True, "flush publish failed"),
],
)
@pytest.mark.anyio
async def test_run_agent_handles_pending_file_args_when_stream_or_flush_raises(stream_error: bool, flush_publish_error: bool, expected_error: str):
run_manager = RunManager()
record = await run_manager.create("thread-file-stream-error")
async def publish(_run_id: str, event: str, _data: Any):
if flush_publish_error and event == "messages":
raise RuntimeError("flush publish failed")
bridge = SimpleNamespace(
publish=AsyncMock(side_effect=publish),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
class DummyAgent:
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
del graph_input, config, stream_mode, subgraphs
yield (
"messages",
(
AIMessageChunk(
content="",
id="ai-file",
tool_call_chunks=[
{
"id": "call-file",
"index": 0,
"name": "write_file",
"args": '{"path":"report.md","content":"partial',
}
],
),
{},
),
)
if stream_error:
raise RuntimeError("stream failed")
await run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=None),
agent_factory=lambda **_kwargs: DummyAgent(),
graph_input={},
config={},
stream_modes=["messages-tuple", "values"],
)
message_events = [call.args for call in bridge.publish.await_args_list if call.args[1] == "messages"]
assert len(message_events) == 1
assert message_events[0][2][0]["tool_call_chunks"][0]["args"].endswith('"content":"partial')
error_events = [call.args for call in bridge.publish.await_args_list if call.args[1] == "error"]
assert error_events[0][2]["message"] == expected_error
@pytest.mark.anyio
async def test_run_agent_keeps_file_chunks_unbatched_without_values_mode():
run_manager = RunManager()
record = await run_manager.create("thread-file-messages-only")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
chunks = [
AIMessageChunk(
content="",
id="ai-file",
tool_call_chunks=[
{
"id": "call-file",
"index": 0,
"name": "write_file",
"args": '{"path":"report.md","content":"Hel',
}
],
),
AIMessageChunk(
content="",
id="ai-file",
tool_call_chunks=[{"index": 0, "name": None, "args": 'lo"}'}],
),
]
class DummyAgent:
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
del graph_input, config, stream_mode, subgraphs
for chunk in chunks:
yield (chunk, {})
await run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=None),
agent_factory=lambda **_kwargs: DummyAgent(),
graph_input={},
config={},
stream_modes=["messages-tuple"],
)
message_events = [call.args[2] for call in bridge.publish.await_args_list if call.args[1] == "messages"]
assert len(message_events) == 2
assert message_events[0][0]["tool_call_chunks"][0]["args"].endswith('"content":"Hel')
assert message_events[1][0]["tool_call_chunks"][0]["args"] == 'lo"}'
@pytest.mark.anyio
async def test_run_agent_threads_explicit_app_config_into_config_only_factory():
run_manager = RunManager()
record = await run_manager.create("thread-1")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
app_config = object()
captured: dict[str, object] = {}
class DummyAgent:
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
captured["astream_context"] = config["context"]
yield {"messages": []}
def factory(*, config):
captured["factory_context"] = config["context"]
return DummyAgent()
await run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=None, app_config=app_config),
agent_factory=factory,
graph_input={},
config={},
)
await asyncio.sleep(0)
assert captured["factory_context"]["app_config"] is app_config
assert captured["astream_context"]["app_config"] is app_config
fetched = await run_manager.get(record.run_id)
assert fetched is not None
assert fetched.status == RunStatus.success
bridge.publish_end.assert_awaited_once_with(record.run_id)
bridge.cleanup.assert_awaited_once_with(record.run_id, delay=60)
@pytest.mark.anyio
async def test_run_agent_threads_pre_existing_message_ids_into_runtime_context():
run_manager = RunManager()
record = await run_manager.create("thread-1")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
captured: dict[str, object] = {}
class DummyCheckpointer:
async def aget_tuple(self, _config):
return SimpleNamespace(
config={"configurable": {"checkpoint_id": "checkpoint-1"}},
checkpoint={"channel_values": {}},
metadata={},
pending_writes=[],
)
class DummyAgent:
async def aget_state(self, _config):
return SimpleNamespace(
values={
"messages": [
HumanMessage(id="h1", content="question"),
AIMessage(id="a1", content="answer"),
]
},
config={
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
"checkpoint_id": "checkpoint-1",
}
},
parent_config=None,
metadata={},
next=(),
tasks=(),
created_at=None,
)
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
captured["context"] = config["context"]
yield {"messages": []}
def factory(*, config):
return DummyAgent()
await run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=DummyCheckpointer()),
agent_factory=factory,
graph_input={},
config={},
)
context = captured["context"]
assert context[CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] == frozenset({"h1", "a1"})
@pytest.mark.anyio
async def test_run_agent_marks_rollback_unusable_when_capture_fails():
"""A failed pre-run capture must disable rollback entirely.
Rollback now restores messages by forking the pre-run checkpoint through
the graph (raw checkpoint blobs cannot reconstruct Delta-channel history),
so there is no safe fallback when materialization fails: the worker must
report ``snapshot_capture_failed=True`` with no rollback point rather
than restoring an empty or partial message history.
"""
run_manager = RunManager()
record = await run_manager.create("thread-1")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
class DummyCheckpointer:
async def aget_tuple(self, _config):
return SimpleNamespace(
config={
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
"checkpoint_id": "checkpoint-1",
}
},
checkpoint={"id": "checkpoint-1", "channel_values": {}},
metadata={"source": "loop"},
pending_writes=[],
)
class DummyAgent:
async def aget_state(self, _config):
# Cancel after the worker crosses its startup barrier so this test
# exercises the running rollback path, not pending cancellation.
record.abort_action = "rollback"
record.abort_event.set()
raise RuntimeError("materialization failed")
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
yield {"messages": []}
with patch(
"deerflow.runtime.runs.worker._rollback_to_pre_run_checkpoint",
new_callable=AsyncMock,
) as rollback:
await run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=DummyCheckpointer()),
agent_factory=lambda *, config: DummyAgent(),
graph_input={},
config={},
)
rollback.assert_awaited_once()
rollback_kwargs = rollback.await_args.kwargs
assert rollback_kwargs["snapshot_capture_failed"] is True
assert rollback_kwargs["rollback_point"] is None
@pytest.mark.anyio
async def test_run_agent_overrides_spoofed_pre_existing_message_ids_without_snapshot():
run_manager = RunManager()
record = await run_manager.create("thread-1")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
captured: dict[str, object] = {}
class DummyAgent:
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
captured["context"] = config["context"]
yield {"messages": []}
def factory(*, config):
return DummyAgent()
await run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=None),
agent_factory=factory,
graph_input={},
config={"context": {CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY: {"spoofed"}}},
)
context = captured["context"]
assert context[CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] == frozenset()
@pytest.mark.anyio
async def test_run_agent_marks_llm_error_fallback_as_error_status():
run_manager = RunManager()
record = await run_manager.create("thread-1")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
class DummyAgent:
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
yield {
"messages": [
AIMessage(
content="The configured LLM provider is temporarily unavailable after multiple retries.",
additional_kwargs={
"deerflow_error_fallback": True,
"error_type": "APIConnectionError",
"error_reason": "transient",
"error_detail": "Connection error.",
},
)
]
}
def factory(*, config):
return DummyAgent()
await run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=None),
agent_factory=factory,
graph_input={},
config={},
)
fetched = await run_manager.get(record.run_id)
assert fetched is not None
assert fetched.status == RunStatus.error
assert fetched.error == "Connection error."
bridge.publish_end.assert_awaited_once_with(record.run_id)
@pytest.mark.anyio
async def test_run_agent_defaults_root_run_name_from_assistant_id():
run_manager = RunManager()
record = await run_manager.create("thread-1", assistant_id="lead_agent")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
captured: dict[str, object] = {}
class DummyAgent:
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
captured["astream_run_name"] = config["run_name"]
yield {"messages": []}
def factory(*, config):
captured["factory_run_name"] = config["run_name"]
return DummyAgent()
await run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=None),
agent_factory=factory,
graph_input={},
config={},
)
assert captured["factory_run_name"] == "lead_agent"
assert captured["astream_run_name"] == "lead_agent"
@pytest.mark.anyio
async def test_run_agent_defaults_root_run_name_from_context_agent_name():
run_manager = RunManager()
record = await run_manager.create("thread-1", assistant_id="lead_agent")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
captured: dict[str, object] = {}
class DummyAgent:
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
captured["astream_run_name"] = config["run_name"]
yield {"messages": []}
def factory(*, config):
captured["factory_run_name"] = config["run_name"]
return DummyAgent()
await run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=None),
agent_factory=factory,
graph_input={},
config={"context": {"agent_name": "finalis"}},
)
assert captured["factory_run_name"] == "finalis"
assert captured["astream_run_name"] == "finalis"
@pytest.mark.anyio
async def test_run_agent_defaults_root_run_name_from_configurable_agent_name():
run_manager = RunManager()
record = await run_manager.create("thread-1", assistant_id="lead_agent")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
captured: dict[str, object] = {}
class DummyAgent:
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
captured["astream_run_name"] = config["run_name"]
yield {"messages": []}
def factory(*, config):
captured["factory_run_name"] = config["run_name"]
return DummyAgent()
await run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=None),
agent_factory=factory,
graph_input={},
config={"configurable": {"agent_name": "finalis"}},
)
assert captured["factory_run_name"] == "finalis"
assert captured["astream_run_name"] == "finalis"
@pytest.mark.anyio
async def test_rollback_forks_pre_run_checkpoint_without_deleting_thread(monkeypatch):
checkpointer = FakeCheckpointer()
mock_graph = _stub_mutation_graph(
monkeypatch,
restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}},
)
rollback_point = _make_rollback_point(
pending_writes=[
("task-a", "messages", {"content": "first"}),
("task-a", "status", "done"),
("task-b", "events", {"type": "tool"}),
],
)
await _rollback_to_pre_run_checkpoint(
accessor=_rollback_accessor(checkpointer),
checkpointer=checkpointer,
thread_id="thread-1",
run_id="run-1",
rollback_point=rollback_point,
snapshot_capture_failed=False,
)
checkpointer.adelete_thread.assert_not_awaited()
mock_graph.aupdate_state.assert_awaited_once()
update_config, update_values = mock_graph.aupdate_state.await_args.args[:2]
assert update_config["configurable"]["checkpoint_id"] == "ckpt-1"
overwrite = update_values["messages"]
assert isinstance(overwrite, Overwrite)
assert overwrite.value == ["before"]
assert mock_graph.aupdate_state.await_args.kwargs["as_node"] == "rollback_restore"
assert checkpointer.aput_writes.await_args_list == [
call(
{"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}},
[("messages", {"content": "first"}), ("status", "done")],
task_id="task-a",
),
call(
{"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}},
[("events", {"type": "tool"})],
task_id="task-b",
),
]
@pytest.mark.anyio
async def test_rollback_skips_when_rollback_point_has_no_checkpoint_id(monkeypatch):
"""A rollback point without a checkpoint id cannot anchor a fork - skip safely."""
checkpointer = FakeCheckpointer()
mock_graph = _stub_mutation_graph(
monkeypatch,
restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}},
)
await _rollback_to_pre_run_checkpoint(
accessor=_rollback_accessor(checkpointer),
checkpointer=checkpointer,
thread_id="thread-1",
run_id="run-1",
rollback_point=_make_rollback_point(checkpoint_id=None),
snapshot_capture_failed=False,
)
checkpointer.adelete_thread.assert_not_awaited()
mock_graph.aupdate_state.assert_not_awaited()
checkpointer.aput_writes.assert_not_awaited()
class _ArtifactChannelState(TypedDict):
messages: Annotated[list[AnyMessage], add_messages]
artifacts: Annotated[list[str], merge_artifacts]
def _build_message_and_artifact_graph(checkpointer: Any):
async def _append(state: dict[str, Any]) -> dict[str, Any]:
n = len(state.get("messages") or [])
return {"messages": [HumanMessage(content=f"turn-{n}")], "artifacts": [f"a{n}"]}
builder = StateGraph(_ArtifactChannelState)
builder.add_node("append", _append)
builder.set_entry_point("append")
builder.set_finish_point("append")
return builder.compile(checkpointer=checkpointer)
class _ExtensionFullState(TypedDict):
"""Full-mode schema with a channel base ThreadState does not know about.
Mirrors a custom ``AgentMiddleware.state_schema`` contribution (memory,
todos, uploads, …): absent from the base schema the mutation graph falls
back to, so a base-schema restore silently drops it.
"""
messages: Annotated[list[AnyMessage], add_messages]
memory_notes: NotRequired[str]
class _ExtensionDeltaState(TypedDict):
messages: Annotated[list[AnyMessage], DeltaChannel(merge_message_writes, snapshot_frequency=1000)]
memory_notes: NotRequired[str]
@pytest.mark.anyio
@pytest.mark.parametrize(
"mode,state_schema",
[("full", _ExtensionFullState), ("delta", _ExtensionDeltaState)],
)
async def test_rollback_preserves_middleware_contributed_channels(mode, state_schema):
"""The rollback mutation graph must compile with the thread's effective schema.
Regression: ``_rollback_to_pre_run_checkpoint`` used to build the mutation
graph from the base ThreadState fallback, silently dropping channels
contributed by custom middleware from the restored checkpoint.
"""
checkpointer = InMemorySaver()
async def _step(state: dict[str, Any]) -> dict[str, Any]:
n = len(state.get("messages") or [])
return {"messages": [HumanMessage(content=f"turn-{n}")], "memory_notes": f"note-{n}"}
builder = StateGraph(state_schema)
builder.add_node("step", _step)
builder.set_entry_point("step")
builder.set_finish_point("step")
graph = builder.compile(checkpointer=checkpointer)
accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode=mode)
thread_config = {"configurable": {"thread_id": "thread-1"}}
await graph.ainvoke({}, thread_config)
rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config)
assert rollback_point is not None
await graph.ainvoke({}, thread_config)
await _rollback_to_pre_run_checkpoint(
accessor=accessor,
checkpointer=checkpointer,
thread_id="thread-1",
run_id="run-1",
rollback_point=rollback_point,
snapshot_capture_failed=False,
)
latest_snapshot = await accessor.aget(thread_config)
assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0"]
assert latest_snapshot.values["memory_notes"] == "note-0"
@pytest.mark.anyio
async def test_rollback_restores_non_message_channels_via_fork_inheritance():
"""Rollback overwrites only messages; sibling reducer channels must be
restored to their pre-run values by forking the pre-run checkpoint
(non-updated channels inherit from the parent). Locks in the load-bearing
assumption of the rollback design."""
checkpointer = InMemorySaver()
graph = _build_message_and_artifact_graph(checkpointer)
accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="full")
thread_config = {"configurable": {"thread_id": "thread-1"}}
await graph.ainvoke({}, thread_config)
rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config)
assert rollback_point is not None
await graph.ainvoke({}, thread_config)
cancelled_snapshot = await accessor.aget(thread_config)
assert cancelled_snapshot.values["artifacts"] == ["a0", "a1"]
assert [message.content for message in cancelled_snapshot.values["messages"]] == ["turn-0", "turn-1"]
await _rollback_to_pre_run_checkpoint(
accessor=accessor,
checkpointer=checkpointer,
thread_id="thread-1",
run_id="run-1",
rollback_point=rollback_point,
snapshot_capture_failed=False,
)
latest_snapshot = await accessor.aget(thread_config)
assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0"]
assert latest_snapshot.values["artifacts"] == ["a0"]
@pytest.mark.anyio
async def test_rollback_restored_checkpoint_becomes_latest_with_real_checkpointer():
"""Full-mode: the restored fork becomes latest, keeps pre-run values and
pre-run pending writes, and drops writes attached to the cancelled run."""
checkpointer = InMemorySaver()
graph = _build_message_append_graph(_FullChannelState, checkpointer)
accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="full")
thread_config = {"configurable": {"thread_id": "thread-1"}}
await graph.ainvoke({}, thread_config)
pre_run_snapshot = await accessor.aget(thread_config)
pre_run_checkpoint_id = pre_run_snapshot.config["configurable"]["checkpoint_id"]
pre_run_config = {
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
"checkpoint_id": pre_run_checkpoint_id,
}
}
checkpointer.put_writes(pre_run_config, [("title", "pending-before")], task_id="task-before")
rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config)
assert rollback_point is not None
assert rollback_point.pending_writes == (("task-before", "title", "pending-before"),)
await graph.ainvoke({}, thread_config)
cancelled_snapshot = await accessor.aget(thread_config)
cancelled_checkpoint_id = cancelled_snapshot.config["configurable"]["checkpoint_id"]
checkpointer.put_writes(
{"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": cancelled_checkpoint_id}},
[("title", "pending-after")],
task_id="task-after",
)
await _rollback_to_pre_run_checkpoint(
accessor=accessor,
checkpointer=checkpointer,
thread_id="thread-1",
run_id="run-1",
rollback_point=rollback_point,
snapshot_capture_failed=False,
)
latest_snapshot = await accessor.aget(thread_config)
restored_checkpoint_id = latest_snapshot.config["configurable"]["checkpoint_id"]
assert restored_checkpoint_id not in (pre_run_checkpoint_id, cancelled_checkpoint_id)
assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0"]
latest_tuple = await checkpointer.aget_tuple(thread_config)
assert latest_tuple.pending_writes == [("task-before", "title", "pending-before")]
restored_tuple = await checkpointer.aget_tuple({"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": restored_checkpoint_id}})
assert restored_tuple.parent_config["configurable"]["checkpoint_id"] == pre_run_checkpoint_id
@pytest.mark.anyio
async def test_rollback_deletes_thread_when_no_snapshot_exists():
checkpointer = FakeCheckpointer()
await _rollback_to_pre_run_checkpoint(
accessor=_rollback_accessor(checkpointer),
checkpointer=checkpointer,
thread_id="thread-1",
run_id="run-1",
rollback_point=None,
snapshot_capture_failed=False,
)
checkpointer.adelete_thread.assert_awaited_once_with("thread-1")
checkpointer.aput_writes.assert_not_awaited()
@pytest.mark.anyio
async def test_rollback_raises_when_restore_config_has_no_checkpoint_id(monkeypatch):
checkpointer = FakeCheckpointer()
_stub_mutation_graph(
monkeypatch,
restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}},
)
with pytest.raises(RuntimeError, match="did not return checkpoint_id"):
await _rollback_to_pre_run_checkpoint(
accessor=_rollback_accessor(checkpointer),
checkpointer=checkpointer,
thread_id="thread-1",
run_id="run-1",
rollback_point=_make_rollback_point(pending_writes=[("task-a", "messages", "value")]),
snapshot_capture_failed=False,
)
checkpointer.adelete_thread.assert_not_awaited()
checkpointer.aput_writes.assert_not_awaited()
@pytest.mark.anyio
async def test_capture_rollback_point_normalizes_missing_checkpoint_ns():
"""Snapshots whose config omits ``checkpoint_ns`` must anchor the root namespace."""
class _TupleCheckpointer:
async def aget_tuple(self, _config):
return SimpleNamespace(pending_writes=[])
class _SnapshotAgent:
async def aget_state(self, _config):
return SimpleNamespace(
values={"messages": [HumanMessage(id="h1", content="before")]},
config={"configurable": {"thread_id": "thread-1", "checkpoint_id": "ckpt-1"}},
metadata={"source": "loop"},
)
checkpointer = _TupleCheckpointer()
accessor = CheckpointStateAccessor(graph=_SnapshotAgent(), checkpointer=checkpointer, mode="full")
rollback_point = await _capture_rollback_point(accessor, checkpointer, {"configurable": {"thread_id": "thread-1"}})
assert rollback_point is not None
assert rollback_point.config["configurable"]["checkpoint_ns"] == ""
assert rollback_point.config["configurable"]["checkpoint_id"] == "ckpt-1"
assert [message.content for message in rollback_point.messages] == ["before"]
assert rollback_point.metadata == {"source": "loop"}
@pytest.mark.anyio
async def test_capture_rollback_point_returns_none_without_checkpoint():
"""A thread with no checkpoints has no rollback point (delete/reset contract)."""
class _EmptyAgent:
async def aget_state(self, _config):
return SimpleNamespace(
values={},
config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}},
metadata={},
)
checkpointer = SimpleNamespace(aget_tuple=AsyncMock(return_value=None))
accessor = CheckpointStateAccessor(graph=_EmptyAgent(), checkpointer=checkpointer, mode="full")
assert await _capture_rollback_point(accessor, checkpointer, {"configurable": {"thread_id": "thread-1"}}) is None
@pytest.mark.anyio
async def test_rollback_raises_on_malformed_pending_write_not_a_tuple(monkeypatch):
"""pending_writes containing a non-3-tuple item should raise RuntimeError."""
checkpointer = FakeCheckpointer()
_stub_mutation_graph(
monkeypatch,
restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}},
)
with pytest.raises(RuntimeError, match="rollback failed: pending_write is not a 3-tuple"):
await _rollback_to_pre_run_checkpoint(
accessor=_rollback_accessor(checkpointer),
checkpointer=checkpointer,
thread_id="thread-1",
run_id="run-1",
rollback_point=_make_rollback_point(
pending_writes=[
("task-a", "messages", "valid"), # valid
["only", "two"], # malformed: only 2 elements
],
),
snapshot_capture_failed=False,
)
checkpointer.aput_writes.assert_not_awaited()
@pytest.mark.anyio
async def test_rollback_raises_on_malformed_pending_write_non_string_channel(monkeypatch):
"""pending_writes containing a non-string channel should raise RuntimeError."""
checkpointer = FakeCheckpointer()
_stub_mutation_graph(
monkeypatch,
restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}},
)
with pytest.raises(RuntimeError, match="rollback failed: pending_write has non-string channel"):
await _rollback_to_pre_run_checkpoint(
accessor=_rollback_accessor(checkpointer),
checkpointer=checkpointer,
thread_id="thread-1",
run_id="run-1",
rollback_point=_make_rollback_point(
pending_writes=[
("task-a", 123, "value"), # malformed: channel is not a string
],
),
snapshot_capture_failed=False,
)
checkpointer.aput_writes.assert_not_awaited()
@pytest.mark.anyio
async def test_rollback_propagates_aput_writes_failure(monkeypatch):
"""If aput_writes fails, the exception should propagate (not be swallowed)."""
checkpointer = FakeCheckpointer()
checkpointer.aput_writes.side_effect = RuntimeError("Database connection lost")
_stub_mutation_graph(
monkeypatch,
restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}},
)
with pytest.raises(RuntimeError, match="Database connection lost"):
await _rollback_to_pre_run_checkpoint(
accessor=_rollback_accessor(checkpointer),
checkpointer=checkpointer,
thread_id="thread-1",
run_id="run-1",
rollback_point=_make_rollback_point(
pending_writes=[
("task-a", "messages", "value"),
],
),
snapshot_capture_failed=False,
)
checkpointer.aput_writes.assert_awaited_once()
@pytest.mark.anyio
async def test_rollback_linearizes_delta_restore_onto_cancelled_head():
"""Delta rollback replaces state on the head instead of creating a fork.
The cancelled path has already attached writes to the pre-run checkpoint,
so forking that checkpoint would replay sibling writes. The restored
checkpoint must instead descend from the cancelled head and replace the
captured messages there.
"""
checkpointer = InMemorySaver()
graph = _build_message_append_graph(_DeltaChannelState, checkpointer)
accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="delta")
thread_config = {"configurable": {"thread_id": "thread-1"}}
await graph.ainvoke({}, thread_config)
await graph.ainvoke({}, thread_config)
rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config)
assert rollback_point is not None
assert [message.content for message in rollback_point.messages] == ["turn-0", "turn-1"]
pre_run_checkpoint_id = rollback_point.config["configurable"]["checkpoint_id"]
await graph.ainvoke({}, thread_config) # the run that gets cancelled
cancelled_snapshot = await accessor.aget(thread_config)
cancelled_checkpoint_id = cancelled_snapshot.config["configurable"]["checkpoint_id"]
assert len(cancelled_snapshot.values["messages"]) == 3
await _rollback_to_pre_run_checkpoint(
accessor=accessor,
checkpointer=checkpointer,
thread_id="thread-1",
run_id="run-1",
rollback_point=rollback_point,
snapshot_capture_failed=False,
)
latest_snapshot = await accessor.aget(thread_config)
restored_checkpoint_id = latest_snapshot.config["configurable"]["checkpoint_id"]
assert restored_checkpoint_id not in (pre_run_checkpoint_id, cancelled_checkpoint_id)
assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0", "turn-1"]
restored_tuple = await checkpointer.aget_tuple({"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": restored_checkpoint_id}})
assert restored_tuple.parent_config["configurable"]["checkpoint_id"] == cancelled_checkpoint_id
assert restored_tuple.metadata.get("source") == "update"
# A non-snapshot Delta checkpoint must not persist the full message list.
raw_messages = restored_tuple.checkpoint.get("channel_values", {}).get("messages")
assert not isinstance(raw_messages, list)
@pytest.mark.anyio
async def test_rollback_restores_pre_run_pending_writes_for_delta_checkpoints():
"""Pre-run pending writes are re-attached to the restored checkpoint; writes
attached to the cancelled run are not."""
checkpointer = InMemorySaver()
graph = _build_message_append_graph(_DeltaChannelState, checkpointer)
accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="delta")
thread_config = {"configurable": {"thread_id": "thread-1"}}
await graph.ainvoke({}, thread_config)
pre_run_snapshot = await accessor.aget(thread_config)
pre_run_checkpoint_id = pre_run_snapshot.config["configurable"]["checkpoint_id"]
checkpointer.put_writes(
{"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": pre_run_checkpoint_id}},
[("title", "in-flight-title")],
task_id="task-in-flight",
)
rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config)
assert rollback_point is not None
assert rollback_point.pending_writes == (("task-in-flight", "title", "in-flight-title"),)
await graph.ainvoke({}, thread_config) # cancelled run
cancelled_snapshot = await accessor.aget(thread_config)
cancelled_checkpoint_id = cancelled_snapshot.config["configurable"]["checkpoint_id"]
checkpointer.put_writes(
{"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": cancelled_checkpoint_id}},
[("title", "cancelled-title")],
task_id="task-cancelled",
)
await _rollback_to_pre_run_checkpoint(
accessor=accessor,
checkpointer=checkpointer,
thread_id="thread-1",
run_id="run-1",
rollback_point=rollback_point,
snapshot_capture_failed=False,
)
latest_tuple = await checkpointer.aget_tuple(thread_config)
assert latest_tuple.pending_writes == [("task-in-flight", "title", "in-flight-title")]
assert [message.content for message in (await accessor.aget(thread_config)).values["messages"]] == ["turn-0"]
@pytest.mark.anyio
async def test_rollback_linearizes_delta_restore_sqlite_reopen(tmp_path):
"""Same linear restore contract against a disk-backed saver, verified after a
close/reopen so only persisted bytes can satisfy the assertions."""
db_path = tmp_path / "rollback.sqlite3"
async with AsyncSqliteSaver.from_conn_string(str(db_path)) as checkpointer:
await checkpointer.setup()
graph = _build_message_append_graph(_DeltaChannelState, checkpointer)
accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="delta")
thread_config = {"configurable": {"thread_id": "thread-1"}}
await graph.ainvoke({}, thread_config)
await graph.ainvoke({}, thread_config)
rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config)
assert rollback_point is not None
pre_run_checkpoint_id = rollback_point.config["configurable"]["checkpoint_id"]
await graph.ainvoke({}, thread_config) # cancelled run
cancelled_snapshot = await accessor.aget(thread_config)
cancelled_checkpoint_id = cancelled_snapshot.config["configurable"]["checkpoint_id"]
await _rollback_to_pre_run_checkpoint(
accessor=accessor,
checkpointer=checkpointer,
thread_id="thread-1",
run_id="run-1",
rollback_point=rollback_point,
snapshot_capture_failed=False,
)
restored_snapshot = await accessor.aget(thread_config)
restored_checkpoint_id = restored_snapshot.config["configurable"]["checkpoint_id"]
assert restored_checkpoint_id not in (pre_run_checkpoint_id, cancelled_checkpoint_id)
assert [message.content for message in restored_snapshot.values["messages"]] == ["turn-0", "turn-1"]
restored_tuple = await checkpointer.aget_tuple({"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": restored_checkpoint_id}})
assert restored_tuple.parent_config["configurable"]["checkpoint_id"] == cancelled_checkpoint_id
async with AsyncSqliteSaver.from_conn_string(str(db_path)) as checkpointer:
graph = _build_message_append_graph(_DeltaChannelState, checkpointer)
accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="delta")
thread_config = {"configurable": {"thread_id": "thread-1"}}
latest_snapshot = await accessor.aget(thread_config)
assert latest_snapshot.config["configurable"]["checkpoint_id"] == restored_checkpoint_id
assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0", "turn-1"]
restored_tuple = await checkpointer.aget_tuple({"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": restored_checkpoint_id}})
assert restored_tuple.parent_config["configurable"]["checkpoint_id"] == cancelled_checkpoint_id
raw_messages = restored_tuple.checkpoint.get("channel_values", {}).get("messages")
assert not isinstance(raw_messages, list)
def test_agent_factory_supports_app_config_detects_supported_signature():
def factory(*, config, app_config=None):
return (config, app_config)
assert _agent_factory_supports_app_config(factory) is True
def test_build_runtime_context_defaults_to_thread_and_run_id():
ctx = _build_runtime_context("thread-1", "run-1", None)
assert ctx == {"thread_id": "thread-1", "run_id": "run-1"}
def test_build_runtime_context_merges_caller_context():
"""Regression for issue #2677: keys from ``config['context']`` (e.g. ``agent_name``)
must be merged into the Runtime's context so that ``ToolRuntime.context`` — which
is what ``setup_agent`` reads — can see them."""
caller_context = {"agent_name": "my-agent", "is_bootstrap": True, "model_name": "gpt-4"}
ctx = _build_runtime_context("thread-1", "run-1", caller_context)
assert ctx["thread_id"] == "thread-1"
assert ctx["run_id"] == "run-1"
assert ctx["agent_name"] == "my-agent"
assert ctx["is_bootstrap"] is True
assert ctx["model_name"] == "gpt-4"
def test_build_runtime_context_caller_cannot_override_thread_id_or_run_id():
"""A malicious or buggy caller must not be able to overwrite the worker-assigned
``thread_id`` / ``run_id`` by stuffing them into ``config['context']``."""
caller_context = {"thread_id": "spoofed", "run_id": "spoofed", "agent_name": "ok"}
ctx = _build_runtime_context("real-thread", "real-run", caller_context)
assert ctx["thread_id"] == "real-thread"
assert ctx["run_id"] == "real-run"
assert ctx["agent_name"] == "ok"
def test_build_runtime_context_ignores_caller_pre_existing_message_ids():
caller_context = {CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY: {"spoofed"}}
ctx = _build_runtime_context("thread-1", "run-1", caller_context)
assert CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY not in ctx
def test_build_runtime_context_ignores_non_dict_caller_context():
ctx = _build_runtime_context("thread-1", "run-1", "not-a-dict")
assert ctx == {"thread_id": "thread-1", "run_id": "run-1"}
def test_agent_factory_supports_app_config_returns_false_when_signature_lookup_fails(monkeypatch):
class BrokenCallable:
def __call__(self, **kwargs):
return kwargs
monkeypatch.setattr("deerflow.runtime.runs.worker.inspect.signature", lambda _obj: (_ for _ in ()).throw(ValueError("boom")))
assert _agent_factory_supports_app_config(BrokenCallable()) is False
# ---------------------------------------------------------------------------
# _extract_llm_error_fallback_message coverage
# ---------------------------------------------------------------------------
def test_try_extract_from_message_finds_fallback_on_message_object():
msg = AIMessage(
content="fallback",
additional_kwargs={
"deerflow_error_fallback": True,
"error_detail": "Connection error.",
"error_reason": "transient",
},
)
assert _try_extract_from_message(msg) == "Connection error."
def test_try_extract_from_message_finds_fallback_on_dict():
msg = {
"content": "fallback",
"additional_kwargs": {
"deerflow_error_fallback": True,
"error_detail": "Quota exceeded.",
},
}
assert _try_extract_from_message(msg) == "Quota exceeded."
def test_try_extract_from_message_returns_none_for_normal_message():
msg = AIMessage(content="hello")
assert _try_extract_from_message(msg) is None
def test_extract_llm_error_fallback_message_large_state_chunk_no_fallback():
"""Normal-size state dict without fallback markers must not raise and should return None."""
large_state = {
"messages": [
AIMessage(content="Hello!"),
{"role": "user", "content": "Hi there"},
],
"foo": "x" * 10_000,
"bar": {"nested": {"deep": {"data": list(range(1000))}}},
"baz": [{"id": i, "payload": "y" * 1000} for i in range(500)],
}
assert _extract_llm_error_fallback_message(large_state) is None
def test_extract_llm_error_fallback_message_finds_fallback_in_messages_list():
state = {
"messages": [
AIMessage(content="Hello!"),
AIMessage(
content="Unavailable.",
additional_kwargs={
"deerflow_error_fallback": True,
"error_detail": "Connection error.",
},
),
],
"other_state": "large_value" * 1000,
}
assert _extract_llm_error_fallback_message(state) == "Connection error."
def test_extract_llm_error_fallback_message_finds_fallback_in_raw_message():
msg = AIMessage(
content="Unavailable.",
additional_kwargs={
"deerflow_error_fallback": True,
"error_reason": "quota",
},
)
assert _extract_llm_error_fallback_message(msg) == "quota"
def test_extract_llm_error_fallback_message_finds_fallback_in_tuple():
item = (
"messages",
AIMessage(
content="Unavailable.",
additional_kwargs={
"deerflow_error_fallback": True,
"error_detail": "Circuit open.",
},
),
)
assert _extract_llm_error_fallback_message(item) == "Circuit open."
def test_extract_llm_error_fallback_message_returns_none_for_empty_values():
assert _extract_llm_error_fallback_message({}) is None
assert _extract_llm_error_fallback_message([]) is None
assert _extract_llm_error_fallback_message(None) is None
assert _extract_llm_error_fallback_message("string") is None
def test_extract_llm_error_fallback_message_finds_fallback_in_updates_mode():
"""stream_mode='updates' yields dicts keyed by node name (e.g. {'call_model': {...}}).
Fallback marker is nested inside the node's state update, not at the top level."""
update_chunk = {
"call_model": {
"messages": [
AIMessage(
content="Unavailable.",
additional_kwargs={
"deerflow_error_fallback": True,
"error_detail": "Connection error.",
},
)
]
}
}
assert _extract_llm_error_fallback_message(update_chunk) == "Connection error."
def test_extract_llm_error_fallback_message_updates_mode_no_fallback():
"""Normal updates chunk without any fallback should return None safely."""
update_chunk = {
"__interrupt__": [
{
"value": "ask_human",
"resumable": True,
"ns": ["agent"],
"when": "during",
}
]
}
assert _extract_llm_error_fallback_message(update_chunk) is None
# ---------------------------------------------------------------------------
# pre_existing_ids filtering — stale fallback markers from prior runs
# ---------------------------------------------------------------------------
def test_try_extract_skips_message_with_pre_existing_id():
"""Fallback marker on a message whose id is in pre_existing_ids must be ignored."""
msg = AIMessage(
id="stale-1",
content="Unavailable.",
additional_kwargs={
"deerflow_error_fallback": True,
"error_detail": "Connection error.",
},
)
assert _try_extract_from_message(msg, {"stale-1"}) is None
# Without the filter, the same message would still surface the marker.
assert _try_extract_from_message(msg) == "Connection error."
def test_try_extract_still_finds_fresh_message_when_others_are_stale():
"""A non-stale message with a fallback marker must still match."""
msg = AIMessage(
id="fresh-1",
content="Unavailable.",
additional_kwargs={
"deerflow_error_fallback": True,
"error_detail": "Connection error.",
},
)
assert _try_extract_from_message(msg, {"stale-1", "stale-2"}) == "Connection error."
def test_try_extract_skips_dict_message_with_pre_existing_id():
msg = {
"id": "stale-2",
"content": "Unavailable.",
"additional_kwargs": {
"deerflow_error_fallback": True,
"error_detail": "Quota exceeded.",
},
}
assert _try_extract_from_message(msg, {"stale-2"}) is None
assert _try_extract_from_message(msg) == "Quota exceeded."
def test_extract_llm_error_fallback_message_skips_stale_history():
"""A state chunk replaying a stale fallback marker from a prior run must return None."""
state = {
"messages": [
AIMessage(id="stale-1", content="Hi"),
AIMessage(
id="stale-fallback",
content="Unavailable.",
additional_kwargs={
"deerflow_error_fallback": True,
"error_detail": "Connection error.",
},
),
]
}
assert _extract_llm_error_fallback_message(state, {"stale-1", "stale-fallback"}) is None
def test_extract_llm_error_fallback_message_returns_fresh_marker_alongside_stale_history():
"""Stale history is ignored, but a brand-new fallback in the same chunk is reported."""
state = {
"messages": [
AIMessage(id="stale-1", content="Hi"),
AIMessage(
id="stale-fallback",
content="Old failure.",
additional_kwargs={
"deerflow_error_fallback": True,
"error_detail": "Old error.",
},
),
AIMessage(
id="fresh-fallback",
content="New failure.",
additional_kwargs={
"deerflow_error_fallback": True,
"error_detail": "Fresh error.",
},
),
]
}
assert _extract_llm_error_fallback_message(state, {"stale-1", "stale-fallback"}) == "Fresh error."
def test_extract_llm_error_fallback_message_default_filter_is_empty():
"""Passing no pre_existing_ids must preserve the original (pre-fix) behavior."""
state = {
"messages": [
AIMessage(
id="any",
content="Unavailable.",
additional_kwargs={
"deerflow_error_fallback": True,
"error_detail": "Connection error.",
},
)
]
}
assert _extract_llm_error_fallback_message(state) == "Connection error."
def test_collect_pre_existing_message_ids_pulls_ids_from_values():
values = {
"messages": [
AIMessage(id="a", content="x"),
AIMessage(id="b", content="y"),
AIMessage(content="no-id-here"), # ignored
]
}
assert _collect_pre_existing_message_ids(values) == {"a", "b"}
def test_collect_pre_existing_message_ids_handles_missing_pieces():
assert _collect_pre_existing_message_ids(None) == set()
assert _collect_pre_existing_message_ids({}) == set()
assert _collect_pre_existing_message_ids({"messages": None}) == set()
@pytest.mark.anyio
async def test_run_agent_ignores_stale_llm_error_fallback_from_prior_run():
"""A stale fallback marker checkpointed by an earlier run on the same thread
must NOT cause a successful current run to be reported as ``error``.
This guards against the regression where one IndexError-driven failure (now
classified transient and surfaced as a ``deerflow_error_fallback`` AIMessage)
persisted in thread history and tripped ``RunStatus.error`` on every
subsequent run that re-played the messages channel via ``stream_mode="values"``.
"""
run_manager = RunManager()
record = await run_manager.create("thread-1")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
stale_fallback = AIMessage(
id="stale-fallback",
content="Old failure.",
additional_kwargs={
"deerflow_error_fallback": True,
"error_type": "IndexError",
"error_reason": "transient",
"error_detail": "list index out of range",
},
)
class StaleHistoryCheckpointer:
async def aget_tuple(self, config):
checkpoint = empty_checkpoint()
checkpoint["id"] = "ckpt-stale"
checkpoint["channel_values"] = {"messages": [stale_fallback]}
return SimpleNamespace(
config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "ckpt-stale"}},
checkpoint=checkpoint,
metadata={},
pending_writes=[],
)
class DummyAgent:
async def aget_state(self, _config):
return SimpleNamespace(
values={"messages": [stale_fallback]},
config={
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
"checkpoint_id": "ckpt-stale",
}
},
parent_config=None,
metadata={},
next=(),
tasks=(),
created_at=None,
)
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
# Replay the prior fallback message (as LangGraph would when using
# stream_mode="values") and then yield a fresh successful AIMessage.
yield {
"messages": [
stale_fallback,
AIMessage(id="fresh-ok", content="Hello — the run succeeded."),
]
}
def factory(*, config):
return DummyAgent()
await run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=StaleHistoryCheckpointer()),
agent_factory=factory,
graph_input={},
config={},
)
fetched = await run_manager.get(record.run_id)
assert fetched is not None
assert fetched.status == RunStatus.success, f"Stale fallback marker from prior run should not flip current run to error, got status={fetched.status} error={fetched.error!r}"
bridge.publish_end.assert_awaited_once_with(record.run_id)
class _FakeCheckpointTuple:
"""Minimal stand-in for ``CheckpointTuple`` used by ``_ensure_interrupted_title``."""
def __init__(self, *, checkpoint: dict, metadata: dict, config: dict | None = None):
self.checkpoint = checkpoint
self.metadata = metadata
self.config = config or {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}}
class _TitleCheckpointer:
"""Captures ``aput`` arguments and exposes ``get_next_version`` like DB savers."""
def __init__(self, *, tuple_value: _FakeCheckpointTuple | None, put_result: dict | None = None):
self.aget_tuple = AsyncMock(return_value=tuple_value)
self.aput = AsyncMock(return_value=put_result or {})
def get_next_version(self, current, _channel):
if current is None:
return 1
if isinstance(current, int):
return current + 1
if isinstance(current, str):
try:
return str(int(current) + 1)
except ValueError:
return f"{current}.1"
return 1
@pytest.mark.anyio
async def test_interrupted_title_finalization_blocks_new_same_thread_run(monkeypatch):
"""A cancelled run must remain active while its title-only checkpoint is finalizing."""
from deerflow.agents.middlewares.title_middleware import TitleMiddleware
monkeypatch.setattr(
TitleMiddleware,
"_generate_title_result",
lambda self, state, allow_partial_exchange=False: {"title": "Old Prompt"},
)
initial_checkpoint = {
"id": "ckpt-old",
"ts": "2026-06-29T00:00:00Z",
"channel_values": {"messages": [{"type": "human", "content": "old prompt"}]},
"channel_versions": {"messages": 1},
}
class _BlockingTitleCheckpointer:
def __init__(self) -> None:
self.latest_checkpoint = copy.deepcopy(initial_checkpoint)
self.latest_metadata = {"source": "loop", "step": 1}
self.title_write_started = asyncio.Event()
self.release_title_write = asyncio.Event()
async def aget_tuple(self, config):
del config
return _FakeCheckpointTuple(
checkpoint=copy.deepcopy(self.latest_checkpoint),
metadata=dict(self.latest_metadata),
config={
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
"checkpoint_id": self.latest_checkpoint["id"],
}
},
)
async def aput(self, config, checkpoint, metadata, new_versions):
del config, new_versions
self.title_write_started.set()
await self.release_title_write.wait()
self.latest_checkpoint = copy.deepcopy(checkpoint)
self.latest_metadata = dict(metadata)
return {
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
"checkpoint_id": checkpoint["id"],
}
}
def get_next_version(self, current, _channel):
return (current or 0) + 1
run_manager = RunManager()
record = await run_manager.create("thread-1")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
checkpointer = _BlockingTitleCheckpointer()
class _AbortingAgent:
metadata = {"model_name": "fake-test-model"}
checkpointer: Any | None = None
store: Any | None = None
interrupt_before_nodes = None
interrupt_after_nodes = None
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
del graph_input, config, stream_mode, subgraphs
record.abort_event.set()
if False:
yield # pragma: no cover
def factory(*, config):
del config
return _AbortingAgent()
task = asyncio.create_task(
run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=checkpointer),
agent_factory=factory,
graph_input={"messages": [{"role": "user", "content": "old prompt"}]},
config={},
)
)
record.task = task
try:
await asyncio.wait_for(checkpointer.title_write_started.wait(), timeout=1.0)
with pytest.raises(ConflictError, match="active run"):
await run_manager.create_or_reject("thread-1", multitask_strategy="reject")
finally:
checkpointer.release_title_write.set()
await task
records = await run_manager.list_by_thread("thread-1")
assert [item.run_id for item in records] == [record.run_id]
assert checkpointer.latest_checkpoint["channel_values"]["messages"] == [{"type": "human", "content": "old prompt"}]
@pytest.mark.anyio
async def test_finalizing_run_only_blocks_reject_strategy():
"""A finalizing run must not break interrupt/rollback superseding semantics."""
async def _seed_finalizing_run():
run_manager = RunManager()
record = await run_manager.create("thread-1")
release_cleanup = asyncio.Event()
cleanup_cancelled = asyncio.Event()
async def _cleanup_task():
try:
await release_cleanup.wait()
except asyncio.CancelledError:
cleanup_cancelled.set()
raise
task = asyncio.create_task(_cleanup_task())
record.task = task
await run_manager.set_status(record.run_id, RunStatus.interrupted)
await run_manager.set_finalizing(record.run_id, True)
return run_manager, record, task, release_cleanup, cleanup_cancelled
for strategy in ("interrupt", "rollback"):
run_manager, record, task, release_cleanup, cleanup_cancelled = await _seed_finalizing_run()
try:
replacement = await run_manager.create_or_reject("thread-1", multitask_strategy=strategy)
await asyncio.sleep(0)
assert replacement.run_id != record.run_id
assert record.status == RunStatus.interrupted
assert record.finalizing is True
assert not cleanup_cancelled.is_set()
assert not task.done()
finally:
release_cleanup.set()
with suppress(asyncio.CancelledError):
await task
run_manager, _record, task, release_cleanup, _cleanup_cancelled = await _seed_finalizing_run()
try:
with pytest.raises(ConflictError, match="active run"):
await run_manager.create_or_reject("thread-1", multitask_strategy="reject")
finally:
release_cleanup.set()
with suppress(asyncio.CancelledError):
await task
@pytest.mark.anyio
async def test_admitted_pending_replacement_does_not_steal_interrupted_title_recovery(monkeypatch):
"""The old run must still write the fallback title before releasing a serialized replacement."""
from deerflow.agents.middlewares.title_middleware import TitleMiddleware
monkeypatch.setattr(
TitleMiddleware,
"_generate_title_result",
lambda self, state, allow_partial_exchange=False: {"title": "Old Prompt"},
)
initial_checkpoint = {
"id": "ckpt-old",
"ts": "2026-06-29T00:00:00Z",
"channel_values": {"messages": [{"type": "human", "content": "Old prompt"}]},
"channel_versions": {"messages": 1},
}
run_manager = RunManager()
old_record = await run_manager.create("thread-1")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
checkpointer = _TitleCheckpointer(
tuple_value=_FakeCheckpointTuple(
checkpoint=initial_checkpoint,
metadata={"source": "loop", "step": 1},
config={
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
"checkpoint_id": "ckpt-old",
}
},
),
)
old_title_gate_entered = asyncio.Event()
release_old_title_gate = asyncio.Event()
original_wait_for_prior_finalizing = run_manager.wait_for_prior_finalizing
async def _wait_for_prior_finalizing(thread_id, run_id, **kwargs):
if run_id == old_record.run_id and old_record.status == RunStatus.interrupted and old_record.finalizing:
old_title_gate_entered.set()
await release_old_title_gate.wait()
return await original_wait_for_prior_finalizing(thread_id, run_id, **kwargs)
run_manager.wait_for_prior_finalizing = _wait_for_prior_finalizing # type: ignore[method-assign]
class _AbortingAgent:
metadata = {"model_name": "fake-test-model"}
checkpointer: Any | None = None
store: Any | None = None
interrupt_before_nodes = None
interrupt_after_nodes = None
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
del graph_input, config, stream_mode, subgraphs
old_record.abort_event.set()
if False:
yield # pragma: no cover
def factory(*, config):
del config
return _AbortingAgent()
old_task = asyncio.create_task(
run_agent(
bridge,
run_manager,
old_record,
ctx=RunContext(checkpointer=checkpointer),
agent_factory=factory,
graph_input={"messages": [{"role": "user", "content": "Old prompt"}]},
config={},
)
)
old_record.task = old_task
try:
await asyncio.wait_for(old_title_gate_entered.wait(), timeout=1.0)
replacement_record = await run_manager.create_or_reject("thread-1", multitask_strategy="interrupt")
assert replacement_record.status == RunStatus.pending
release_old_title_gate.set()
await old_task
finally:
release_old_title_gate.set()
if not old_task.done():
old_task.cancel()
with suppress(asyncio.CancelledError):
await old_task
checkpointer.aput.assert_awaited_once()
_, written_checkpoint, _, _ = checkpointer.aput.await_args.args
assert written_checkpoint["channel_values"]["title"] == "Old Prompt"
@pytest.mark.anyio
async def test_interrupted_title_does_not_overwrite_checkpoint_from_admitted_replacement(monkeypatch):
"""A replacement run admitted by multitask interrupt must not lose its newer checkpoint."""
from deerflow.agents.middlewares.title_middleware import TitleMiddleware
monkeypatch.setattr(
TitleMiddleware,
"_generate_title_result",
lambda self, state, allow_partial_exchange=False: {"title": "Old prompt"},
)
old_checkpoint = {
"id": "ckpt-old",
"ts": "2026-06-29T00:00:00Z",
"channel_values": {"messages": [{"type": "human", "content": "Old prompt"}]},
"channel_versions": {"messages": 1},
}
replacement_messages = [
{"type": "human", "content": "Old prompt"},
{"type": "human", "content": "Replacement prompt"},
]
replacement_checkpoint = {
"id": "ckpt-replacement",
"ts": "2026-06-29T00:00:01Z",
"channel_values": {"messages": replacement_messages},
"channel_versions": {"messages": 2},
}
class _ReplacementRaceCheckpointer:
def __init__(self) -> None:
self.latest_checkpoint = copy.deepcopy(old_checkpoint)
self.latest_metadata = {"source": "loop", "step": 1}
self.title_write_started = asyncio.Event()
self.replacement_checkpoint_written = asyncio.Event()
async def aget_tuple(self, config):
del config
return _FakeCheckpointTuple(
checkpoint=copy.deepcopy(self.latest_checkpoint),
metadata=dict(self.latest_metadata),
config={
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
"checkpoint_id": self.latest_checkpoint["id"],
}
},
)
async def aput(self, config, checkpoint, metadata, new_versions):
del config, new_versions
self.title_write_started.set()
await self.replacement_checkpoint_written.wait()
self.latest_checkpoint = copy.deepcopy(checkpoint)
self.latest_metadata = dict(metadata)
return {
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
"checkpoint_id": checkpoint["id"],
}
}
def get_next_version(self, current, _channel):
return (current or 0) + 1
run_manager = RunManager()
old_record = await run_manager.create("thread-1")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
checkpointer = _ReplacementRaceCheckpointer()
old_agent_started = asyncio.Event()
class _BlockingAgent:
metadata = {"model_name": "fake-test-model"}
checkpointer: Any | None = None
store: Any | None = None
interrupt_before_nodes = None
interrupt_after_nodes = None
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
del graph_input, config, stream_mode, subgraphs
old_agent_started.set()
while True:
await asyncio.sleep(0.05)
if False:
yield # pragma: no cover
def factory(*, config):
del config
return _BlockingAgent()
old_task = asyncio.create_task(
run_agent(
bridge,
run_manager,
old_record,
ctx=RunContext(checkpointer=checkpointer),
agent_factory=factory,
graph_input={"messages": [{"role": "user", "content": "Old prompt"}]},
config={},
)
)
old_record.task = old_task
try:
await asyncio.wait_for(old_agent_started.wait(), timeout=1.0)
replacement_record = await run_manager.create_or_reject("thread-1", multitask_strategy="interrupt")
assert replacement_record.run_id != old_record.run_id
await run_manager.set_status(replacement_record.run_id, RunStatus.running)
with suppress(asyncio.TimeoutError):
await asyncio.wait_for(checkpointer.title_write_started.wait(), timeout=0.25)
checkpointer.latest_checkpoint = copy.deepcopy(replacement_checkpoint)
checkpointer.latest_metadata = {"source": "loop", "step": 2}
checkpointer.replacement_checkpoint_written.set()
await old_task
finally:
checkpointer.replacement_checkpoint_written.set()
if not old_task.done():
old_task.cancel()
with suppress(asyncio.CancelledError):
await old_task
assert checkpointer.latest_checkpoint["channel_values"]["messages"] == replacement_messages
@pytest.mark.anyio
async def test_replacement_run_waits_for_prior_finalizing_run():
"""Replacement workers must not enter the graph while an older run is finalizing."""
run_manager = RunManager()
old_record = await run_manager.create("thread-1")
replacement_record = await run_manager.create("thread-1")
await run_manager.set_finalizing(old_record.run_id, True)
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
replacement_started = asyncio.Event()
class _ReplacementAgent:
metadata = {"model_name": "fake-test-model"}
checkpointer: Any | None = None
store: Any | None = None
interrupt_before_nodes = None
interrupt_after_nodes = None
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
del graph_input, config, stream_mode, subgraphs
replacement_started.set()
if False:
yield # pragma: no cover
def factory(*, config):
del config
return _ReplacementAgent()
task = asyncio.create_task(
run_agent(
bridge,
run_manager,
replacement_record,
ctx=RunContext(checkpointer=None),
agent_factory=factory,
graph_input={"messages": [{"role": "user", "content": "Replacement prompt"}]},
config={},
)
)
replacement_record.task = task
try:
await asyncio.sleep(0.1)
assert not replacement_started.is_set()
await run_manager.set_finalizing(old_record.run_id, False)
await asyncio.wait_for(replacement_started.wait(), timeout=1.0)
await task
finally:
await run_manager.set_finalizing(old_record.run_id, False)
if not task.done():
task.cancel()
with suppress(asyncio.CancelledError):
await task
@pytest.mark.anyio
async def test_ensure_interrupted_title_reloads_latest_checkpoint_before_write():
"""If the checkpoint advances before the title write, preserve the newer messages."""
from deerflow.config.title_config import TitleConfig
old_checkpoint = {
"id": "ckpt-old",
"ts": "2026-06-29T00:00:00Z",
"channel_values": {"messages": [{"type": "human", "content": "Old prompt"}]},
"channel_versions": {"messages": 1},
}
new_messages = [{"type": "human", "content": "New prompt"}]
new_checkpoint = {
"id": "ckpt-new",
"ts": "2026-06-29T00:00:01Z",
"channel_values": {"messages": new_messages},
"channel_versions": {"messages": 2},
}
class _AdvancingTitleCheckpointer:
def __init__(self) -> None:
self.read_count = 0
self.aput = AsyncMock(return_value={})
async def aget_tuple(self, config):
del config
self.read_count += 1
checkpoint = old_checkpoint if self.read_count == 1 else new_checkpoint
return _FakeCheckpointTuple(
checkpoint=copy.deepcopy(checkpoint),
metadata={"source": "loop", "step": self.read_count},
config={
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
"checkpoint_id": checkpoint["id"],
}
},
)
def get_next_version(self, current, _channel):
return (current or 0) + 1
checkpointer = _AdvancingTitleCheckpointer()
app_config = SimpleNamespace(title=TitleConfig(enabled=True, max_chars=40, max_words=20))
title = await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=app_config)
assert title == "New prompt"
_, written_checkpoint, _, _ = checkpointer.aput.await_args.args
assert written_checkpoint["channel_values"]["messages"] == new_messages
assert written_checkpoint["channel_values"]["title"] == "New prompt"
@pytest.mark.anyio
async def test_ensure_interrupted_title_bumps_channel_version_and_declares_it_in_new_versions(monkeypatch):
"""Regression for #3859 review: DB-backed savers (Sqlite/Postgres) strip inline
``channel_values`` from ``put`` and only persist blobs for channels listed in
``new_versions``. The helper must therefore bump ``channel_versions["title"]``
and pass ``{"title": next_version}`` so the fallback title actually survives
a fresh ``aget_tuple`` after the worker's finally hook.
"""
from deerflow.agents.middlewares.title_middleware import TitleMiddleware
monkeypatch.setattr(
TitleMiddleware,
"_generate_title_result",
lambda self, state, allow_partial_exchange=False: {"title": "Generated Title"},
)
initial_checkpoint = {
"id": "ckpt-1",
"ts": "2026-06-29T00:00:00Z",
"channel_values": {"messages": [{"type": "human", "content": "hi"}]},
"channel_versions": {"messages": 5},
}
checkpointer = _TitleCheckpointer(
tuple_value=_FakeCheckpointTuple(
checkpoint=initial_checkpoint,
metadata={"source": "loop", "step": 7},
),
)
title = await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None)
assert title == "Generated Title"
checkpointer.aput.assert_awaited_once()
write_config, written_checkpoint, written_metadata, new_versions = checkpointer.aput.await_args.args
# The title channel must be declared in new_versions — without this, DB
# savers drop the inline channel_values["title"] from the persisted blob.
assert new_versions == {"title": 1}
# Channel versions on the checkpoint itself must also reflect the bump,
# so a subsequent aget_tuple reconstructs channel_values with the title.
assert written_checkpoint["channel_versions"]["title"] == 1
# Pre-existing channel versions must be preserved.
assert written_checkpoint["channel_versions"]["messages"] == 5
# The fallback title rides into channel_values for the (legacy / single-table)
# savers that inline the snapshot.
assert written_checkpoint["channel_values"]["title"] == "Generated Title"
assert written_metadata["source"] == "update"
assert written_metadata["step"] == 8
assert written_metadata["writes"] == {"runtime_interrupt_title": {"title": "Generated Title"}}
# The write is parented to the checkpoint it was derived from - without
# the parent pointer the saver stores a root checkpoint, severing
# Delta-channel replay ancestry and full-mode history walks.
assert write_config == {"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "ckpt-1"}}
@pytest.mark.anyio
async def test_ensure_interrupted_title_writes_graph_input_fallback_without_checkpoint(monkeypatch):
"""When no checkpoint exists, graph_input should still seed the fallback title write."""
from deerflow.agents.middlewares.title_middleware import TitleMiddleware
captured_state: dict[str, Any] = {}
def _generate(self, state, allow_partial_exchange=False):
del self
captured_state.update(state)
assert allow_partial_exchange is True
return {"title": "Graph Input Title"}
monkeypatch.setattr(TitleMiddleware, "_generate_title_result", _generate)
checkpointer = _TitleCheckpointer(tuple_value=None)
title = await _ensure_interrupted_title(
checkpointer=checkpointer,
thread_id="thread-1",
app_config=None,
graph_input={"messages": [{"role": "user", "content": "Please name this thread"}]},
)
assert title == "Graph Input Title"
assert captured_state["messages"] == [{"role": "user", "content": "Please name this thread"}]
checkpointer.aput.assert_awaited_once()
_, written_checkpoint, written_metadata, new_versions = checkpointer.aput.await_args.args
assert written_checkpoint["channel_values"]["title"] == "Graph Input Title"
assert written_metadata["writes"] == {"runtime_interrupt_title": {"title": "Graph Input Title"}}
assert set(new_versions.keys()) == {"title"}
@pytest.mark.anyio
async def test_ensure_interrupted_title_bumps_existing_string_version(monkeypatch):
"""When the checkpointer lacks ``get_next_version`` and the prior title
version is a string (some savers use UUID-shaped versions), the helper must
still produce a strictly different value rather than overwriting in place.
"""
from deerflow.agents.middlewares.title_middleware import TitleMiddleware
monkeypatch.setattr(
TitleMiddleware,
"_generate_title_result",
lambda self, state, allow_partial_exchange=False: {"title": "T"},
)
initial_checkpoint = {
"id": "ckpt-1",
"ts": "2026-06-29T00:00:00Z",
"channel_values": {"messages": [{"type": "human", "content": "hi"}]},
"channel_versions": {"title": "v3"},
}
class _NoGetNextVersion:
def __init__(self):
self.aget_tuple = AsyncMock(
return_value=_FakeCheckpointTuple(
checkpoint=initial_checkpoint,
metadata={},
),
)
self.aput = AsyncMock(return_value={})
checkpointer = _NoGetNextVersion()
await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None)
_, written_checkpoint, _, new_versions = checkpointer.aput.await_args.args
bumped = written_checkpoint["channel_versions"]["title"]
assert bumped != "v3", "title version must change so DB savers persist the update"
assert new_versions == {"title": bumped}
@pytest.mark.anyio
async def test_ensure_interrupted_title_skips_when_title_already_set():
"""If the checkpoint already carries a title, no new checkpoint is written."""
checkpointer = _TitleCheckpointer(
tuple_value=_FakeCheckpointTuple(
checkpoint={
"id": "ckpt-1",
"channel_values": {"messages": [], "title": "Already there"},
"channel_versions": {"title": 1},
},
metadata={},
),
)
title = await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None)
assert title == "Already there"
checkpointer.aput.assert_not_awaited()
@pytest.mark.anyio
async def test_ensure_interrupted_title_round_trip_with_real_sqlite_checkpointer(tmp_path):
"""Full round-trip against a real ``AsyncSqliteSaver`` on a disk-backed DB.
Mirrors what Gateway constructs in production via ``make_checkpointer`` when
``database.backend == "sqlite"``, then closes and re-opens the saver to
simulate a fresh connection. The fallback title must survive that boundary —
this is the scenario the #3874 review flagged as broken before the
``new_versions={"title": ...}`` fix.
"""
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.base import empty_checkpoint
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from deerflow.config.title_config import TitleConfig
db_path = str(tmp_path / "ckpt.db")
thread_cfg = {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}}
# 1. Seed a first-turn checkpoint that has a human message and NO title —
# the same shape the agent leaves behind when interrupted mid-stream.
async with AsyncSqliteSaver.from_conn_string(db_path) as writer:
await writer.setup()
ck = empty_checkpoint()
ck["channel_values"] = {
"messages": [HumanMessage(content="Why is the sky blue?").model_dump()],
}
ck["channel_versions"] = {"messages": 1}
await writer.aput(thread_cfg, ck, {"source": "loop", "step": 1, "writes": {}}, {"messages": 1})
# 2. Run the worker helper through a *fresh* saver instance — this is what
# the lifespan-owned checkpointer pool does for each request.
title_config = TitleConfig(enabled=True, max_chars=40, max_words=20)
app_config = SimpleNamespace(title=title_config)
async with AsyncSqliteSaver.from_conn_string(db_path) as worker_saver:
title = await _ensure_interrupted_title(
checkpointer=worker_saver,
thread_id="thread-1",
app_config=app_config,
)
assert title, "fallback title must be generated from the seeded user message"
# 3. Open ANOTHER fresh saver and confirm the title survives — this is the
# invariant the #3874 review was guarding: ``new_versions={}`` would
# cause DB savers to drop the title blob, so a fresh aget_tuple would
# read back without it.
async with AsyncSqliteSaver.from_conn_string(db_path) as reader:
tup = await reader.aget_tuple(thread_cfg)
assert tup is not None
persisted = tup.checkpoint.get("channel_values", {}).get("title")
assert persisted == title
# ---------------------------------------------------------------------------
# _bump_channel_version — invariant: the returned version MUST differ from
# the prior value, no matter the checkpointer's versioning scheme.
# ---------------------------------------------------------------------------
class _CheckpointerWithIntVersion:
"""A checkpointer whose ``get_next_version`` increments integers (default LangGraph behavior)."""
@staticmethod
def get_next_version(current, _channel):
return (current or 0) + 1
class _CheckpointerWithBrokenGetNextVersion:
"""A checkpointer whose ``get_next_version`` raises — must fall back, not propagate."""
@staticmethod
def get_next_version(current, _channel):
raise RuntimeError("simulated saver bug")
def test_bump_channel_version_uses_checkpointer_get_next_version_when_available():
"""Happy path — saver's ``get_next_version`` result is preferred over our fallback."""
assert _bump_channel_version(_CheckpointerWithIntVersion(), 5) == 6
def test_bump_channel_version_falls_back_on_broken_get_next_version():
"""A raising ``get_next_version`` must not propagate; the defensive path bumps from prior.
Without this, a saver bug would leave ``new_versions={"title": v}`` no-op
on DB savers — the very class of bug the #3874 review flagged.
"""
bumped = _bump_channel_version(_CheckpointerWithBrokenGetNextVersion(), 7)
assert bumped == 8
# ---------------------------------------------------------------------------
# _ensure_interrupted_title — additional defensive boundaries
# ---------------------------------------------------------------------------
@pytest.mark.anyio
async def test_ensure_interrupted_title_handles_none_messages_channel(monkeypatch):
"""A partially-initialized checkpoint with ``messages=None`` must not crash."""
from deerflow.config.title_config import TitleConfig
initial_checkpoint = {
"id": "ckpt-1",
"channel_values": {"messages": None},
"channel_versions": {},
}
checkpointer = _TitleCheckpointer(
tuple_value=_FakeCheckpointTuple(checkpoint=initial_checkpoint, metadata={}),
)
app_config = SimpleNamespace(title=TitleConfig(enabled=True, max_chars=40, max_words=20))
assert await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=app_config) is None
checkpointer.aput.assert_not_awaited()
@pytest.mark.anyio
async def test_ensure_interrupted_title_propagates_aput_error_to_caller(monkeypatch):
"""Exceptions from ``aput`` propagate — the caller (worker.run_agent finally block) is responsible for swallowing them.
This test pins the contract: the helper itself does NOT silently eat saver errors,
so a structural saver regression remains visible in the logs at the call site.
"""
from deerflow.agents.middlewares.title_middleware import TitleMiddleware
monkeypatch.setattr(
TitleMiddleware,
"_generate_title_result",
lambda self, state, allow_partial_exchange=False: {"title": "Generated"},
)
initial_checkpoint = {
"id": "ckpt-1",
"channel_values": {"messages": [{"type": "human", "content": "hi"}]},
"channel_versions": {"messages": 1},
}
checkpointer = _TitleCheckpointer(
tuple_value=_FakeCheckpointTuple(checkpoint=initial_checkpoint, metadata={}),
)
checkpointer.aput.side_effect = RuntimeError("simulated DB write failure")
with pytest.raises(RuntimeError, match="simulated DB write failure"):
await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None)
@pytest.mark.anyio
async def test_ensure_interrupted_title_idempotent_across_repeated_calls(monkeypatch):
"""Second invocation against the now-titled checkpoint must not re-write.
Regression anchor for the case where a brittle helper might re-trigger
on subsequent finally-hook runs (e.g. retries) and rewrite the title.
"""
from langgraph.checkpoint.memory import InMemorySaver
from deerflow.agents.middlewares.title_middleware import TitleMiddleware
monkeypatch.setattr(
TitleMiddleware,
"_generate_title_result",
lambda self, state, allow_partial_exchange=False: {"title": "First Title"},
)
checkpointer = InMemorySaver()
cfg = {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}}
ck = empty_checkpoint()
ck["channel_values"] = {"messages": [{"type": "human", "content": "hi"}]}
ck["channel_versions"] = {"messages": 1}
await checkpointer.aput(cfg, ck, {"source": "loop", "step": 1, "writes": {}}, {"messages": 1})
first = await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None)
assert first == "First Title"
# Second call: title is now present, so the helper short-circuits without
# rewriting — even if the middleware were to suggest a different title.
monkeypatch.setattr(
TitleMiddleware,
"_generate_title_result",
lambda self, state, allow_partial_exchange=False: {"title": "Different Title"},
)
second = await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None)
assert second == "First Title"
tup = await checkpointer.aget_tuple(cfg)
assert tup.checkpoint["channel_values"]["title"] == "First Title"
@pytest.mark.anyio
async def test_ensure_interrupted_title_preserves_non_title_channel_versions(monkeypatch):
"""Bumping ``channel_versions["title"]`` must not modify other channels' versions.
Regression anchor: an earlier draft built ``new_versions`` from
``dict(channel_versions)`` and would have erroneously declared every
channel as "needs new blob" on DB savers.
"""
from deerflow.agents.middlewares.title_middleware import TitleMiddleware
monkeypatch.setattr(
TitleMiddleware,
"_generate_title_result",
lambda self, state, allow_partial_exchange=False: {"title": "Generated"},
)
initial_checkpoint = {
"id": "ckpt-1",
"channel_values": {
"messages": [{"type": "human", "content": "hi"}],
"artifacts": [],
"todos": None,
},
"channel_versions": {"messages": 5, "artifacts": 3, "todos": 1},
}
checkpointer = _TitleCheckpointer(
tuple_value=_FakeCheckpointTuple(checkpoint=initial_checkpoint, metadata={}),
)
await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None)
_, written_checkpoint, _, new_versions = checkpointer.aput.await_args.args
# Only the title channel is declared in new_versions.
assert set(new_versions.keys()) == {"title"}
# Other channel versions are preserved verbatim on the written checkpoint.
assert written_checkpoint["channel_versions"]["messages"] == 5
assert written_checkpoint["channel_versions"]["artifacts"] == 3
assert written_checkpoint["channel_versions"]["todos"] == 1
@pytest.mark.anyio
async def test_worker_finally_block_swallows_helper_exceptions(monkeypatch):
"""The worker's interrupted-title hook must remain non-fatal — any exception
from the helper (DB saver bug, middleware bug, etc.) must not propagate past
the run boundary or prevent the subsequent threads_meta sync block from
running. This pins the integration of helper + finally try/except, not just
the helper itself.
"""
import deerflow.runtime.runs.worker as worker_module
helper_called = asyncio.Event()
async def _boom(*_args, **_kwargs):
helper_called.set()
raise RuntimeError("forced helper failure")
monkeypatch.setattr(worker_module, "_ensure_interrupted_title", _boom)
run_manager = RunManager()
record = await run_manager.create("thread-1")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
class _MinimalCheckpointer:
async def aget_tuple(self, config):
return None
async def aput(self, *args, **kwargs):
return {}
captured_status: dict[str, Any] = {}
class _ThreadStore:
async def update_display_name(self, thread_id, title):
captured_status["display_name"] = (thread_id, title)
async def update_status(self, thread_id, status):
captured_status["status"] = (thread_id, status)
class _AbortingAgent:
def __init__(self) -> None:
self.metadata = {"model_name": "fake-test-model"}
self.checkpointer: Any | None = None
self.store: Any | None = None
self.interrupt_before_nodes = None
self.interrupt_after_nodes = None
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
# Abort immediately so the run lands in the interrupted branch.
record.abort_event.set()
if False:
yield # pragma: no cover — make this an async generator
return
def factory(*, config):
del config
return _AbortingAgent()
await run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=_MinimalCheckpointer(), thread_store=_ThreadStore()),
agent_factory=factory,
graph_input={"messages": []},
config={},
)
# The helper raised, but the run still reaches the threads_meta status sync
# and ``publish_end`` — i.e. the SSE stream is closed cleanly and the row
# reflects the run outcome.
assert helper_called.is_set()
assert captured_status.get("status") == ("thread-1", "interrupted")
bridge.publish_end.assert_awaited_once_with(record.run_id)
Reference in New Issue View Git Blame Copy Permalink