wehub-resource-sync
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)
|