mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-10 14:58:46 +00:00
dcb2e687d5
* feat(channels): add GitHub event-driven agents (#3754) Add a webhook-driven GitHub channel with fail-closed webhook routing, deterministic per-agent PR/issue threads, mention-gated trigger fan-out, GitHub App token injection for sandboxed gh/git commands, and backend/AGENTS.md documentation. * fix(llm-middleware): classify bare IndexError as transient Upstream chat providers occasionally return 200 OK with an empty generations list (observed against Volces "coding" on ark.cn-beijing.volces.com). When that happens, langchain_core.language_models.chat_models.ainvoke raises ``IndexError: list index out of range`` at ``llm_result.generations[0][0].message`` and kills the run. Treat a bare IndexError reaching the middleware as a transient upstream-payload glitch and route it through the existing retry/backoff path instead of failing the whole agent run. The retry budget and backoff schedule are unchanged. Adds three regression tests covering the classifier and both the recover-on-retry and exhausted-retries paths. * fix(runtime): ignore stale LLM fallback markers from prior runs When a run on a thread ends with the LLM-error-handling middleware emitting a `deerflow_error_fallback`-marked AIMessage (e.g. after the IndexError empty-generations classification fix lands), that message is persisted to the thread's checkpoint as part of the messages channel. LangGraph replays the full message history in `stream_mode="values"` chunks, so every subsequent run on the same thread re-streams the stale fallback marker — and the worker's chunk scanner faithfully picks it up, flipping `RunStatus.success` to `RunStatus.error` for runs that themselves had no LLM failure at all. Snapshot the set of pre-existing message ids from the pre-run checkpoint and thread it through `_extract_llm_error_fallback_message` / `_try_extract_from_message` as a filter. Markers on history messages are ignored; markers on fresh messages produced during this run still trip the error path. Falls back to an empty set when the checkpointer is absent or the snapshot can't be captured, preserving the prior behavior on first-run / no-state paths. Adds unit tests for the new filter (helper-level and `_collect_pre_existing_message_ids`) plus an integration test exercising the full `run_agent` path with a stale history checkpointer. * fix(channels): make github channel fire-and-forget to avoid httpx.ReadTimeout on long runs GitHub agent runs (clone -> edit -> test -> push -> PR) routinely exceed the langgraph_sdk default 300s read deadline. The manager's runs.wait call kept an HTTP stream open for the entire run lifetime, so the long run blew up with httpx.ReadTimeout and the outer except branch then released the dedupe key and emitted a false 'internal error' outbound. The GitHub channel's outbound send is log-only by design: agents post to the issue/PR via the gh CLI in the sandbox when they choose to comment or create a PR. There is nothing for the manager to ferry back, so the long-poll was pure overhead. This change adds ChannelRunPolicy.fire_and_forget (default False) and sets it True for the github channel. When fire_and_forget is True, _handle_chat dispatches via client.runs.create (short POST, returns once the run is pending) instead of client.runs.wait, and skips the response-extraction + outbound-publish block. ConflictError on a busy thread still trips the standard THREAD_BUSY_MESSAGE path so behavior on the busy case is preserved for any future non-github fire-and-forget channel. Other (non-github) channels are unchanged: their policy defaults fire_and_forget=False and they continue to dispatch via runs.wait. Adds 6 regression tests in tests/test_channels.py::TestGithubFireAndForget: - Default ChannelRunPolicy.fire_and_forget is False. - The github policy registers fire_and_forget=True. - github inbound calls runs.create, not runs.wait, with the right kwargs. - github inbound publishes no outbound on success. - ConflictError from runs.create still emits THREAD_BUSY_MESSAGE. - Non-github channels (slack) still dispatch via runs.wait. * test(lead-agent): accept user_id kwarg in skill-policy test stubs The two GitHub-channel tests added in #3754 stubbed _load_enabled_skills_for_tool_policy with a lambda that only accepted `available_skills` and `app_config`, but the real function (and its call site in agent.py) also passes `user_id`. This raised TypeError on every run, failing backend-unit-tests. Add `user_id=None` to match the three sibling stubs in the same file. * refactor(gateway): disambiguate context-key set names The two frozensets _INTERNAL_ONLY_CONTEXT_KEYS and _CONTEXT_ONLY_KEYS shared a confusable "CONTEXT_ONLY" token in different orders, and the first broke the _CONTEXT_<X>_KEYS pattern of its sibling _CONTEXT_CONFIGURABLE_KEYS. Rename to make the distinct axes explicit: _CONTEXT_INTERNAL_CALLER_KEYS - WHO: internal callers (scheduler) only _CONTEXT_RUNTIME_ONLY_KEYS - WHERE: runtime context only, never configurable Pure rename, no behavior change.
1860 lines
68 KiB
Python
1860 lines
68 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,
|
|
_collect_pre_existing_message_ids,
|
|
_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
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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_snapshot():
|
|
snapshot = {
|
|
"checkpoint": {
|
|
"channel_values": {
|
|
"messages": [
|
|
AIMessage(id="a", content="x"),
|
|
AIMessage(id="b", content="y"),
|
|
AIMessage(content="no-id-here"), # ignored
|
|
]
|
|
}
|
|
}
|
|
}
|
|
assert _collect_pre_existing_message_ids(snapshot) == {"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({"checkpoint": None}) == set()
|
|
assert _collect_pre_existing_message_ids({"checkpoint": {}}) == set()
|
|
assert _collect_pre_existing_message_ids({"checkpoint": {"channel_values": None}}) == set()
|
|
assert _collect_pre_existing_message_ids({"checkpoint": {"channel_values": {"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 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"}}
|
|
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)
|