mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-10 23:08:45 +00:00
2e15e3fe0d
* fix: generate title for interrupted first turn * test(title): cover partial-exchange + dict-form messages Harden the interrupted-run fallback path added in 19fc34fd: - TitleMiddleware._should_generate_title now accepts a lone first-turn user message when allow_partial_exchange=True, so the worker can still derive a title if cancellation lands before any AI chunk is checkpointed. - runtime/runs/worker._ensure_interrupted_title computes the next checkpoint step defensively (treat missing/non-int step as 0) and renames a shadowed ckpt_config local for readability. - Add four unit tests in tests/test_title_middleware_core_logic.py: partial-exchange allows user-only, partial-exchange still respects an existing title, dict-form messages are recognized, and the sync fallback path derives a title from dict-form messages — matching what channel_values stores in the checkpoint. Refs #3859. * fix: persist interrupted-title via channel_versions bump Address PR #3874 review feedback: ``_ensure_interrupted_title`` previously called ``aput(..., new_versions={})``. LangGraph's DB-backed savers (``PostgresSaver`` and the v4 ``SqliteSaver`` blob layout) strip inline ``channel_values`` from ``put`` and only persist blobs for channels named in ``new_versions`` — so the fallback ``title`` channel was dropped on read-back and ``threads_meta.display_name`` stayed ``"Untitled"`` after refresh on those backends. The original in-memory e2e passed because ``InMemorySaver`` keeps the inline snapshot verbatim. Fix mirrors ``_rollback_to_pre_run_checkpoint`` in the same file: bump ``channel_versions["title"]`` (via ``checkpointer.get_next_version`` when available, else int/string fallbacks), persist the new version on the checkpoint, and declare it in ``new_versions`` so the DB savers actually write the blob. Regression coverage in ``tests/test_run_worker_rollback.py``: - ``test_ensure_interrupted_title_bumps_channel_version_and_declares_it_in_new_versions`` — exact ``aput`` invariants: ``new_versions == {"title": 1}``, the written checkpoint's ``channel_versions["title"]`` is bumped, and the pre-existing ``messages`` version is preserved. - ``test_ensure_interrupted_title_bumps_existing_string_version`` — string-shaped prior version (some savers use UUID-style versions); bumped value must differ from the prior, no overwrite-in-place. - ``test_ensure_interrupted_title_skips_when_title_already_set`` — title short-circuit; no extra ``aput``. - ``test_ensure_interrupted_title_returns_none_when_no_checkpoint`` — no checkpoint yet; returns ``None`` without writing. - ``test_ensure_interrupted_title_round_trip_with_real_sqlite_checkpointer`` — full round-trip against a real ``AsyncSqliteSaver`` on a disk-backed DB, then closes and re-opens the saver to simulate a fresh connection. The fallback title must still be present on the second ``aget_tuple``. This is the exact scenario the review flagged. Validated locally with the full backend suite: 5195 passed, 18 skipped. Refs #3859. Addresses review on #3874. * test(worker, title): harden interrupted-title fallback for every saver Defensive coverage on top of the channel_versions fix (commit 05253957), addressing edge cases surfaced during a second-pass review of #3874. Worker: - Extract version bump into ``_bump_channel_version(checkpointer, current)`` with explicit fallbacks for int / float / numeric-string / UUID-shaped string / None / bool, AND a wrap-around defense when the saver's ``get_next_version`` raises or returns an unchanged value. The invariant is: returned version MUST differ from the prior. Without this, a saver bug (or a custom backend) could leave ``new_versions={"title": v}`` no-op on DB savers — the very class of bug the original review pointed out. Title middleware: - Coerce ``state.get("messages")`` from ``None`` to ``[]`` on both ``_should_generate_title`` and ``_build_title_prompt``. A partially-initialized checkpoint can carry ``messages=None`` on the channel_values channel (the worker reads raw channel_values, not BaseMessages), and the default kwarg only protects against a missing key. Repro: ``TypeError: 'NoneType' object is not iterable`` from the next() generator — confirmed by reverting the fix and watching ``test_*_handles_none_messages_channel`` go red. Tests (TDD-verified red→green for the new asserts): - ``test_run_worker_rollback.py``: * ``_bump_channel_version`` — 8 tests covering every version type (int, float, numeric string, UUID-style string, None, bool) and every saver-side fault mode (no ``get_next_version`` / raising / stuck on identity). * ``test_ensure_interrupted_title_*`` — 5 additional helper boundary tests: title.enabled=false short-circuit; empty messages list; messages=None; aput-error propagation (helper contract: caller swallows, not the helper); idempotency on a real InMemorySaver across two invocations. * ``test_ensure_interrupted_title_preserves_non_title_channel_versions`` — pins that ``new_versions`` only contains ``"title"`` and that other channels' versions are untouched (regression anchor for a sloppier draft that bumped every channel). * ``test_worker_finally_block_swallows_helper_exceptions`` — pins the integration contract: even if the helper raises, the worker's threads_meta status sync still runs and ``publish_end`` is still awaited so the SSE stream closes cleanly. - ``test_title_middleware_core_logic.py``: * 4 additional tests: ``messages=None`` on both ``_should_generate_title`` and ``_build_title_prompt``; the ``role: user`` / ``role: assistant`` (OpenAI-style) dict normalization; partial-exchange path with a dict-form message. Verification: - ``PYTHONPATH=. uv run pytest tests/ -x --ignore=tests/blocking_io -q`` → 5215 passed, 18 skipped. - ``ruff check`` + ``ruff format --check`` clean on every touched file. - Red/green TDD verification: temporarily reverted the ``new_versions={}`` fix → 4 new tests went red as expected; restored and the suite is green again. Same red/green dance for the ``messages=None`` coercion. Refs #3859. Addresses second-pass review on #3874. * fix(title): ignore dict context reminders in fallback * fix(worker): link interrupted-title checkpoint to its parent The title-bump checkpoint written by ``_ensure_interrupted_title`` was landing without a ``parent_checkpoint_id`` — a real orphan in the LangGraph history graph. Reproduction (disk-backed AsyncSqliteSaver): [seed] checkpoint_id = 1f173dbc... [helper] wrote title = "Why is the sky blue?" [issue 1] new checkpoint = 1f173dbc..., parent = None [issue 1] is new checkpoint orphaned? True Root cause: ``_ensure_interrupted_title`` built ``write_config`` as ``{"thread_id": ..., "checkpoint_ns": ...}`` only. ``BaseCheckpointSaver`` implementations read ``configurable.checkpoint_id`` from that config as the *parent* id when inserting (see ``langgraph/checkpoint/sqlite/aio.py`` ``aput``: ``config["configurable"].get("checkpoint_id")`` becomes the ``parent_checkpoint_id`` column). With no value, the saver writes NULL — the new checkpoint is a tree root. Consequences: - Any future LangGraph ``runs.resume_from`` / time-travel feature has no backward edge to walk past the title-bump. - History-visualization UIs built on ``alist()`` render the title-bump as a sibling of the prior checkpoint, not its descendant. Fix: read ``checkpoint_id`` off the tuple's own config and thread it into ``write_config["configurable"]["checkpoint_id"]`` before calling ``aput``, the same pattern every middleware-driven write uses. Three new regression tests against real ``AsyncSqliteSaver`` (disk-backed, fresh connections so we exercise the on-disk read path): - ``test_ensure_interrupted_title_links_new_checkpoint_to_its_parent`` — asserts ``latest.parent_config["configurable"]["checkpoint_id"]`` equals the seeded checkpoint id. TDD red-green verified: reverting the fix flips this test red with ``AssertionError: title-bump checkpoint must have a parent_config``. - ``test_ensure_interrupted_title_appears_in_history_with_audit_marker`` — pins the audit contract: the title-bump entry in ``alist()`` carries ``metadata.source == "update"`` and ``metadata.writes`` contains ``runtime_interrupt_title``. This is a deliberate design choice — we do NOT hide the entry from history (audit trail belongs in the saver), but its source and writes marker MUST be unambiguous so UIs/tools can identify it. - ``test_ensure_interrupted_title_survives_immediate_next_turn`` — cancel → immediate user follow-up scenario. Simulates the agent's next turn appending a (user, ai) pair without touching the title channel, then opens a fresh saver and verifies the title is still present after the next-turn checkpoint write. Pins the channel-version-blob invariant established by commit 05253957 — without the ``new_versions={"title": v}`` declaration there, the title blob would vanish from the DB and this test would read back ``None``. Verification: - ``PYTHONPATH=. uv run pytest tests/ -x --ignore=tests/blocking_io -q`` → 5222 passed, 15 skipped. - ``ruff check`` + ``ruff format --check`` clean on every touched file. - Reproduction script confirms ``parent_checkpoint_id`` is now non-null and the next-turn read-back preserves the fallback title. Refs #3859. * Revert "fix(worker): link interrupted-title checkpoint to its parent" This reverts commit c763ed9334781db1acdce0f5f33d663d8d5f80ff. * test: trim over-engineered test coverage Reduce review surface area on PR #3874 by dropping defensive tests that don't pin a real invariant. After self-review: - ``_bump_channel_version``: 8 tests → 2 (happy path + saver-error fallback). Dropped float / bool / numeric-string / UUID-string / missing-get-next-version / stuck-get-next-version branches — those are speculative scaffolding for savers we don't ship. - ``_ensure_interrupted_title``: dropped ``returns_none_when_title_disabled``, ``returns_none_with_no_user_message``, ``returns_none_when_no_checkpoint`` — boundary guards already exercised by the e2e test and the ``handles_none_messages_channel`` regression anchor. Net: -107 lines of test code. Remaining coverage still pins every red-green-verified invariant (channel_versions bump, string-version bump, idempotency, sqlite round-trip, non-title channel preservation, aput-error contract, worker finally swallowing, partial-exchange). Verification: 5209 passed, 15 skipped. * fix: harden interrupted title finalization * fix: serialize interrupted title finalization * fix: preserve interrupt semantics during title finalization * fix: preserve delayed interrupted title recovery --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
1657 lines
61 KiB
Python
1657 lines
61 KiB
Python
import asyncio
|
|
import copy
|
|
from contextlib import suppress
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock, call
|
|
|
|
import pytest
|
|
from langchain_core.messages import AIMessage
|
|
from langgraph.checkpoint.base import empty_checkpoint
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
|
|
from deerflow.runtime.runs.manager import ConflictError, RunManager
|
|
from deerflow.runtime.runs.schemas import RunStatus
|
|
from deerflow.runtime.runs.worker import (
|
|
RunContext,
|
|
_agent_factory_supports_app_config,
|
|
_build_runtime_context,
|
|
_bump_channel_version,
|
|
_ensure_interrupted_title,
|
|
_extract_llm_error_fallback_message,
|
|
_install_runtime_context,
|
|
_rollback_to_pre_run_checkpoint,
|
|
_try_extract_from_message,
|
|
run_agent,
|
|
)
|
|
|
|
|
|
class FakeCheckpointer:
|
|
def __init__(self, *, put_result):
|
|
self.adelete_thread = AsyncMock()
|
|
self.aput = AsyncMock(return_value=put_result)
|
|
self.aput_writes = AsyncMock()
|
|
|
|
|
|
def _make_checkpoint(checkpoint_id: str, messages: list[str], version: int):
|
|
checkpoint = empty_checkpoint()
|
|
checkpoint["id"] = checkpoint_id
|
|
checkpoint["channel_values"] = {"messages": messages}
|
|
checkpoint["channel_versions"] = {"messages": version}
|
|
return checkpoint
|
|
|
|
|
|
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
|
|
|
|
|
|
@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_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_restores_snapshot_without_deleting_thread():
|
|
checkpointer = FakeCheckpointer(put_result={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}})
|
|
|
|
await _rollback_to_pre_run_checkpoint(
|
|
checkpointer=checkpointer,
|
|
thread_id="thread-1",
|
|
run_id="run-1",
|
|
pre_run_checkpoint_id="ckpt-1",
|
|
pre_run_snapshot={
|
|
"checkpoint_ns": "",
|
|
"checkpoint": {
|
|
"id": "ckpt-1",
|
|
"channel_versions": {"messages": 3},
|
|
"channel_values": {"messages": ["before"]},
|
|
},
|
|
"metadata": {"source": "input"},
|
|
"pending_writes": [
|
|
("task-a", "messages", {"content": "first"}),
|
|
("task-a", "status", "done"),
|
|
("task-b", "events", {"type": "tool"}),
|
|
],
|
|
},
|
|
snapshot_capture_failed=False,
|
|
)
|
|
|
|
checkpointer.adelete_thread.assert_not_awaited()
|
|
checkpointer.aput.assert_awaited_once()
|
|
restore_config, restored_checkpoint, restored_metadata, new_versions = checkpointer.aput.await_args.args
|
|
assert restore_config == {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}}
|
|
assert restored_checkpoint["id"] != "ckpt-1"
|
|
assert "channel_versions" in restored_checkpoint
|
|
assert "channel_values" in restored_checkpoint
|
|
assert restored_checkpoint["channel_versions"] == {"messages": 3}
|
|
assert restored_checkpoint["channel_values"] == {"messages": ["before"]}
|
|
assert restored_metadata == {"source": "input"}
|
|
assert new_versions == {"messages": 3}
|
|
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_restored_checkpoint_becomes_latest_with_real_checkpointer():
|
|
checkpointer = InMemorySaver()
|
|
thread_config = {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}}
|
|
before_checkpoint = _make_checkpoint("0001", ["before"], 1)
|
|
before_config = checkpointer.put(thread_config, before_checkpoint, {"step": 1}, {"messages": 1})
|
|
after_checkpoint = _make_checkpoint("0002", ["after"], 2)
|
|
after_config = checkpointer.put(before_config, after_checkpoint, {"step": 2}, {"messages": 2})
|
|
checkpointer.put_writes(after_config, [("messages", "pending-after")], task_id="task-after")
|
|
|
|
await _rollback_to_pre_run_checkpoint(
|
|
checkpointer=checkpointer,
|
|
thread_id="thread-1",
|
|
run_id="run-1",
|
|
pre_run_checkpoint_id="0001",
|
|
pre_run_snapshot={
|
|
"checkpoint_ns": "",
|
|
"checkpoint": before_checkpoint,
|
|
"metadata": {"step": 1},
|
|
"pending_writes": [("task-before", "messages", "pending-before")],
|
|
},
|
|
snapshot_capture_failed=False,
|
|
)
|
|
|
|
latest = checkpointer.get_tuple(thread_config)
|
|
|
|
assert latest is not None
|
|
assert latest.config["configurable"]["checkpoint_id"] != "0001"
|
|
assert latest.config["configurable"]["checkpoint_id"] != "0002"
|
|
assert latest.checkpoint["channel_values"] == {"messages": ["before"]}
|
|
assert latest.pending_writes == [("task-before", "messages", "pending-before")]
|
|
assert ("task-after", "messages", "pending-after") not in latest.pending_writes
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_rollback_deletes_thread_when_no_snapshot_exists():
|
|
checkpointer = FakeCheckpointer(put_result=None)
|
|
|
|
await _rollback_to_pre_run_checkpoint(
|
|
checkpointer=checkpointer,
|
|
thread_id="thread-1",
|
|
run_id="run-1",
|
|
pre_run_checkpoint_id=None,
|
|
pre_run_snapshot=None,
|
|
snapshot_capture_failed=False,
|
|
)
|
|
|
|
checkpointer.adelete_thread.assert_awaited_once_with("thread-1")
|
|
checkpointer.aput.assert_not_awaited()
|
|
checkpointer.aput_writes.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_rollback_raises_when_restore_config_has_no_checkpoint_id():
|
|
checkpointer = FakeCheckpointer(put_result={"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}})
|
|
|
|
with pytest.raises(RuntimeError, match="did not return checkpoint_id"):
|
|
await _rollback_to_pre_run_checkpoint(
|
|
checkpointer=checkpointer,
|
|
thread_id="thread-1",
|
|
run_id="run-1",
|
|
pre_run_checkpoint_id="ckpt-1",
|
|
pre_run_snapshot={
|
|
"checkpoint_ns": "",
|
|
"checkpoint": {"id": "ckpt-1", "channel_versions": {}},
|
|
"metadata": {},
|
|
"pending_writes": [("task-a", "messages", "value")],
|
|
},
|
|
snapshot_capture_failed=False,
|
|
)
|
|
|
|
checkpointer.adelete_thread.assert_not_awaited()
|
|
checkpointer.aput.assert_awaited_once()
|
|
checkpointer.aput_writes.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_rollback_normalizes_none_checkpoint_ns_to_root_namespace():
|
|
checkpointer = FakeCheckpointer(put_result={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}})
|
|
|
|
await _rollback_to_pre_run_checkpoint(
|
|
checkpointer=checkpointer,
|
|
thread_id="thread-1",
|
|
run_id="run-1",
|
|
pre_run_checkpoint_id="ckpt-1",
|
|
pre_run_snapshot={
|
|
"checkpoint_ns": None,
|
|
"checkpoint": {"id": "ckpt-1", "channel_versions": {}},
|
|
"metadata": {},
|
|
"pending_writes": [],
|
|
},
|
|
snapshot_capture_failed=False,
|
|
)
|
|
|
|
checkpointer.aput.assert_awaited_once()
|
|
restore_config, restored_checkpoint, restored_metadata, new_versions = checkpointer.aput.await_args.args
|
|
assert restore_config == {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}}
|
|
assert restored_checkpoint["id"] != "ckpt-1"
|
|
assert restored_checkpoint["channel_versions"] == {}
|
|
assert restored_metadata == {}
|
|
assert new_versions == {}
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_rollback_raises_on_malformed_pending_write_not_a_tuple():
|
|
"""pending_writes containing a non-3-tuple item should raise RuntimeError."""
|
|
checkpointer = FakeCheckpointer(put_result={"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(
|
|
checkpointer=checkpointer,
|
|
thread_id="thread-1",
|
|
run_id="run-1",
|
|
pre_run_checkpoint_id="ckpt-1",
|
|
pre_run_snapshot={
|
|
"checkpoint_ns": "",
|
|
"checkpoint": {"id": "ckpt-1", "channel_versions": {}},
|
|
"metadata": {},
|
|
"pending_writes": [
|
|
("task-a", "messages", "valid"), # valid
|
|
["only", "two"], # malformed: only 2 elements
|
|
],
|
|
},
|
|
snapshot_capture_failed=False,
|
|
)
|
|
|
|
# aput succeeded but aput_writes should not be called due to malformed data
|
|
checkpointer.aput.assert_awaited_once()
|
|
checkpointer.aput_writes.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_rollback_raises_on_malformed_pending_write_non_string_channel():
|
|
"""pending_writes containing a non-string channel should raise RuntimeError."""
|
|
checkpointer = FakeCheckpointer(put_result={"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(
|
|
checkpointer=checkpointer,
|
|
thread_id="thread-1",
|
|
run_id="run-1",
|
|
pre_run_checkpoint_id="ckpt-1",
|
|
pre_run_snapshot={
|
|
"checkpoint_ns": "",
|
|
"checkpoint": {"id": "ckpt-1", "channel_versions": {}},
|
|
"metadata": {},
|
|
"pending_writes": [
|
|
("task-a", 123, "value"), # malformed: channel is not a string
|
|
],
|
|
},
|
|
snapshot_capture_failed=False,
|
|
)
|
|
|
|
checkpointer.aput.assert_awaited_once()
|
|
checkpointer.aput_writes.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_rollback_propagates_aput_writes_failure():
|
|
"""If aput_writes fails, the exception should propagate (not be swallowed)."""
|
|
checkpointer = FakeCheckpointer(put_result={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}})
|
|
# Simulate aput_writes failure
|
|
checkpointer.aput_writes.side_effect = RuntimeError("Database connection lost")
|
|
|
|
with pytest.raises(RuntimeError, match="Database connection lost"):
|
|
await _rollback_to_pre_run_checkpoint(
|
|
checkpointer=checkpointer,
|
|
thread_id="thread-1",
|
|
run_id="run-1",
|
|
pre_run_checkpoint_id="ckpt-1",
|
|
pre_run_snapshot={
|
|
"checkpoint_ns": "",
|
|
"checkpoint": {"id": "ckpt-1", "channel_versions": {}},
|
|
"metadata": {},
|
|
"pending_writes": [
|
|
("task-a", "messages", "value"),
|
|
],
|
|
},
|
|
snapshot_capture_failed=False,
|
|
)
|
|
|
|
# aput succeeded, aput_writes was called but failed
|
|
checkpointer.aput.assert_awaited_once()
|
|
checkpointer.aput_writes.assert_awaited_once()
|
|
|
|
|
|
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_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
|
|
|
|
|
|
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"}}
|
|
assert write_config == {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}}
|
|
|
|
|
|
@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
|
|
|
|
async def _boom(*_args, **_kwargs):
|
|
raise RuntimeError("forced helper failure")
|
|
|
|
monkeypatch.setattr(worker_module, "_ensure_interrupted_title", _boom)
|
|
|
|
run_manager = RunManager()
|
|
record = await run_manager.create("thread-1")
|
|
record.status = RunStatus.interrupted
|
|
|
|
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 captured_status.get("status") == ("thread-1", "interrupted")
|
|
bridge.publish_end.assert_awaited_once_with(record.run_id)
|