mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-14 16:08:41 +00:00
572744975d
* fix(tools): run tool assembly off-loop at async entry points get_available_tools() may block on MCP cache initialization while it is called on async agent-assembly paths (task_tool, durable batch execution), stalling the calling event loop for the full discovery duration. Dispatch the (unchanged, synchronous) assembly call to a worker thread via asyncio.to_thread at the two async entry points so the loop keeps processing requests, SSE frames, cancellations, and timers. Fixes #5172 * fix(tools): offload lead-agent assembly off-loop and pin with blocking-io anchors Review follow-up for #5224: - run_agent now dispatches agent_factory(...) through asyncio.to_thread, so lead-agent assembly (including both get_available_tools call sites in _assemble_lead_agent) runs off the event loop — the Gateway headline scenario from issue #5172. - _ensure_sync_invocable_tool takes a double-checked threading.Lock, making the in-place tool.func wrap on the shared tool singletons explicitly single-shot now that assembly can run concurrently on worker threads. - Add backend/tests/blocking_io/test_tool_assembly_offloop.py: blocking-probe anchors for task_tool and SubagentBatchService._execute_item under the strict Blockbuster gate, plus a meta-check proving the gate trips on the exact syscall class (ExtensionsConfig.from_file on the loop). Verified the anchor goes red when the offload is flattened back to a plain call. * fix(gateway): build checkpoint state accessor off-loop; anchor run_agent offload Review follow-up for #5224: - Add abuild_checkpoint_state_accessor (asyncio.to_thread around the unchanged sync builder) and switch every async call site to it: the stateless_wait route, thread_runs, both threads call sites, and the build_thread_checkpoint_state_accessor boundary. The agent-factory assembly re-enters get_available_tools() and may block on MCP cache initialization; repeat calls hit _state_accessor_graph_cache and only pay the thread hop. - Add a third blocking-io anchor driving the real run_agent with minimal RunManager/bridge stubs; the factory performs a real production blocking read (ExtensionsConfig.from_file()) and the test asserts assembly never runs on the main thread. Verified the anchor goes red when the run_agent offload is flattened back to a plain call. - Adapt the test_threads_router checkpoint-builder patch sites to the new async name. * refactor(tools): carry assembly offloads on a dedicated bounded pool Review follow-up for #5224: - Add utils/assembly_io.py: a dedicated ThreadPoolExecutor (default 8 workers, DEER_FLOW_ASSEMBLY_WORKERS-overridable, mirroring utils/file_io.py and tools/sync.py) with run_assembly(), which copies contextvars explicitly. A hung stdio MCP server parks its worker for the full MCP timeout; carrying assembly hops on the loop's default executor would let a few parked assemblies queue every other to_thread/run_in_executor(None, ...) caller behind them. - Switch all four offloads (run_agent, task_tool, batch _execute_item, abuild_checkpoint_state_accessor) to run_assembly(). - State the cold-path behavior in the accessor docstring: the graph cache validates factory identity, so non-identity-stable factories may duplicate lead-agent assembly across concurrent readers (MCP discovery stays process-wide single-flight); the pool bounds the duplicates. - Add a fourth blocking-io anchor driving build_thread_checkpoint_state_ accessor with a per-resolution fresh factory (always a cache miss) and the real production blocking read; enumerate all four offloads in the gate's module docstring. Verified the anchor goes red when abuild_checkpoint_state_accessor is flattened back to a plain call. * fix(subagents): revalidate batch item before launch; make assembly pool observable Review follow-up for #5224: - _execute_item() revalidates the durable state right after assembly and before executor.execute_async(): renew_item_lease() returns valid=False when cancel_batch() terminalized the item or the lease was lost while assembly was parked, and the launch is skipped (the canceller already finalized the item). Previously the launch was unconditional and the poll loop's cancellation checks only started after execution began. - Regression test driving the real SQLite repository: a blocking assembly probe parks _execute_item, cancel_batch() lands, and the launch is skipped with the item staying cancelled. Verified the test goes red when the revalidation is removed. - run_assembly() tracks pending assemblies and logs a throttled WARNING once the pending count exceeds the worker count, so assembly starvation (workers parked on a hung MCP server) is distinguishable from idle. - The run_agent blocking-io anchor now binds a sentinel extension snapshot via ctx.extensions and asserts the factory observed it through get_agent_build_extensions(), pinning run_assembly()'s ContextVar propagation. Verified red when ctx.run is dropped. - Document the assembly pool in backend/AGENTS.md. * fix(utils): decrement the assembly pending count on the pool thread The pending-assembly counter behind the starvation warning decremented from the asyncio future's done callback, which never fires once the submitting loop is closed while its worker is still running: the count ratcheted up permanently and eventually fired the starvation warning with no starvation behind it (reproduced at 97dc9bec by review). Decrement instead from the dispatched work item: run_assembly() wraps func so a finally drops the count under the pending lock on the pool thread, and the done callback is gone. Pin the counter with tests/test_assembly_io.py: a healthy call returns the count to zero, and an abandoned loop (stopped while the worker is parked) does not wedge it — the abandoned case goes red against the old done-callback decrement. * docs(utils): fix the pending-counter comment after the decrement move The comment still described the removed done-callback decrement, contradicting _work()'s own comment; state the actual mechanism (increment on the loop before dispatch, decrement from the dispatched work item's finally on a pool thread). * test(gateway): retarget checkpoint-accessor stubs to the services seam thread_runs and runs now call abuild_checkpoint_state_accessor, so the upstream wait-reader, regenerate-prepare, and idempotency tests must stub the sync builder where abuild resolves it (app.gateway.services); stubbing the removed router re-exports fails with AttributeError at setup. The async seam semantics are unchanged: run_assembly invokes the stubbed sync builder off-loop and propagates its return values and exceptions. Move the agent/tool assembly off-load note from backend/AGENTS.md to deerflow/utils/AGENTS.md (next to assembly_io.py) so the effective instruction chain for agents/middlewares no longer grows past the AG002 hard limit. * fix(runtime): serialize same-key accessor assembly and release queued-cancel slots Address the three review follow-ups on the assembly off-load: - assembly_io: a job cancelled while still queued never runs its work item, so the dispatched finally never fired and _pending_assemblies stayed elevated until a false starvation warning. Exactly-once cleanup now rides the concurrent future's cancelled() state — cancel() only succeeds before the executor starts the item, so cancelled() is true precisely when the finally will never run — plus a submit-failure release; the one-worker queued-cancellation case is pinned red/green. - services: overlapping cold readers sharing one cache key could both run full agent assembly. _state_accessor_graph now serializes per key through a thread-side KeyedLockTable (pool threads, no running loop) and re-validates factory/app-config identity under the lock, so the factory runs exactly once while identity changes still rebuild. Cache dict access is lock-guarded now that construction runs off-loop. - guidance inventory: register deerflow/utils/AGENTS.md in EXPECTED_GUIDANCE_PATHS so test_repository_has_the_approved_scoped_ guidance_shape matches the relocated assembly note (CI shard 4). * test(keyed-lock): pin KeyedLockTable reclamation and waiter bypass directly Thread-side counterparts of the async table's own tests: overlapping hold() calls serialize (a late arrival joins the live entry instead of creating a second lock that bypasses a queued waiter), the last check-in pops the entry, and many unique keys leave the registry empty. Both regressions verified red — popping unconditionally trips the late-arrival test, never reclaiming trips the many-keys test. --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
1502 lines
54 KiB
Python
1502 lines
54 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import copy
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
from _router_auth_helpers import call_unwrapped
|
|
from fastapi import HTTPException
|
|
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
|
|
from langgraph.checkpoint.base import empty_checkpoint, uuid6
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
|
|
from deerflow.runtime import RunStatus
|
|
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY
|
|
|
|
|
|
def _checkpoint(
|
|
checkpoint_id: str,
|
|
messages: list[object],
|
|
*,
|
|
metadata: dict | None = None,
|
|
goal: dict | None = None,
|
|
next_tasks: tuple[str, ...] = (),
|
|
):
|
|
channel_values = {"messages": messages}
|
|
if goal is not None:
|
|
channel_values["goal"] = goal
|
|
return SimpleNamespace(
|
|
config={
|
|
"configurable": {
|
|
"thread_id": "thread-1",
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": checkpoint_id,
|
|
"checkpoint_map": None,
|
|
}
|
|
},
|
|
checkpoint={"channel_values": channel_values},
|
|
metadata=metadata or {},
|
|
next=next_tasks,
|
|
)
|
|
|
|
|
|
async def _put_memory_checkpoint(
|
|
checkpointer: InMemorySaver,
|
|
thread_id: str,
|
|
messages: list[object],
|
|
*,
|
|
step: int,
|
|
parent_config: dict | None = None,
|
|
metadata: dict | None = None,
|
|
) -> dict:
|
|
checkpoint = empty_checkpoint()
|
|
checkpoint["id"] = str(uuid6())
|
|
checkpoint["channel_values"] = {"messages": messages}
|
|
checkpoint["channel_versions"] = {"messages": step}
|
|
checkpoint_metadata = {
|
|
"step": step,
|
|
"source": "loop",
|
|
"writes": {"test": {"messages": messages}},
|
|
"parents": {},
|
|
}
|
|
checkpoint_metadata.update(metadata or {})
|
|
return await checkpointer.aput(
|
|
parent_config or {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}},
|
|
checkpoint,
|
|
checkpoint_metadata,
|
|
{"messages": step},
|
|
)
|
|
|
|
|
|
async def _collect_checkpoints(checkpointer: InMemorySaver, config: dict) -> list:
|
|
return [checkpoint async for checkpoint in checkpointer.alist(config)]
|
|
|
|
|
|
class FakeCheckpointer:
|
|
def __init__(self, history, *, latest=None, materialized_history=None, materialized_latest=None):
|
|
self.history = history
|
|
self.latest = latest
|
|
self.materialized_history = materialized_history
|
|
self.materialized_latest = materialized_latest
|
|
self.alist_limits = []
|
|
|
|
async def aget_tuple(self, config):
|
|
checkpoint_id = config.get("configurable", {}).get("checkpoint_id")
|
|
if checkpoint_id:
|
|
return next((item for item in self.history if item.config["configurable"]["checkpoint_id"] == checkpoint_id), None)
|
|
return self.latest or (self.history[0] if self.history else None)
|
|
|
|
async def alist(self, config, limit=200):
|
|
self.alist_limits.append(limit)
|
|
for item in self.history[:limit]:
|
|
yield item
|
|
|
|
|
|
def _snapshot(checkpoint_id: str, messages: list[object], *, metadata: dict | None = None):
|
|
return SimpleNamespace(
|
|
values={"messages": messages},
|
|
config={
|
|
"configurable": {
|
|
"thread_id": "thread-1",
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": checkpoint_id,
|
|
"checkpoint_map": None,
|
|
}
|
|
},
|
|
metadata=metadata or {},
|
|
)
|
|
|
|
|
|
class FakeAccessor:
|
|
def __init__(self, checkpointer: FakeCheckpointer):
|
|
self.checkpointer = checkpointer
|
|
|
|
@staticmethod
|
|
def _from_raw(checkpoint):
|
|
return SimpleNamespace(
|
|
values=dict(checkpoint.checkpoint.get("channel_values", {})),
|
|
config=checkpoint.config,
|
|
metadata=checkpoint.metadata,
|
|
parent_config=getattr(checkpoint, "parent_config", None),
|
|
next=getattr(checkpoint, "next", ()),
|
|
)
|
|
|
|
async def aget(self, config):
|
|
materialized_latest = getattr(self.checkpointer, "materialized_latest", None)
|
|
if materialized_latest is not None and not config.get("configurable", {}).get("checkpoint_id"):
|
|
return materialized_latest
|
|
raw = await self.checkpointer.aget_tuple(config)
|
|
return self._from_raw(raw) if raw is not None else SimpleNamespace(values={}, config={}, metadata={})
|
|
|
|
async def ahistory(self, config, *, limit=None):
|
|
alist_limits = getattr(self.checkpointer, "alist_limits", None)
|
|
if alist_limits is not None:
|
|
alist_limits.append(limit)
|
|
history = getattr(self.checkpointer, "materialized_history", None)
|
|
if history is None:
|
|
if hasattr(self.checkpointer, "history"):
|
|
history = [self._from_raw(item) for item in self.checkpointer.history]
|
|
else:
|
|
history = [self._from_raw(item) async for item in self.checkpointer.alist(config, limit=limit)]
|
|
return history[:limit]
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _patch_checkpoint_accessor(monkeypatch):
|
|
from app.gateway import services
|
|
from app.gateway.routers import thread_runs
|
|
|
|
def build_accessor(request, *, thread_id, assistant_id=None, checkpoint_id=None):
|
|
config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
|
|
if checkpoint_id is not None:
|
|
config["configurable"]["checkpoint_id"] = checkpoint_id
|
|
return FakeAccessor(request.app.state.checkpointer), config
|
|
|
|
async def build_thread_accessor(request, *, thread_id, checkpoint_id=None):
|
|
return build_accessor(request, thread_id=thread_id, checkpoint_id=checkpoint_id)
|
|
|
|
monkeypatch.setattr(services, "build_checkpoint_state_accessor", build_accessor)
|
|
monkeypatch.setattr(thread_runs, "build_thread_checkpoint_state_accessor", build_thread_accessor)
|
|
|
|
|
|
class FakeEventStore:
|
|
def __init__(self, rows):
|
|
self.rows = rows
|
|
|
|
async def list_messages(self, thread_id, *, limit=50, before_seq=None, after_seq=None):
|
|
return self.rows[-limit:]
|
|
|
|
|
|
class FakeRunManager:
|
|
def __init__(self, records):
|
|
self.records = records
|
|
|
|
async def list_by_thread(self, thread_id, *, user_id=None, limit=100):
|
|
return self.records[:limit]
|
|
|
|
async def get(self, run_id, *, user_id=None):
|
|
return next((record for record in self.records if record.run_id == run_id), None)
|
|
|
|
|
|
def _request(checkpointer, event_store, *, run_manager=None, user_id="user-1"):
|
|
from app.gateway.auth_disabled import AUTH_SOURCE_SESSION
|
|
|
|
return SimpleNamespace(
|
|
app=SimpleNamespace(
|
|
state=SimpleNamespace(
|
|
checkpointer=checkpointer,
|
|
run_event_store=event_store,
|
|
run_manager=run_manager or FakeRunManager([]),
|
|
)
|
|
),
|
|
state=SimpleNamespace(user=SimpleNamespace(id=user_id), auth_source=AUTH_SOURCE_SESSION),
|
|
)
|
|
|
|
|
|
def test_run_wait_readers_return_materialized_final_values() -> None:
|
|
from app.gateway import services
|
|
from app.gateway.routers import runs, thread_runs
|
|
|
|
snapshot = SimpleNamespace(
|
|
values={
|
|
"messages": [
|
|
HumanMessage(id="h1", content="question"),
|
|
AIMessage(id="a1", content="answer"),
|
|
]
|
|
},
|
|
config={
|
|
"configurable": {
|
|
"thread_id": "thread-1",
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": "ckpt-2",
|
|
}
|
|
},
|
|
parent_config={"configurable": {"checkpoint_id": "ckpt-1"}},
|
|
metadata={"step": 2},
|
|
next=(),
|
|
tasks=(),
|
|
created_at=None,
|
|
)
|
|
accessor = SimpleNamespace(aget=AsyncMock(return_value=snapshot))
|
|
record = SimpleNamespace(
|
|
run_id="run-1",
|
|
thread_id="thread-1",
|
|
task=None,
|
|
status=RunStatus.success,
|
|
error=None,
|
|
)
|
|
request = SimpleNamespace()
|
|
body = thread_runs.RunCreateRequest(
|
|
assistant_id="lead-agent",
|
|
config={"configurable": {"thread_id": "thread-1"}},
|
|
)
|
|
|
|
async def _scenario() -> tuple[dict, dict]:
|
|
with (
|
|
patch.object(thread_runs, "get_stream_bridge", return_value=object()),
|
|
patch.object(thread_runs, "get_run_manager", return_value=object()),
|
|
patch.object(thread_runs, "start_run", AsyncMock(return_value=record)),
|
|
patch.object(
|
|
services,
|
|
"build_checkpoint_state_accessor",
|
|
create=True,
|
|
return_value=(accessor, snapshot.config),
|
|
),
|
|
patch.object(runs, "get_stream_bridge", return_value=object()),
|
|
patch.object(runs, "get_run_manager", return_value=object()),
|
|
patch.object(runs, "start_run", AsyncMock(return_value=record)),
|
|
patch.object(
|
|
services,
|
|
"build_checkpoint_state_accessor",
|
|
create=True,
|
|
return_value=(accessor, snapshot.config),
|
|
),
|
|
):
|
|
thread_result = await call_unwrapped(thread_runs.wait_run, "thread-1", body, request)
|
|
stateless_result = await call_unwrapped(runs.stateless_wait, body, request)
|
|
return thread_result, stateless_result
|
|
|
|
thread_result, stateless_result = asyncio.run(_scenario())
|
|
|
|
assert [message["id"] for message in thread_result["messages"]] == ["h1", "a1"]
|
|
assert [message["id"] for message in stateless_result["messages"]] == ["h1", "a1"]
|
|
|
|
|
|
def test_run_wait_readers_preserve_terminal_error_without_checkpoint() -> None:
|
|
from app.gateway import services
|
|
from app.gateway.routers import runs, thread_runs
|
|
|
|
snapshot = SimpleNamespace(
|
|
values={},
|
|
config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}},
|
|
parent_config=None,
|
|
metadata={},
|
|
next=(),
|
|
tasks=(),
|
|
created_at=None,
|
|
)
|
|
accessor = SimpleNamespace(aget=AsyncMock(return_value=snapshot))
|
|
record = SimpleNamespace(
|
|
run_id="run-1",
|
|
thread_id="thread-1",
|
|
task=None,
|
|
status=RunStatus.error,
|
|
error="run failed before checkpoint",
|
|
)
|
|
request = SimpleNamespace()
|
|
body = thread_runs.RunCreateRequest(config={"configurable": {"thread_id": "thread-1"}})
|
|
|
|
async def _scenario() -> tuple[dict, dict]:
|
|
with (
|
|
patch.object(thread_runs, "get_stream_bridge", return_value=object()),
|
|
patch.object(thread_runs, "get_run_manager", return_value=object()),
|
|
patch.object(thread_runs, "start_run", AsyncMock(return_value=record)),
|
|
patch.object(
|
|
services,
|
|
"build_checkpoint_state_accessor",
|
|
return_value=(accessor, snapshot.config),
|
|
),
|
|
patch.object(runs, "get_stream_bridge", return_value=object()),
|
|
patch.object(runs, "get_run_manager", return_value=object()),
|
|
patch.object(runs, "start_run", AsyncMock(return_value=record)),
|
|
patch.object(
|
|
services,
|
|
"build_checkpoint_state_accessor",
|
|
return_value=(accessor, snapshot.config),
|
|
),
|
|
):
|
|
thread_result = await call_unwrapped(thread_runs.wait_run, "thread-1", body, request)
|
|
stateless_result = await call_unwrapped(runs.stateless_wait, body, request)
|
|
return thread_result, stateless_result
|
|
|
|
thread_result, stateless_result = asyncio.run(_scenario())
|
|
|
|
expected = {"status": "error", "error": "run failed before checkpoint"}
|
|
assert thread_result == expected
|
|
assert stateless_result == expected
|
|
|
|
|
|
@pytest.mark.parametrize("route_name", ["thread", "stateless"])
|
|
def test_run_wait_readers_preserve_terminal_error_when_accessor_builder_fails(route_name: str) -> None:
|
|
from app.gateway import services
|
|
from app.gateway.routers import runs, thread_runs
|
|
|
|
record = SimpleNamespace(
|
|
run_id="run-1",
|
|
thread_id="thread-1",
|
|
task=None,
|
|
status=RunStatus.error,
|
|
error="run failed before checkpoint",
|
|
)
|
|
request = SimpleNamespace()
|
|
body = thread_runs.RunCreateRequest(config={"configurable": {"thread_id": "thread-1"}})
|
|
|
|
async def _scenario() -> dict:
|
|
if route_name == "thread":
|
|
with (
|
|
patch.object(thread_runs, "get_stream_bridge", return_value=object()),
|
|
patch.object(thread_runs, "get_run_manager", return_value=object()),
|
|
patch.object(thread_runs, "start_run", AsyncMock(return_value=record)),
|
|
patch.object(
|
|
services,
|
|
"build_checkpoint_state_accessor",
|
|
side_effect=RuntimeError("graph construction failed"),
|
|
),
|
|
):
|
|
return await call_unwrapped(thread_runs.wait_run, "thread-1", body, request)
|
|
|
|
with (
|
|
patch.object(runs, "get_stream_bridge", return_value=object()),
|
|
patch.object(runs, "get_run_manager", return_value=object()),
|
|
patch.object(runs, "start_run", AsyncMock(return_value=record)),
|
|
patch.object(
|
|
services,
|
|
"build_checkpoint_state_accessor",
|
|
side_effect=RuntimeError("graph construction failed"),
|
|
),
|
|
):
|
|
return await call_unwrapped(runs.stateless_wait, body, request)
|
|
|
|
result = asyncio.run(_scenario())
|
|
|
|
assert result == {"status": "error", "error": "run failed before checkpoint"}
|
|
|
|
|
|
def test_prepare_regenerate_payload_returns_clean_input_and_base_checkpoint():
|
|
from app.gateway.routers import thread_runs
|
|
|
|
human = HumanMessage(
|
|
id="human-1",
|
|
content="<current_uploads>injected</current_uploads>\n\n/data-analysis analyze data.csv",
|
|
additional_kwargs={
|
|
ORIGINAL_USER_CONTENT_KEY: "/data-analysis analyze data.csv",
|
|
"files": [{"filename": "data.csv", "path": "/mnt/user-data/uploads/data.csv"}],
|
|
},
|
|
)
|
|
ai = AIMessage(id="ai-1", content="answer v1")
|
|
base = _checkpoint("ckpt-base", [])
|
|
after_human = _checkpoint("ckpt-human", [human])
|
|
latest = _checkpoint("ckpt-ai", [human, ai])
|
|
checkpointer = FakeCheckpointer([latest, after_human, base])
|
|
event_store = FakeEventStore(
|
|
[
|
|
{
|
|
"run_id": "run-old",
|
|
"event_type": "llm.ai.response",
|
|
"category": "message",
|
|
"content": {"id": "ai-1", "type": "ai", "content": "answer v1"},
|
|
"metadata": {"caller": "lead_agent"},
|
|
}
|
|
]
|
|
)
|
|
|
|
original_builder = thread_runs.build_thread_checkpoint_state_accessor
|
|
thread_builder = AsyncMock(side_effect=original_builder)
|
|
with patch.object(thread_runs, "build_thread_checkpoint_state_accessor", thread_builder):
|
|
response = asyncio.run(thread_runs._prepare_regenerate_payload("thread-1", "ai-1", _request(checkpointer, event_store)))
|
|
|
|
assert [call.kwargs["thread_id"] for call in thread_builder.await_args_list] == ["thread-1", "thread-1"]
|
|
|
|
assert response.checkpoint == {
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": "ckpt-base",
|
|
"checkpoint_map": None,
|
|
}
|
|
assert response.target_run_id == "run-old"
|
|
assert response.metadata == {
|
|
"regenerate_from_message_id": "ai-1",
|
|
"regenerate_from_run_id": "run-old",
|
|
"regenerate_checkpoint_id": "ckpt-base",
|
|
}
|
|
assert "title" not in response.input
|
|
regenerated_human = response.input["messages"][0]
|
|
assert regenerated_human["id"] == "human-1"
|
|
assert regenerated_human["content"] == [{"type": "text", "text": "/data-analysis analyze data.csv"}]
|
|
assert regenerated_human["additional_kwargs"] == {"files": [{"filename": "data.csv", "path": "/mnt/user-data/uploads/data.csv"}]}
|
|
|
|
|
|
def test_prepare_regenerate_payload_preserves_latest_thread_title():
|
|
from app.gateway.routers import thread_runs
|
|
|
|
human = HumanMessage(id="human-1", content="question")
|
|
ai = AIMessage(id="ai-1", content="answer v1")
|
|
base = _checkpoint("ckpt-base", [])
|
|
latest = _checkpoint("ckpt-ai", [human, ai])
|
|
latest.checkpoint["channel_values"]["title"] = "User renamed title"
|
|
checkpointer = FakeCheckpointer([latest, base])
|
|
event_store = FakeEventStore(
|
|
[
|
|
{
|
|
"run_id": "run-old",
|
|
"event_type": "llm.ai.response",
|
|
"category": "message",
|
|
"content": {"id": "ai-1", "type": "ai", "content": "answer v1"},
|
|
"metadata": {"caller": "lead_agent"},
|
|
}
|
|
]
|
|
)
|
|
|
|
response = asyncio.run(
|
|
thread_runs._prepare_regenerate_payload(
|
|
"thread-1",
|
|
"ai-1",
|
|
_request(checkpointer, event_store),
|
|
)
|
|
)
|
|
|
|
assert response.checkpoint["checkpoint_id"] == "ckpt-base"
|
|
assert response.input["title"] == "User renamed title"
|
|
|
|
|
|
def test_prepare_regenerate_payload_supports_latest_interrupted_response_missing_from_checkpoint():
|
|
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
|
|
|
human = HumanMessage(
|
|
id="human-1",
|
|
content="question",
|
|
additional_kwargs={"run_id": "run-interrupted"},
|
|
)
|
|
base = _checkpoint("ckpt-base", [])
|
|
latest = _checkpoint("ckpt-human", [human])
|
|
checkpointer = FakeCheckpointer([latest, base])
|
|
run_manager = FakeRunManager(
|
|
[
|
|
SimpleNamespace(
|
|
run_id="run-interrupted",
|
|
thread_id="thread-1",
|
|
status=RunStatus.interrupted,
|
|
)
|
|
]
|
|
)
|
|
|
|
response = asyncio.run(
|
|
_prepare_regenerate_payload(
|
|
"thread-1",
|
|
"lc_run--partial-response",
|
|
_request(
|
|
checkpointer,
|
|
FakeEventStore([]),
|
|
run_manager=run_manager,
|
|
),
|
|
)
|
|
)
|
|
|
|
assert response.checkpoint["checkpoint_id"] == "ckpt-base"
|
|
assert response.target_run_id == "run-interrupted"
|
|
assert response.metadata == {
|
|
"regenerate_from_message_id": "lc_run--partial-response",
|
|
"regenerate_from_run_id": "run-interrupted",
|
|
"regenerate_checkpoint_id": "ckpt-base",
|
|
}
|
|
assert response.input["messages"][0]["id"] == "human-1"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("status", "record_thread_id"),
|
|
[
|
|
(RunStatus.success, "thread-1"),
|
|
(RunStatus.interrupted, "another-thread"),
|
|
],
|
|
)
|
|
def test_prepare_regenerate_payload_does_not_accept_unverified_missing_response(
|
|
status: RunStatus,
|
|
record_thread_id: str,
|
|
):
|
|
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
|
|
|
human = HumanMessage(
|
|
id="human-1",
|
|
content="question",
|
|
additional_kwargs={"run_id": "source-run"},
|
|
)
|
|
checkpointer = FakeCheckpointer(
|
|
[
|
|
_checkpoint("ckpt-human", [human]),
|
|
_checkpoint("ckpt-base", []),
|
|
]
|
|
)
|
|
run_manager = FakeRunManager(
|
|
[
|
|
SimpleNamespace(
|
|
run_id="source-run",
|
|
thread_id=record_thread_id,
|
|
status=status,
|
|
)
|
|
]
|
|
)
|
|
|
|
with pytest.raises(HTTPException) as exc:
|
|
asyncio.run(
|
|
_prepare_regenerate_payload(
|
|
"thread-1",
|
|
"missing-response",
|
|
_request(
|
|
checkpointer,
|
|
FakeEventStore([]),
|
|
run_manager=run_manager,
|
|
),
|
|
)
|
|
)
|
|
|
|
assert exc.value.status_code == 404
|
|
assert exc.value.detail == "Message missing-response not found"
|
|
|
|
|
|
def test_prepare_regenerate_payload_does_not_mutate_legacy_single_checkpoint_branch():
|
|
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
|
|
|
checkpointer = InMemorySaver()
|
|
source_thread_id = "source-thread"
|
|
branch_thread_id = "legacy-branch"
|
|
source_run_id = "source-run"
|
|
human = HumanMessage(id="human-1", content="question", additional_kwargs={"run_id": source_run_id})
|
|
ai = AIMessage(id="ai-1", content="answer")
|
|
|
|
async def _seed() -> str:
|
|
source_base_config = await _put_memory_checkpoint(checkpointer, source_thread_id, [], step=0)
|
|
after_human = await _put_memory_checkpoint(
|
|
checkpointer,
|
|
source_thread_id,
|
|
[human],
|
|
step=1,
|
|
parent_config=source_base_config,
|
|
)
|
|
source_head_config = await _put_memory_checkpoint(
|
|
checkpointer,
|
|
source_thread_id,
|
|
[human, ai],
|
|
step=2,
|
|
parent_config=after_human,
|
|
)
|
|
source_head = await checkpointer.aget_tuple(source_head_config)
|
|
assert source_head is not None
|
|
|
|
legacy_head = copy.deepcopy(source_head.checkpoint)
|
|
legacy_head_id = str(uuid6())
|
|
legacy_head["id"] = legacy_head_id
|
|
legacy_metadata = copy.deepcopy(source_head.metadata)
|
|
legacy_metadata.update(
|
|
{
|
|
"source": "branch",
|
|
"deerflow_branch": True,
|
|
"branch_parent_thread_id": source_thread_id,
|
|
"branch_parent_checkpoint_id": source_head_config["configurable"]["checkpoint_id"],
|
|
"branch_parent_message_id": "ai-1",
|
|
}
|
|
)
|
|
await checkpointer.aput(
|
|
{"configurable": {"thread_id": branch_thread_id, "checkpoint_ns": ""}},
|
|
legacy_head,
|
|
legacy_metadata,
|
|
dict(legacy_head["channel_versions"]),
|
|
)
|
|
return legacy_head_id
|
|
|
|
legacy_head_id = asyncio.run(_seed())
|
|
request = _request(checkpointer, FakeEventStore([]))
|
|
|
|
with pytest.raises(HTTPException) as exc:
|
|
asyncio.run(_prepare_regenerate_payload(branch_thread_id, "ai-1", request))
|
|
|
|
assert exc.value.status_code == 409
|
|
assert exc.value.detail == "Could not find an addressable checkpoint before the target user message"
|
|
latest = asyncio.run(checkpointer.aget_tuple({"configurable": {"thread_id": branch_thread_id, "checkpoint_ns": ""}}))
|
|
assert latest is not None
|
|
assert latest.config["configurable"]["checkpoint_id"] == legacy_head_id
|
|
branch_history = asyncio.run(_collect_checkpoints(checkpointer, {"configurable": {"thread_id": branch_thread_id, "checkpoint_ns": ""}}))
|
|
assert [item.config["configurable"]["checkpoint_id"] for item in branch_history] == [legacy_head_id]
|
|
|
|
|
|
def test_prepare_regenerate_payload_rejects_legacy_branch_when_source_checkpoint_is_missing():
|
|
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
|
|
|
checkpointer = InMemorySaver()
|
|
branch_thread_id = "legacy-orphan"
|
|
human = HumanMessage(id="human-1", content="question", additional_kwargs={"run_id": "source-run"})
|
|
ai = AIMessage(id="ai-1", content="answer")
|
|
|
|
async def _seed() -> None:
|
|
await _put_memory_checkpoint(
|
|
checkpointer,
|
|
branch_thread_id,
|
|
[human, ai],
|
|
step=1,
|
|
metadata={
|
|
"source": "branch",
|
|
"deerflow_branch": True,
|
|
"branch_parent_thread_id": "deleted-source",
|
|
"branch_parent_checkpoint_id": "missing-checkpoint",
|
|
"branch_parent_message_id": "ai-1",
|
|
},
|
|
)
|
|
|
|
asyncio.run(_seed())
|
|
with pytest.raises(HTTPException) as exc:
|
|
asyncio.run(
|
|
_prepare_regenerate_payload(
|
|
branch_thread_id,
|
|
"ai-1",
|
|
_request(checkpointer, FakeEventStore([])),
|
|
)
|
|
)
|
|
|
|
assert exc.value.status_code == 409
|
|
assert exc.value.detail == "Could not find an addressable checkpoint before the target user message"
|
|
branch_history = asyncio.run(_collect_checkpoints(checkpointer, {"configurable": {"thread_id": branch_thread_id, "checkpoint_ns": ""}}))
|
|
assert len(branch_history) == 1
|
|
|
|
|
|
def test_prepare_edit_regenerate_payload_returns_new_human_and_edit_metadata():
|
|
from app.gateway.routers import thread_runs
|
|
|
|
human = HumanMessage(
|
|
id="human-1",
|
|
content="<current_uploads>injected</current_uploads>\n\noriginal question",
|
|
name="researcher",
|
|
additional_kwargs={
|
|
ORIGINAL_USER_CONTENT_KEY: "original question",
|
|
"files": [{"filename": "data.csv", "path": "/mnt/user-data/uploads/data.csv"}],
|
|
"referenced_message_contexts": [{"message_id": "ai-prev", "quote": "quoted"}],
|
|
"hide_from_ui": False,
|
|
"run_id": "old-run",
|
|
"timestamp": "2026-07-22T00:00:00Z",
|
|
"middleware_private": "do-not-copy",
|
|
},
|
|
)
|
|
ai = AIMessage(id="ai-1", content="answer v1")
|
|
base = _checkpoint("ckpt-base", [])
|
|
after_human = _checkpoint("ckpt-human", [human])
|
|
latest = _checkpoint("ckpt-ai", [human, ai])
|
|
checkpointer = FakeCheckpointer([latest, after_human, base])
|
|
event_store = FakeEventStore(
|
|
[
|
|
{
|
|
"run_id": "run-old",
|
|
"event_type": "llm.ai.response",
|
|
"category": "message",
|
|
"content": {"id": "ai-1", "type": "ai", "content": "answer v1"},
|
|
"metadata": {"caller": "lead_agent"},
|
|
}
|
|
]
|
|
)
|
|
run_manager = FakeRunManager(
|
|
[
|
|
SimpleNamespace(
|
|
run_id="run-old",
|
|
status=RunStatus.success,
|
|
metadata={},
|
|
last_ai_message="answer v1",
|
|
)
|
|
]
|
|
)
|
|
|
|
response = asyncio.run(
|
|
thread_runs._prepare_edit_regenerate_payload(
|
|
"thread-1",
|
|
"human-1",
|
|
" updated question\nwith details ",
|
|
_request(checkpointer, event_store, run_manager=run_manager),
|
|
)
|
|
)
|
|
|
|
assert response.checkpoint == {
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": "ckpt-base",
|
|
"checkpoint_map": None,
|
|
}
|
|
assert response.target_run_id == "run-old"
|
|
assert response.replacement_human_message_id != "human-1"
|
|
assert response.source_message_ids == ["human-1", "ai-1"]
|
|
assert response.metadata == {
|
|
"replay_kind": "edit",
|
|
"regenerate_from_message_id": "ai-1",
|
|
"regenerate_from_run_id": "run-old",
|
|
"regenerate_checkpoint_id": "ckpt-base",
|
|
"edit_from_message_id": "human-1",
|
|
"edit_message_id": response.replacement_human_message_id,
|
|
"edit_version_group_id": "human-1",
|
|
}
|
|
replacement = response.input["messages"][0]
|
|
assert replacement == {
|
|
"type": "human",
|
|
"id": response.replacement_human_message_id,
|
|
"name": "researcher",
|
|
"content": [{"type": "text", "text": "updated question\nwith details"}],
|
|
"additional_kwargs": {
|
|
"files": [{"filename": "data.csv", "path": "/mnt/user-data/uploads/data.csv"}],
|
|
"referenced_message_contexts": [{"message_id": "ai-prev", "quote": "quoted"}],
|
|
},
|
|
}
|
|
|
|
|
|
def _first_turn_checkpointer() -> FakeCheckpointer:
|
|
"""First-turn history where the user message's id is swapped mid-run.
|
|
|
|
``DynamicContextMiddleware`` moves the first user message to ``{id}__user``
|
|
and gives ``{id}`` to the injected reminder, so every checkpoint written
|
|
before that node ran holds the same prompt under an id the replay-base
|
|
lookup cannot match (#4531).
|
|
"""
|
|
system = SystemMessage(id="human-1", content="<system-reminder>date</system-reminder>")
|
|
swapped_human = HumanMessage(id="human-1__user", content="original question")
|
|
raw_human = HumanMessage(id="human-1", content="original question")
|
|
ai = AIMessage(id="ai-1", content="answer v1")
|
|
return FakeCheckpointer(
|
|
[
|
|
_checkpoint("ckpt-head", [system, swapped_human, ai]),
|
|
_checkpoint("ckpt-after-inject", [system, swapped_human], next_tasks=("LoopDetectionMiddleware.before_agent",)),
|
|
_checkpoint("ckpt-mid", [raw_human], next_tasks=("DynamicContextMiddleware.before_agent",)),
|
|
_checkpoint("ckpt-input", [], next_tasks=("__start__",)),
|
|
_checkpoint("ckpt-empty", []),
|
|
]
|
|
)
|
|
|
|
|
|
def _answer_run_fixtures() -> tuple[FakeEventStore, FakeRunManager]:
|
|
event_store = FakeEventStore(
|
|
[
|
|
{
|
|
"run_id": "run-old",
|
|
"event_type": "llm.ai.response",
|
|
"category": "message",
|
|
"content": {"id": "ai-1", "type": "ai", "content": "answer v1"},
|
|
"metadata": {"caller": "lead_agent"},
|
|
}
|
|
]
|
|
)
|
|
run_manager = FakeRunManager([SimpleNamespace(run_id="run-old", status=RunStatus.success, metadata={}, last_ai_message="answer v1")])
|
|
return event_store, run_manager
|
|
|
|
|
|
def test_prepare_edit_regenerate_payload_preserves_a_rename_the_replay_base_predates():
|
|
"""Replaying a turn must not roll the thread back to an older title (#4457)."""
|
|
from app.gateway.routers import thread_runs
|
|
|
|
human = HumanMessage(id="human-1", content="original question")
|
|
ai = AIMessage(id="ai-1", content="answer v1")
|
|
base = _checkpoint("ckpt-base", [])
|
|
base.checkpoint["channel_values"]["title"] = "auto generated title"
|
|
latest = _checkpoint("ckpt-ai", [human, ai])
|
|
latest.checkpoint["channel_values"]["title"] = "User renamed title"
|
|
event_store, run_manager = _answer_run_fixtures()
|
|
|
|
response = asyncio.run(
|
|
thread_runs._prepare_edit_regenerate_payload(
|
|
"thread-1",
|
|
"human-1",
|
|
"edited question",
|
|
_request(FakeCheckpointer([latest, base]), event_store, run_manager=run_manager),
|
|
)
|
|
)
|
|
|
|
assert response.checkpoint["checkpoint_id"] == "ckpt-base"
|
|
assert response.input["title"] == "User renamed title"
|
|
|
|
|
|
def test_prepare_edit_regenerate_payload_lets_an_untitled_base_name_the_edited_turn():
|
|
"""A base with no title belongs to a thread that has not been named yet.
|
|
|
|
Pinning the current title there would keep a name generated from the prompt
|
|
the edit just replaced, so leave the channel empty and let the title
|
|
middleware name the rewritten turn.
|
|
"""
|
|
from app.gateway.routers import thread_runs
|
|
|
|
human = HumanMessage(id="human-1", content="original question")
|
|
ai = AIMessage(id="ai-1", content="answer v1")
|
|
latest = _checkpoint("ckpt-ai", [human, ai])
|
|
latest.checkpoint["channel_values"]["title"] = "title of the replaced question"
|
|
event_store, run_manager = _answer_run_fixtures()
|
|
|
|
response = asyncio.run(
|
|
thread_runs._prepare_edit_regenerate_payload(
|
|
"thread-1",
|
|
"human-1",
|
|
"edited question",
|
|
_request(FakeCheckpointer([latest, _checkpoint("ckpt-base", [])]), event_store, run_manager=run_manager),
|
|
)
|
|
)
|
|
|
|
assert response.checkpoint["checkpoint_id"] == "ckpt-base"
|
|
assert "title" not in response.input
|
|
|
|
|
|
def test_prepare_regenerate_payload_replays_the_pre_swap_user_message_id():
|
|
"""Replaying `{id}__user` would make the reminder middleware skip the turn.
|
|
|
|
The replay base predates the injection, so the turn must re-enter the graph
|
|
under the id the client originally sent or it loses its date/memory block.
|
|
"""
|
|
from app.gateway.routers import thread_runs
|
|
|
|
event_store, run_manager = _answer_run_fixtures()
|
|
|
|
response = asyncio.run(
|
|
thread_runs._prepare_regenerate_payload(
|
|
"thread-1",
|
|
"ai-1",
|
|
_request(_first_turn_checkpointer(), event_store, run_manager=run_manager),
|
|
)
|
|
)
|
|
|
|
assert response.checkpoint["checkpoint_id"] == "ckpt-empty"
|
|
assert response.input["messages"][0]["id"] == "human-1"
|
|
|
|
|
|
def test_prepare_edit_regenerate_payload_skips_mid_run_replay_base_on_first_turn():
|
|
"""The replay base must predate the turn, not sit inside the run that produced it.
|
|
|
|
``ckpt-mid`` still holds the original prompt (under its pre-swap id) and owns
|
|
the injection node's pending writes, so replaying from it re-adds the prompt
|
|
the edit is meant to replace.
|
|
"""
|
|
from app.gateway.routers import thread_runs
|
|
|
|
event_store, run_manager = _answer_run_fixtures()
|
|
|
|
response = asyncio.run(
|
|
thread_runs._prepare_edit_regenerate_payload(
|
|
"thread-1",
|
|
"human-1__user",
|
|
"edited question",
|
|
_request(_first_turn_checkpointer(), event_store, run_manager=run_manager),
|
|
)
|
|
)
|
|
|
|
assert response.checkpoint["checkpoint_id"] == "ckpt-empty"
|
|
assert response.metadata["regenerate_checkpoint_id"] == "ckpt-empty"
|
|
|
|
|
|
def test_prepare_edit_regenerate_payload_prefers_checkpoint_lineage():
|
|
"""Edit replay must resolve its base the same lineage-first way regenerate does.
|
|
|
|
A chronological scan cannot tell sibling branches apart (#4358), so the head
|
|
checkpoint has to reach the lineage walk.
|
|
"""
|
|
from app.gateway.routers import thread_runs
|
|
|
|
event_store, run_manager = _answer_run_fixtures()
|
|
checkpointer = _first_turn_checkpointer()
|
|
from app.gateway.checkpoint_lineage import CheckpointParentMissingError
|
|
|
|
walk = AsyncMock(side_effect=CheckpointParentMissingError("no parent link"))
|
|
|
|
with patch.object(thread_runs, "find_checkpoint_before_message", walk):
|
|
response = asyncio.run(
|
|
thread_runs._prepare_edit_regenerate_payload(
|
|
"thread-1",
|
|
"human-1__user",
|
|
"edited question",
|
|
_request(checkpointer, event_store, run_manager=run_manager),
|
|
)
|
|
)
|
|
|
|
assert walk.await_count == 1
|
|
head_checkpoint = walk.await_args.args[1]
|
|
assert head_checkpoint.config["configurable"]["checkpoint_id"] == "ckpt-head"
|
|
# Legacy checkpoints without parent links still degrade to the bounded scan.
|
|
assert response.checkpoint["checkpoint_id"] == "ckpt-empty"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("replacement_text", "detail"),
|
|
[
|
|
(" \n\t", "Edited message cannot be empty"),
|
|
(" original question ", "Edited message is unchanged"),
|
|
],
|
|
)
|
|
def test_prepare_edit_regenerate_payload_rejects_empty_or_unchanged_text(replacement_text: str, detail: str):
|
|
from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload
|
|
|
|
human = HumanMessage(id="human-1", content="original question")
|
|
ai = AIMessage(id="ai-1", content="answer")
|
|
checkpointer = FakeCheckpointer(
|
|
[
|
|
_checkpoint("ckpt-ai", [human, ai]),
|
|
_checkpoint("ckpt-human", [human]),
|
|
_checkpoint("ckpt-base", []),
|
|
]
|
|
)
|
|
event_store = FakeEventStore(
|
|
[
|
|
{
|
|
"run_id": "run-old",
|
|
"event_type": "llm.ai.response",
|
|
"category": "message",
|
|
"content": {"id": "ai-1", "type": "ai", "content": "answer"},
|
|
"metadata": {"caller": "lead_agent"},
|
|
}
|
|
]
|
|
)
|
|
run_manager = FakeRunManager([SimpleNamespace(run_id="run-old", status=RunStatus.success, metadata={}, last_ai_message="answer")])
|
|
|
|
with pytest.raises(HTTPException) as exc:
|
|
asyncio.run(
|
|
_prepare_edit_regenerate_payload(
|
|
"thread-1",
|
|
"human-1",
|
|
replacement_text,
|
|
_request(checkpointer, event_store, run_manager=run_manager),
|
|
)
|
|
)
|
|
|
|
assert exc.value.status_code == 409
|
|
assert exc.value.detail == detail
|
|
|
|
|
|
def test_prepare_edit_regenerate_payload_requires_latest_human_turn():
|
|
from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload
|
|
|
|
old_human = HumanMessage(id="human-old", content="old question")
|
|
old_ai = AIMessage(id="ai-old", content="old answer")
|
|
latest_human = HumanMessage(id="human-latest", content="latest question")
|
|
latest_ai = AIMessage(id="ai-latest", content="latest answer")
|
|
checkpointer = FakeCheckpointer(
|
|
[
|
|
_checkpoint("ckpt-latest", [old_human, old_ai, latest_human, latest_ai]),
|
|
_checkpoint("ckpt-base", []),
|
|
]
|
|
)
|
|
|
|
with pytest.raises(HTTPException) as exc:
|
|
asyncio.run(
|
|
_prepare_edit_regenerate_payload(
|
|
"thread-1",
|
|
"human-old",
|
|
"edited old question",
|
|
_request(checkpointer, FakeEventStore([])),
|
|
)
|
|
)
|
|
|
|
assert exc.value.status_code == 409
|
|
assert exc.value.detail == "Only the latest completed user turn can be edited"
|
|
|
|
|
|
def test_prepare_edit_regenerate_payload_allows_answered_historical_clarification():
|
|
from app.gateway.routers import thread_runs
|
|
|
|
old_human = HumanMessage(id="human-old", content="ambiguous request")
|
|
clarification = ToolMessage(
|
|
id="tool-clarify",
|
|
tool_call_id="call-clarify",
|
|
content="need input",
|
|
artifact={
|
|
"human_input": {
|
|
"version": 1,
|
|
"kind": "human_input_request",
|
|
"request_id": "clarify-1",
|
|
"prompt": "Which format?",
|
|
}
|
|
},
|
|
)
|
|
clarification_answer = HumanMessage(
|
|
id="human-clarify-answer",
|
|
content="Use Markdown",
|
|
additional_kwargs={
|
|
"hide_from_ui": True,
|
|
"human_input_response": {
|
|
"version": 1,
|
|
"kind": "human_input_response",
|
|
"request_id": "clarify-1",
|
|
"source": "user",
|
|
"value": "Use Markdown",
|
|
},
|
|
},
|
|
)
|
|
latest_human = HumanMessage(id="human-latest", content="latest question")
|
|
latest_ai = AIMessage(id="ai-latest", content="latest answer")
|
|
historical_messages = [old_human, clarification, clarification_answer]
|
|
checkpointer = FakeCheckpointer(
|
|
[
|
|
_checkpoint("ckpt-ai", [*historical_messages, latest_human, latest_ai]),
|
|
_checkpoint("ckpt-human", [*historical_messages, latest_human]),
|
|
_checkpoint("ckpt-base", historical_messages),
|
|
]
|
|
)
|
|
event_store = FakeEventStore(
|
|
[
|
|
{
|
|
"run_id": "run-latest",
|
|
"event_type": "llm.ai.response",
|
|
"category": "message",
|
|
"content": {"id": "ai-latest", "type": "ai", "content": "latest answer"},
|
|
"metadata": {"caller": "lead_agent"},
|
|
}
|
|
]
|
|
)
|
|
run_manager = FakeRunManager(
|
|
[
|
|
SimpleNamespace(
|
|
run_id="run-latest",
|
|
status=RunStatus.success,
|
|
metadata={},
|
|
last_ai_message="latest answer",
|
|
)
|
|
]
|
|
)
|
|
|
|
response = asyncio.run(
|
|
thread_runs._prepare_edit_regenerate_payload(
|
|
"thread-1",
|
|
"human-latest",
|
|
"updated latest question",
|
|
_request(checkpointer, event_store, run_manager=run_manager),
|
|
)
|
|
)
|
|
|
|
assert response.checkpoint["checkpoint_id"] == "ckpt-base"
|
|
assert response.target_run_id == "run-latest"
|
|
assert response.source_message_ids == ["human-latest", "ai-latest"]
|
|
|
|
|
|
def test_prepare_edit_regenerate_payload_rejects_open_clarification_turn():
|
|
from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload
|
|
|
|
human = HumanMessage(id="human-1", content="question")
|
|
tool = ToolMessage(
|
|
id="tool-1",
|
|
tool_call_id="call-1",
|
|
content="need input",
|
|
artifact={"human_input": {"request_id": "clarify-1", "status": "pending"}},
|
|
)
|
|
checkpointer = FakeCheckpointer(
|
|
[
|
|
_checkpoint("ckpt-tool", [human, tool]),
|
|
_checkpoint("ckpt-human", [human]),
|
|
_checkpoint("ckpt-base", []),
|
|
]
|
|
)
|
|
|
|
with pytest.raises(HTTPException) as exc:
|
|
asyncio.run(
|
|
_prepare_edit_regenerate_payload(
|
|
"thread-1",
|
|
"human-1",
|
|
"updated question",
|
|
_request(checkpointer, FakeEventStore([])),
|
|
)
|
|
)
|
|
|
|
assert exc.value.status_code == 409
|
|
assert exc.value.detail == "Only completed assistant text turns can be edited"
|
|
|
|
|
|
def test_prepare_edit_regenerate_payload_rejects_active_goal():
|
|
from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload
|
|
|
|
human = HumanMessage(id="human-1", content="question")
|
|
ai = AIMessage(id="ai-1", content="answer")
|
|
checkpointer = FakeCheckpointer(
|
|
[
|
|
_checkpoint("ckpt-ai", [human, ai], goal={"status": "active", "objective": "finish"}),
|
|
_checkpoint("ckpt-human", [human]),
|
|
_checkpoint("ckpt-base", []),
|
|
]
|
|
)
|
|
|
|
with pytest.raises(HTTPException) as exc:
|
|
asyncio.run(
|
|
_prepare_edit_regenerate_payload(
|
|
"thread-1",
|
|
"human-1",
|
|
"updated question",
|
|
_request(checkpointer, FakeEventStore([])),
|
|
)
|
|
)
|
|
|
|
assert exc.value.status_code == 409
|
|
assert exc.value.detail == "Cannot edit while a goal is active"
|
|
|
|
|
|
def test_prepare_edit_regenerate_payload_requires_successful_source_run():
|
|
from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload
|
|
|
|
human = HumanMessage(id="human-1", content="question")
|
|
ai = AIMessage(id="ai-1", content="answer")
|
|
checkpointer = FakeCheckpointer(
|
|
[
|
|
_checkpoint("ckpt-ai", [human, ai]),
|
|
_checkpoint("ckpt-human", [human]),
|
|
_checkpoint("ckpt-base", []),
|
|
]
|
|
)
|
|
event_store = FakeEventStore(
|
|
[
|
|
{
|
|
"run_id": "run-old",
|
|
"event_type": "llm.ai.response",
|
|
"category": "message",
|
|
"content": {"id": "ai-1", "type": "ai", "content": "answer"},
|
|
"metadata": {"caller": "lead_agent"},
|
|
}
|
|
]
|
|
)
|
|
run_manager = FakeRunManager([SimpleNamespace(run_id="run-old", status=RunStatus.error, metadata={}, last_ai_message="answer")])
|
|
|
|
with pytest.raises(HTTPException) as exc:
|
|
asyncio.run(
|
|
_prepare_edit_regenerate_payload(
|
|
"thread-1",
|
|
"human-1",
|
|
"updated question",
|
|
_request(checkpointer, event_store, run_manager=run_manager),
|
|
)
|
|
)
|
|
|
|
assert exc.value.status_code == 409
|
|
assert exc.value.detail == "Only successful assistant runs can be edited and rerun"
|
|
|
|
|
|
def test_prepare_regenerate_uses_materialized_history_when_raw_messages_are_omitted():
|
|
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
|
|
|
earlier_human = HumanMessage(id="human-0", content="earlier question")
|
|
earlier_ai = AIMessage(id="ai-0", content="earlier answer")
|
|
target_human = HumanMessage(id="human-1", content="question")
|
|
target_ai = AIMessage(id="ai-1", content="answer")
|
|
|
|
raw_latest = _checkpoint("ckpt-ai", [])
|
|
raw_after_human = _checkpoint("ckpt-human", [])
|
|
raw_base = _checkpoint("ckpt-base", [])
|
|
materialized_history = [
|
|
_snapshot("ckpt-ai", [earlier_human, earlier_ai, target_human, target_ai]),
|
|
_snapshot("ckpt-human", [earlier_human, earlier_ai, target_human]),
|
|
_snapshot("ckpt-base", [earlier_human, earlier_ai]),
|
|
]
|
|
checkpointer = FakeCheckpointer(
|
|
[raw_latest, raw_after_human, raw_base],
|
|
latest=raw_latest,
|
|
materialized_history=materialized_history,
|
|
materialized_latest=materialized_history[0],
|
|
)
|
|
event_store = FakeEventStore(
|
|
[
|
|
{
|
|
"run_id": "run-target",
|
|
"event_type": "ai_message",
|
|
"category": "message",
|
|
"content": {"id": "ai-1", "type": "ai", "content": "answer"},
|
|
"metadata": {"caller": "lead_agent"},
|
|
}
|
|
]
|
|
)
|
|
|
|
response = asyncio.run(
|
|
_prepare_regenerate_payload(
|
|
"thread-1",
|
|
"ai-1",
|
|
_request(checkpointer, event_store),
|
|
)
|
|
)
|
|
|
|
assert response.checkpoint["checkpoint_id"] == "ckpt-base"
|
|
assert response.metadata["regenerate_checkpoint_id"] == "ckpt-base"
|
|
assert response.input["messages"][0]["id"] == "human-1"
|
|
assert checkpointer.alist_limits == [400]
|
|
|
|
|
|
def test_prepare_regenerate_rejects_cyclic_lineage_without_chronological_fallback():
|
|
from app.gateway.routers import thread_runs
|
|
|
|
human = HumanMessage(id="human-1", content="question")
|
|
ai = AIMessage(id="ai-1", content="answer")
|
|
|
|
def linked_snapshot(checkpoint_id: str, messages: list[object], parent_id: str) -> SimpleNamespace:
|
|
return SimpleNamespace(
|
|
values={"messages": messages},
|
|
config={
|
|
"configurable": {
|
|
"thread_id": "thread-1",
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": checkpoint_id,
|
|
}
|
|
},
|
|
metadata={},
|
|
parent_config={
|
|
"configurable": {
|
|
"thread_id": "thread-1",
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": parent_id,
|
|
}
|
|
},
|
|
)
|
|
|
|
head = linked_snapshot("head", [human, ai], "cycle")
|
|
cycle = linked_snapshot("cycle", [human], "head")
|
|
wrong_sibling_base = linked_snapshot("wrong-sibling", [], "root")
|
|
by_id = {"head": head, "cycle": cycle}
|
|
|
|
async def aget(config):
|
|
return by_id[config["configurable"]["checkpoint_id"]]
|
|
|
|
accessor = SimpleNamespace(
|
|
aget=aget,
|
|
ahistory=AsyncMock(return_value=[head, wrong_sibling_base]),
|
|
)
|
|
builder = AsyncMock(
|
|
return_value=(
|
|
accessor,
|
|
{"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}},
|
|
)
|
|
)
|
|
|
|
with patch.object(thread_runs, "build_thread_checkpoint_state_accessor", builder):
|
|
with pytest.raises(HTTPException) as exc:
|
|
asyncio.run(
|
|
thread_runs._find_base_checkpoint_before_human(
|
|
"thread-1",
|
|
"human-1",
|
|
_request(FakeCheckpointer([]), FakeEventStore([])),
|
|
head_checkpoint=head,
|
|
)
|
|
)
|
|
|
|
assert exc.value.status_code == 409
|
|
assert exc.value.detail == "Could not safely resolve the checkpoint before the target user message"
|
|
accessor.ahistory.assert_not_awaited()
|
|
|
|
|
|
def test_prepare_regenerate_rejects_dangling_parent_without_chronological_fallback():
|
|
from app.gateway.routers import thread_runs
|
|
|
|
human = HumanMessage(id="human-1", content="question")
|
|
ai = AIMessage(id="ai-1", content="answer")
|
|
head = SimpleNamespace(
|
|
values={"messages": [human, ai]},
|
|
config={
|
|
"configurable": {
|
|
"thread_id": "thread-1",
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": "head",
|
|
}
|
|
},
|
|
metadata={},
|
|
parent_config={
|
|
"configurable": {
|
|
"thread_id": "thread-1",
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": "missing",
|
|
}
|
|
},
|
|
)
|
|
missing = SimpleNamespace(
|
|
values={},
|
|
config=head.parent_config,
|
|
metadata=None,
|
|
created_at=None,
|
|
parent_config=None,
|
|
)
|
|
accessor = SimpleNamespace(
|
|
aget=AsyncMock(return_value=missing),
|
|
ahistory=AsyncMock(return_value=[head, _snapshot("wrong-sibling", [])]),
|
|
)
|
|
builder = AsyncMock(
|
|
return_value=(
|
|
accessor,
|
|
{"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}},
|
|
)
|
|
)
|
|
|
|
with patch.object(thread_runs, "build_thread_checkpoint_state_accessor", builder):
|
|
with pytest.raises(HTTPException) as exc:
|
|
asyncio.run(
|
|
thread_runs._find_base_checkpoint_before_human(
|
|
"thread-1",
|
|
"human-1",
|
|
_request(FakeCheckpointer([]), FakeEventStore([])),
|
|
head_checkpoint=head,
|
|
)
|
|
)
|
|
|
|
assert exc.value.status_code == 409
|
|
assert exc.value.detail == "Could not safely resolve the checkpoint before the target user message"
|
|
accessor.ahistory.assert_not_awaited()
|
|
|
|
|
|
def test_prepare_regenerate_payload_rejects_non_latest_assistant():
|
|
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
|
|
|
human = HumanMessage(id="human-1", content="question")
|
|
old_ai = AIMessage(id="ai-old", content="old")
|
|
latest_ai = AIMessage(id="ai-latest", content="latest")
|
|
base = _checkpoint("ckpt-base", [])
|
|
after_human = _checkpoint("ckpt-human", [human])
|
|
latest = _checkpoint("ckpt-latest", [human, old_ai, latest_ai])
|
|
checkpointer = FakeCheckpointer([latest, after_human, base])
|
|
event_store = FakeEventStore(
|
|
[
|
|
{
|
|
"run_id": "run-old",
|
|
"event_type": "ai_message",
|
|
"category": "message",
|
|
"content": {"id": "ai-old", "type": "ai", "content": "old"},
|
|
"metadata": {"caller": "lead_agent"},
|
|
}
|
|
]
|
|
)
|
|
|
|
with pytest.raises(HTTPException) as exc:
|
|
asyncio.run(_prepare_regenerate_payload("thread-1", "ai-old", _request(checkpointer, event_store)))
|
|
|
|
assert exc.value.status_code == 409
|
|
assert exc.value.detail == "Only the latest assistant message can be regenerated"
|
|
|
|
|
|
def test_prepare_regenerate_payload_falls_back_to_matching_run_when_events_are_missing():
|
|
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
|
|
|
human = HumanMessage(id="human-1", content="question")
|
|
ai = AIMessage(id="ai-1", content="answer")
|
|
base = _checkpoint("ckpt-base", [])
|
|
after_human = _checkpoint("ckpt-human", [human])
|
|
latest = _checkpoint("ckpt-ai", [human, ai])
|
|
checkpointer = FakeCheckpointer([latest, after_human, base])
|
|
run_manager = FakeRunManager(
|
|
[
|
|
SimpleNamespace(run_id="run-latest", status=RunStatus.success, last_ai_message="answer"),
|
|
SimpleNamespace(run_id="run-older", status=RunStatus.error, last_ai_message="answer"),
|
|
]
|
|
)
|
|
|
|
response = asyncio.run(
|
|
_prepare_regenerate_payload(
|
|
"thread-1",
|
|
"ai-1",
|
|
_request(checkpointer, FakeEventStore([]), run_manager=run_manager),
|
|
)
|
|
)
|
|
|
|
assert response.target_run_id == "run-latest"
|
|
assert response.metadata["regenerate_from_run_id"] == "run-latest"
|
|
|
|
|
|
def test_prepare_regenerate_payload_uses_server_stamped_human_run_id_without_parent_events():
|
|
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
|
|
|
human = HumanMessage(id="human-1", content="question", additional_kwargs={"run_id": "parent-run"})
|
|
ai = AIMessage(id="ai-1", content="answer")
|
|
base = _checkpoint("ckpt-base", [])
|
|
after_human = _checkpoint("ckpt-human", [human])
|
|
latest = _checkpoint(
|
|
"ckpt-ai",
|
|
[human, ai],
|
|
metadata={
|
|
"deerflow_branch": True,
|
|
"branch_parent_thread_id": "parent-thread",
|
|
"branch_parent_checkpoint_id": "parent-checkpoint",
|
|
},
|
|
)
|
|
checkpointer = FakeCheckpointer([latest, after_human, base])
|
|
event_store = FakeEventStore([])
|
|
|
|
response = asyncio.run(_prepare_regenerate_payload("thread-1", "ai-1", _request(checkpointer, event_store)))
|
|
|
|
assert response.target_run_id == "parent-run"
|
|
|
|
|
|
def test_prepare_regenerate_payload_rejects_unverified_run_fallback_when_events_are_missing():
|
|
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
|
|
|
human = HumanMessage(id="human-1", content="question")
|
|
ai = AIMessage(id="ai-1", content="answer")
|
|
base = _checkpoint("ckpt-base", [])
|
|
after_human = _checkpoint("ckpt-human", [human])
|
|
latest = _checkpoint("ckpt-ai", [human, ai])
|
|
checkpointer = FakeCheckpointer([latest, after_human, base])
|
|
run_manager = FakeRunManager(
|
|
[
|
|
SimpleNamespace(run_id="run-latest", status=RunStatus.success, last_ai_message="different"),
|
|
]
|
|
)
|
|
|
|
with pytest.raises(HTTPException) as exc:
|
|
asyncio.run(
|
|
_prepare_regenerate_payload(
|
|
"thread-1",
|
|
"ai-1",
|
|
_request(checkpointer, FakeEventStore([]), run_manager=run_manager),
|
|
)
|
|
)
|
|
|
|
assert exc.value.status_code == 409
|
|
assert exc.value.detail == "Could not find source run for assistant message"
|
|
|
|
|
|
def test_prepare_regenerate_payload_requires_addressable_checkpoint_before_human():
|
|
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
|
|
|
human = HumanMessage(id="human-1", content="question")
|
|
ai = AIMessage(id="ai-1", content="answer")
|
|
latest = _checkpoint("ckpt-ai", [human, ai])
|
|
checkpointer = FakeCheckpointer([latest])
|
|
event_store = FakeEventStore(
|
|
[
|
|
{
|
|
"run_id": "run-old",
|
|
"event_type": "llm.ai.response",
|
|
"category": "message",
|
|
"content": {"id": "ai-1", "type": "ai", "content": "answer"},
|
|
"metadata": {"caller": "lead_agent"},
|
|
}
|
|
]
|
|
)
|
|
|
|
with pytest.raises(HTTPException) as exc:
|
|
asyncio.run(_prepare_regenerate_payload("thread-1", "ai-1", _request(checkpointer, event_store)))
|
|
|
|
assert exc.value.status_code == 409
|
|
assert exc.value.detail == "Could not find an addressable checkpoint before the target user message"
|
|
assert checkpointer.alist_limits == [400]
|
|
|
|
|
|
def test_prepare_regenerate_payload_reports_recent_checkpoint_scan_limit():
|
|
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
|
|
|
human = HumanMessage(id="human-1", content="question")
|
|
ai = AIMessage(id="ai-1", content="answer")
|
|
latest = _checkpoint("ckpt-latest", [human, ai])
|
|
history_without_human = [_checkpoint(f"ckpt-{index}", []) for index in range(201)]
|
|
checkpointer = FakeCheckpointer(history_without_human, latest=latest)
|
|
event_store = FakeEventStore(
|
|
[
|
|
{
|
|
"run_id": "run-old",
|
|
"event_type": "llm.ai.response",
|
|
"category": "message",
|
|
"content": {"id": "ai-1", "type": "ai", "content": "answer"},
|
|
"metadata": {"caller": "lead_agent"},
|
|
}
|
|
]
|
|
)
|
|
|
|
with pytest.raises(HTTPException) as exc:
|
|
asyncio.run(_prepare_regenerate_payload("thread-1", "ai-1", _request(checkpointer, event_store)))
|
|
|
|
assert exc.value.status_code == 409
|
|
assert exc.value.detail == "Could not locate target user message in recent checkpoint history (limit=200)"
|
|
assert checkpointer.alist_limits == [400]
|
|
|
|
|
|
def test_find_base_checkpoint_ignores_duration_only_checkpoints() -> None:
|
|
from app.gateway.routers.thread_runs import _find_base_checkpoint_before_human
|
|
|
|
human = HumanMessage(id="human-1", content="question")
|
|
duration_checkpoints = [
|
|
_checkpoint(
|
|
f"duration-{index}",
|
|
[],
|
|
metadata={"writes": {"runtime_run_duration": {"run_ids": [f"run-{index}"]}}},
|
|
)
|
|
for index in range(200)
|
|
]
|
|
base = _checkpoint("ckpt-base", [])
|
|
after_human = _checkpoint("ckpt-human", [human])
|
|
checkpointer = FakeCheckpointer([*duration_checkpoints, after_human, base])
|
|
|
|
result = asyncio.run(_find_base_checkpoint_before_human("thread-1", "human-1", _request(checkpointer, FakeEventStore([]))))
|
|
|
|
assert result.config == base.config
|
|
assert checkpointer.alist_limits == [400]
|