mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-14 16:08:41 +00:00
7389331e65
* feat(extensions): observe task lifecycle and system model calls PR 1 (#4636) gave extensions a middleware chain, and a middleware only sees what passes through the agent graph. Two runtime surfaces stay invisible to it: when a lead run or a subagent begins and ends, and the DeerFlow-owned model calls made outside the graph. This slice adds both, with no new Gateway surface -- routers, services, and the reference extension stay in PR 3. Contract (deerflow-extension-api 0.1.1) --------------------------------------- Two contribution kinds join `middlewares` on the registry: `task_lifecycle` (`on_task_start` / `on_task_stop`, receiving a `TaskInfo` and a conservative `TaskOutcome` of completed / aborted / failed) and `system_model_observer` (`on_system_model_call`, receiving a `SystemOperationKind`, a `SystemModelRequest` snapshot, and a `SystemModelResult` carrying either the response or the provider exception plus a duration). `SystemModelRequest.messages` normalizes to a tuple at construction. Goal evaluation and memory extraction pass a message list while title generation and summarization pass one prompt string, and a bare `str` already satisfies `Sequence` -- without normalization an observer iterating `request.messages` would silently walk characters. Copying also makes the frozen snapshot immutable in fact rather than only by declaration, since observations may run after the call site returns and keeps mutating its own list. Registry marks and rollbacks become per-bucket and positional, so an `install()` that fails after registering two different kinds cannot leave one of them behind. `needs_task_store` now covers all three kinds: a deployment that registers only lifecycle hooks still gets a task store. Task lifecycle -------------- The lead worker notifies start after the run has started and stop after completion persistence and the completion hook, but before clearing the finalizing barrier and publishing the stream end -- holding the barrier across stop is what keeps a same-thread replacement run from overlapping this task's lifecycle. Cancellation raised out of the stop notification is deferred, not propagated in place, so a cancelled run still clears the barrier and emits its end frame. A subagent with a parent `run_id` wraps its execution in the same pair inside `finally`, reporting `parent_task_id` so a delegation tree is reconstructable; a subagent without a `run_id` (embedded client, standalone LangGraph Server) logs and skips rather than inventing a parent. Contributors run in registration order inside one shared 3s budget and every failure is logged and failed open. System model calls ------------------ Four kinds cover the model calls the middleware chain cannot see: goal evaluation, memory extraction, title generation, and summarization. Each site reports both terminal paths without changing the provider exception the host observes, short-circuits on `has_system_model_observers`, and passes the live task store when the runtime has one (detached work gets an isolated store). The sync summarization half stays unobserved on purpose -- it and its only host caller are the sync side of an async-only runtime, so notifying there would block a thread on a call site the host never reaches; the reason is recorded at the call site. The DeerMem backend must stay vendorable and cannot import the extension API, so it reports through a new `MemoryCallbacks.on_memory_llm_result` host hook that the DeerFlow-side callbacks translate into an observation. Notification loop ----------------- Extension resources must be touched on the loop that created them, but subagents can execute on isolated loops and DeerMem runs on a worker thread. The Gateway registers its serving loop before any runtime dependency starts and resets it last through the exit stack, so every startup-failure and cancellation path is covered. Awaited hooks raised on another loop are dispatched across with `run_coroutine_threadsafe` and awaited under the same budget; synchronous sites submit fire-and-forget work. Shutdown stops accepting detached observations before the memory flush -- that flush runs on a worker thread and can emit memory observations -- while keeping the loop alive for awaited task hooks until run and subagent drain completes. Tests ----- `test_extension_task_lifecycle.py`, `test_extension_subagent_lifecycle.py`, and `test_extension_system_model_calls.py` cover ordering, fail-open, budget exhaustion, snapshot binding under a concurrent singleton replacement, the loop-dispatch and shutdown-suspension paths, and both terminal paths at every call site. `test_gateway_run_drain_shutdown.py` pins the stop-before-barrier and drain ordering. * fix(extensions): decide notification fail-open by origin, observe cancellation `_notify_each` only guarded `Exception`, so a contributor letting a `CancelledError` escape — an extension implementing an internal timeout with cancellation, say — skipped its successors and reached the worker's deferred-interrupt path, ending an otherwise successful run as cancelled. Fail-open is about where a failure came from, not its base class: only a genuine cancellation of the host task increments `Task.cancelling()`, so propagate on that and contain everything else. `KeyboardInterrupt` / `SystemExit` still propagate. `observe_system_model_call` skipped observers on cancellation for the same base-class reason, leaving goal / title / summarization silent on a terminal path that is routine — interrupt/rollback admission and shutdown both cancel the run task, with the provider tokens already spent. Awaiting observers there is unreliable (a repeated cancel interrupts that await before any of them runs), so report through the same non-blocking submission the synchronous memory bridge uses, then propagate the cancellation untouched. DeerMem keeps `BaseException` around its provider call, now with the reason recorded: that path runs on a worker thread, where cancelling the awaiting side never interrupts the running thread, so `CancelledError` cannot arrive at all. Its host-hook wrapper narrows to `Exception` — only the hook's own failures are non-fatal, and an observability path must not swallow a process teardown signal. * fix(extensions): warn on budget exhaustion, scope observer logs by task, propagate teardown Review response on #4684: - The memory observation bridge caught BaseException, which would swallow a teardown signal raised while dispatching; it now catches Exception, matching the boundary the DeerMem-side call site documents and tests. - A notification-budget timeout raised mid-hook fell into the generic hook-failure path and logged an asyncio-internal traceback; it now logs a warning like the pre-hook budget skip, while a TimeoutError a contributor raises on its own stays classified as a hook failure. - System model observer logs passed the operation kind as the task id, so log lines said "task goal/title/..."; they now carry the task scope id alongside the kind.
964 lines
36 KiB
Python
964 lines
36 KiB
Python
import asyncio
|
|
import copy
|
|
|
|
import pytest
|
|
from deerflow_extension_api import ExtensionData
|
|
from langchain_core.messages import AIMessage, HumanMessage
|
|
from langgraph.checkpoint.base import empty_checkpoint, uuid6
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
|
|
from deerflow.extensions.registry import ExtensionRegistry
|
|
from deerflow.runtime.checkpoint_state import CheckpointStateAccessor, build_state_mutation_graph
|
|
from deerflow.runtime.goal import GoalEvaluation, attach_goal_evaluation, build_goal_state, latest_visible_assistant_signature, read_thread_goal, write_thread_goal
|
|
from deerflow.runtime.runs import worker
|
|
from deerflow.runtime.runs.manager import RunRecord, RunStartOutcome
|
|
from deerflow.runtime.runs.schemas import DisconnectMode, RunStatus
|
|
|
|
|
|
def _full_accessor(checkpointer) -> CheckpointStateAccessor:
|
|
"""Bind a full-mode accessor over a state-only graph for materialized reads."""
|
|
graph = build_state_mutation_graph("goal_evaluator", "full")
|
|
return CheckpointStateAccessor.bind(graph, checkpointer, mode="full")
|
|
|
|
|
|
class _CollectingBridge:
|
|
def __init__(self) -> None:
|
|
self.events: list[tuple[str, object]] = []
|
|
|
|
async def publish(self, _run_id: str, event: str, payload: object) -> None:
|
|
self.events.append((event, payload))
|
|
|
|
|
|
class _ClearBeforeSecondGoalReadCheckpointer:
|
|
"""Wrap a saver and clear the goal just before the evaluator write rereads.
|
|
|
|
The first ``aget_tuple`` is the evaluator's current-goal read. The second is
|
|
``write_thread_goal`` preparing its read-modify-write. Clearing at that point
|
|
models a user ``/goal clear`` landing between those two operations.
|
|
"""
|
|
|
|
def __init__(self, inner: InMemorySaver, thread_id: str) -> None:
|
|
self.inner = inner
|
|
self.thread_id = thread_id
|
|
self.read_count = 0
|
|
self.cleared = False
|
|
|
|
def get_next_version(self, current, channel):
|
|
return self.inner.get_next_version(current, channel)
|
|
|
|
async def aget_tuple(self, config):
|
|
self.read_count += 1
|
|
if self.read_count == 2 and not self.cleared:
|
|
self.cleared = True
|
|
await write_thread_goal(self.inner, self.thread_id, None, as_node="test_clear")
|
|
return await self.inner.aget_tuple(config)
|
|
|
|
async def aput(self, *args, **kwargs):
|
|
return await self.inner.aput(*args, **kwargs)
|
|
|
|
|
|
class _RaceAfterFirstContinuationCommitCheckpointer:
|
|
"""Wrap a saver and inject a racing user message right after the first
|
|
goal-continuation commit lands.
|
|
|
|
``_prepare_goal_continuation_input``'s real continuation commit (the
|
|
``_persist(..., continuation_count=next_count)`` call that records the
|
|
evaluator's decision to continue) performs this scenario's first
|
|
``aput``. Injecting a racing visible message immediately after that write
|
|
lands lets the worker's trailing visible-conversation-signature re-check
|
|
observe a thread change that happened *after* the continuation was
|
|
committed but *before* that re-check runs -- modelling the
|
|
``thread_changed_before_continuation`` race.
|
|
"""
|
|
|
|
def __init__(self, inner: InMemorySaver, thread_id: str) -> None:
|
|
self.inner = inner
|
|
self.thread_id = thread_id
|
|
self.put_count = 0
|
|
|
|
def get_next_version(self, current, channel):
|
|
return self.inner.get_next_version(current, channel)
|
|
|
|
async def aget_tuple(self, config):
|
|
return await self.inner.aget_tuple(config)
|
|
|
|
async def aput(self, *args, **kwargs):
|
|
result = await self.inner.aput(*args, **kwargs)
|
|
self.put_count += 1
|
|
if self.put_count == 1:
|
|
checkpoint_tuple = await self.inner.aget_tuple({"configurable": {"thread_id": self.thread_id, "checkpoint_ns": ""}})
|
|
checkpoint = getattr(checkpoint_tuple, "checkpoint", {}) or {}
|
|
channel_values = checkpoint.get("channel_values", {}) or {}
|
|
current_messages = channel_values.get("messages", []) or []
|
|
await _write_messages(
|
|
self.inner,
|
|
thread_id=self.thread_id,
|
|
messages=[*current_messages, HumanMessage(content="Actually, stop and wait.")],
|
|
)
|
|
return result
|
|
|
|
|
|
async def _seed_goal_thread(
|
|
checkpointer: InMemorySaver,
|
|
*,
|
|
thread_id: str,
|
|
goal_text: str,
|
|
messages: list | None = None,
|
|
) -> None:
|
|
checkpoint = empty_checkpoint()
|
|
checkpoint["channel_values"] = {
|
|
"messages": messages
|
|
or [
|
|
HumanMessage(content="Please finish this task."),
|
|
AIMessage(content="I made a start, but I am not done."),
|
|
]
|
|
}
|
|
checkpoint["channel_versions"] = {"messages": 1}
|
|
checkpointer.put(
|
|
{"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}},
|
|
checkpoint,
|
|
{"step": 1},
|
|
{"messages": 1},
|
|
)
|
|
await write_thread_goal(checkpointer, thread_id, build_goal_state(goal_text, max_continuations=2))
|
|
|
|
|
|
async def _write_messages(checkpointer: InMemorySaver, *, thread_id: str, messages: list) -> None:
|
|
checkpoint_tuple = await checkpointer.aget_tuple({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}})
|
|
assert checkpoint_tuple is not None
|
|
checkpoint = copy.deepcopy(getattr(checkpoint_tuple, "checkpoint", {}) or {})
|
|
metadata = copy.deepcopy(getattr(checkpoint_tuple, "metadata", {}) or {})
|
|
channel_values = dict(checkpoint.get("channel_values", {}) or {})
|
|
channel_values["messages"] = messages
|
|
checkpoint["channel_values"] = channel_values
|
|
channel_versions = dict(checkpoint.get("channel_versions", {}) or {})
|
|
current_version = channel_versions.get("messages")
|
|
channel_versions["messages"] = checkpointer.get_next_version(current_version, None)
|
|
checkpoint["channel_versions"] = channel_versions
|
|
checkpoint["id"] = str(uuid6())
|
|
metadata["step"] = metadata.get("step", 0) + 1
|
|
metadata["writes"] = {"test": {"messages": messages}}
|
|
await checkpointer.aput(
|
|
{"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}},
|
|
checkpoint,
|
|
metadata,
|
|
{"messages": channel_versions["messages"]},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_goal_worker_returns_hidden_continuation_when_goal_is_unmet(monkeypatch):
|
|
checkpointer = InMemorySaver()
|
|
thread_id = "goal-thread"
|
|
await _seed_goal_thread(checkpointer, thread_id=thread_id, goal_text="Finish all tests")
|
|
bridge = _CollectingBridge()
|
|
task_store = ExtensionData("run-1")
|
|
extensions = ExtensionRegistry().build()
|
|
|
|
async def fake_evaluate_goal_completion(goal, messages, **kwargs):
|
|
assert goal["objective"] == "Finish all tests"
|
|
assert [message.content for message in messages][-1] == "I made a start, but I am not done."
|
|
assert kwargs["task_store"] is task_store
|
|
assert kwargs["extensions"] is extensions
|
|
return GoalEvaluation(
|
|
satisfied=False,
|
|
blocker="goal_not_met_yet",
|
|
reason="Tests have not passed yet.",
|
|
evidence_summary="Implementation is incomplete.",
|
|
)
|
|
|
|
monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion)
|
|
|
|
continuation = await worker._prepare_goal_continuation_input(
|
|
accessor=_full_accessor(checkpointer),
|
|
bridge=bridge,
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id="run-1",
|
|
model_name="test-model",
|
|
app_config=None,
|
|
task_store=task_store,
|
|
extensions=extensions,
|
|
)
|
|
|
|
assert continuation is not None
|
|
[message] = continuation["messages"]
|
|
assert message.additional_kwargs["hide_from_ui"] is True
|
|
assert "Finish all tests" in message.content
|
|
assert "Tests have not passed yet." in message.content
|
|
latest_goal = await read_thread_goal(checkpointer, thread_id)
|
|
assert latest_goal is not None
|
|
assert latest_goal["continuation_count"] == 1
|
|
assert latest_goal["last_evaluation"]["run_id"] == "run-1"
|
|
assert latest_goal["last_evaluation"]["blocker"] == "goal_not_met_yet"
|
|
assert "stand_down_reason" not in latest_goal["last_evaluation"]
|
|
assert bridge.events[0][0] == "values"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_goal_worker_clears_goal_when_evaluator_is_satisfied(monkeypatch):
|
|
checkpointer = InMemorySaver()
|
|
thread_id = "done-goal-thread"
|
|
await _seed_goal_thread(checkpointer, thread_id=thread_id, goal_text="Finish all tests")
|
|
bridge = _CollectingBridge()
|
|
|
|
async def fake_evaluate_goal_completion(_goal, _messages, **_kwargs):
|
|
return GoalEvaluation(
|
|
satisfied=True,
|
|
blocker="none",
|
|
reason="The visible conversation says the task is done.",
|
|
evidence_summary="Done.",
|
|
)
|
|
|
|
monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion)
|
|
|
|
continuation = await worker._prepare_goal_continuation_input(
|
|
accessor=_full_accessor(checkpointer),
|
|
bridge=bridge,
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id="run-2",
|
|
model_name="test-model",
|
|
app_config=None,
|
|
)
|
|
|
|
assert continuation is None
|
|
assert await read_thread_goal(checkpointer, thread_id) is None
|
|
assert bridge.events[0][0] == "values"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_goal_worker_evaluates_materialized_messages_in_delta_mode(monkeypatch):
|
|
"""Delta checkpoints store no ``channel_values.messages``; the goal flow must
|
|
read messages through the mode-matched accessor or it sees an empty list,
|
|
loses the durable-receipt check, and stands down every continuation.
|
|
"""
|
|
checkpointer = InMemorySaver()
|
|
thread_id = "delta-goal-thread"
|
|
accessor = CheckpointStateAccessor.bind(build_state_mutation_graph("goal_evaluator", "delta"), checkpointer, mode="delta")
|
|
await accessor.aupdate(
|
|
{"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}},
|
|
{
|
|
"messages": [
|
|
HumanMessage(content="Please finish this task."),
|
|
AIMessage(content="I made a start, but I am not done."),
|
|
]
|
|
},
|
|
as_node="goal_evaluator",
|
|
)
|
|
await write_thread_goal(checkpointer, thread_id, build_goal_state("Finish all tests", max_continuations=2))
|
|
bridge = _CollectingBridge()
|
|
seen: dict[str, list] = {}
|
|
|
|
async def fake_evaluate_goal_completion(_goal, messages, **_kwargs):
|
|
seen["messages"] = list(messages)
|
|
return GoalEvaluation(
|
|
satisfied=False,
|
|
blocker="goal_not_met_yet",
|
|
reason="Tests have not passed yet.",
|
|
evidence_summary="Implementation is incomplete.",
|
|
)
|
|
|
|
monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion)
|
|
|
|
continuation = await worker._prepare_goal_continuation_input(accessor=accessor, bridge=bridge, checkpointer=checkpointer, thread_id=thread_id, run_id="run-delta", model_name="test-model", app_config=None)
|
|
|
|
assert continuation is not None
|
|
assert [message.content for message in seen["messages"]] == [
|
|
"Please finish this task.",
|
|
"I made a start, but I am not done.",
|
|
]
|
|
latest_goal = await read_thread_goal(checkpointer, thread_id)
|
|
assert latest_goal is not None
|
|
assert latest_goal["continuation_count"] == 1
|
|
assert "stand_down_reason" not in latest_goal["last_evaluation"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_goal_worker_stands_down_for_non_continuable_blocker(monkeypatch):
|
|
checkpointer = InMemorySaver()
|
|
thread_id = "blocked-goal-thread"
|
|
await _seed_goal_thread(checkpointer, thread_id=thread_id, goal_text="Finish all tests")
|
|
bridge = _CollectingBridge()
|
|
|
|
async def fake_evaluate_goal_completion(_goal, _messages, **_kwargs):
|
|
return GoalEvaluation(
|
|
satisfied=False,
|
|
blocker="missing_evidence",
|
|
reason="The transcript does not prove any verification.",
|
|
evidence_summary="No test result is visible.",
|
|
)
|
|
|
|
monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion)
|
|
|
|
continuation = await worker._prepare_goal_continuation_input(
|
|
accessor=_full_accessor(checkpointer),
|
|
bridge=bridge,
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id="run-3",
|
|
model_name="test-model",
|
|
app_config=None,
|
|
)
|
|
|
|
assert continuation is None
|
|
latest_goal = await read_thread_goal(checkpointer, thread_id)
|
|
assert latest_goal is not None
|
|
assert latest_goal["continuation_count"] == 0
|
|
assert latest_goal["last_evaluation"]["blocker"] == "missing_evidence"
|
|
assert latest_goal["last_evaluation"]["stand_down_reason"] == "blocked:missing_evidence"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_goal_worker_stands_down_when_no_progress_repeats(monkeypatch):
|
|
checkpointer = InMemorySaver()
|
|
thread_id = "no-progress-goal-thread"
|
|
messages = [HumanMessage(content="Please finish this task."), AIMessage(content="I made a start, but I am not done.")]
|
|
await _seed_goal_thread(checkpointer, thread_id=thread_id, goal_text="Finish all tests", messages=messages)
|
|
previous_goal = await read_thread_goal(checkpointer, thread_id)
|
|
assert previous_goal is not None
|
|
repeated_evaluation = GoalEvaluation(
|
|
satisfied=False,
|
|
blocker="goal_not_met_yet",
|
|
reason="The same work remains.",
|
|
evidence_summary="No new verification evidence.",
|
|
)
|
|
# Seed the prior evaluation with the SAME visible assistant evidence the worker
|
|
# will recompute, so the no-progress breaker recognises the stalled turn even
|
|
# though the evaluator may reword its free-text reason.
|
|
evidence_signature = latest_visible_assistant_signature(messages)
|
|
await write_thread_goal(
|
|
checkpointer,
|
|
thread_id,
|
|
attach_goal_evaluation(previous_goal, repeated_evaluation, run_id="previous-run", no_progress_count=1, evidence_signature=evidence_signature),
|
|
)
|
|
bridge = _CollectingBridge()
|
|
|
|
async def fake_evaluate_goal_completion(_goal, _messages, **_kwargs):
|
|
return repeated_evaluation
|
|
|
|
monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion)
|
|
|
|
continuation = await worker._prepare_goal_continuation_input(
|
|
accessor=_full_accessor(checkpointer),
|
|
bridge=bridge,
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id="run-4",
|
|
model_name="test-model",
|
|
app_config=None,
|
|
)
|
|
|
|
assert continuation is None
|
|
latest_goal = await read_thread_goal(checkpointer, thread_id)
|
|
assert latest_goal is not None
|
|
assert latest_goal["no_progress_count"] == 2
|
|
assert latest_goal["last_evaluation"]["stand_down_reason"] == "no_progress_detected"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_goal_worker_does_not_resurrect_goal_cleared_during_evaluation(monkeypatch):
|
|
checkpointer = InMemorySaver()
|
|
thread_id = "clear-during-eval-thread"
|
|
await _seed_goal_thread(checkpointer, thread_id=thread_id, goal_text="Finish all tests")
|
|
bridge = _CollectingBridge()
|
|
|
|
async def fake_evaluate_goal_completion(_goal, _messages, **_kwargs):
|
|
await write_thread_goal(checkpointer, thread_id, None, as_node="test")
|
|
return GoalEvaluation(
|
|
satisfied=False,
|
|
blocker="goal_not_met_yet",
|
|
reason="More work remains.",
|
|
evidence_summary="Work remains.",
|
|
)
|
|
|
|
monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion)
|
|
|
|
continuation = await worker._prepare_goal_continuation_input(
|
|
accessor=_full_accessor(checkpointer),
|
|
bridge=bridge,
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id="run-5",
|
|
model_name="test-model",
|
|
app_config=None,
|
|
)
|
|
|
|
assert continuation is None
|
|
assert await read_thread_goal(checkpointer, thread_id) is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_goal_worker_does_not_resurrect_goal_cleared_during_persist():
|
|
checkpointer = InMemorySaver()
|
|
thread_id = "clear-during-persist-thread"
|
|
await _seed_goal_thread(checkpointer, thread_id=thread_id, goal_text="Finish all tests")
|
|
existing_goal = await read_thread_goal(checkpointer, thread_id)
|
|
assert existing_goal is not None
|
|
wrapped_checkpointer = _ClearBeforeSecondGoalReadCheckpointer(checkpointer, thread_id)
|
|
bridge = _CollectingBridge()
|
|
|
|
result = await worker._persist_goal_evaluation(
|
|
bridge=bridge,
|
|
checkpointer=wrapped_checkpointer,
|
|
thread_id=thread_id,
|
|
run_id="run-clear-during-persist",
|
|
goal=existing_goal,
|
|
evaluation=GoalEvaluation(
|
|
satisfied=False,
|
|
blocker="goal_not_met_yet",
|
|
reason="More work remains.",
|
|
evidence_summary="Work remains.",
|
|
),
|
|
no_progress_count=0,
|
|
)
|
|
|
|
assert result is None
|
|
assert wrapped_checkpointer.cleared is True
|
|
assert await read_thread_goal(checkpointer, thread_id) is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_goal_worker_stops_when_abort_is_requested_during_evaluation(monkeypatch):
|
|
checkpointer = InMemorySaver()
|
|
thread_id = "abort-during-eval-thread"
|
|
await _seed_goal_thread(checkpointer, thread_id=thread_id, goal_text="Finish all tests")
|
|
bridge = _CollectingBridge()
|
|
abort_event = asyncio.Event()
|
|
|
|
async def fake_evaluate_goal_completion(_goal, _messages, **_kwargs):
|
|
abort_event.set()
|
|
return GoalEvaluation(
|
|
satisfied=False,
|
|
blocker="goal_not_met_yet",
|
|
reason="More work remains.",
|
|
evidence_summary="Work remains.",
|
|
)
|
|
|
|
monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion)
|
|
|
|
continuation = await worker._prepare_goal_continuation_input(
|
|
accessor=_full_accessor(checkpointer),
|
|
bridge=bridge,
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id="run-abort",
|
|
model_name="test-model",
|
|
app_config=None,
|
|
abort_event=abort_event,
|
|
)
|
|
|
|
assert continuation is None
|
|
latest_goal = await read_thread_goal(checkpointer, thread_id)
|
|
assert latest_goal is not None
|
|
assert latest_goal["continuation_count"] == 0
|
|
assert "last_evaluation" not in latest_goal
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_goal_worker_stands_down_when_thread_changes_after_evaluation(monkeypatch):
|
|
checkpointer = InMemorySaver()
|
|
thread_id = "user-wins-thread"
|
|
await _seed_goal_thread(checkpointer, thread_id=thread_id, goal_text="Finish all tests")
|
|
bridge = _CollectingBridge()
|
|
|
|
async def fake_evaluate_goal_completion(_goal, messages, **_kwargs):
|
|
await _write_messages(
|
|
checkpointer,
|
|
thread_id=thread_id,
|
|
messages=[*messages, HumanMessage(content="Actually, stop and wait.")],
|
|
)
|
|
return GoalEvaluation(
|
|
satisfied=False,
|
|
blocker="goal_not_met_yet",
|
|
reason="More work remains.",
|
|
evidence_summary="Work remains.",
|
|
)
|
|
|
|
monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion)
|
|
|
|
continuation = await worker._prepare_goal_continuation_input(
|
|
accessor=_full_accessor(checkpointer),
|
|
bridge=bridge,
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id="run-6",
|
|
model_name="test-model",
|
|
app_config=None,
|
|
)
|
|
|
|
assert continuation is None
|
|
latest_goal = await read_thread_goal(checkpointer, thread_id)
|
|
assert latest_goal is not None
|
|
assert latest_goal["continuation_count"] == 0
|
|
assert latest_goal["last_evaluation"]["stand_down_reason"] == "thread_changed_after_evaluation"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_goal_worker_stands_down_when_thread_changes_before_continuation(monkeypatch):
|
|
"""A user message racing in right after the continuation commits must not
|
|
double-bump continuation_count.
|
|
|
|
Sibling scenario to ``..._after_evaluation`` above, but the race lands
|
|
later: after the evaluator runs and after _prepare_goal_continuation_input
|
|
commits the real continuation (``_persist(..., continuation_count=next_count)``),
|
|
a racing visible message arrives before the function's trailing re-check.
|
|
That re-check detects the changed thread and stands down via a second
|
|
``_persist(..., continuation_count=next_count, stand_down_reason=...)``
|
|
call using the *same* next_count as the first, already-successful call.
|
|
|
|
Without the fix, that second call re-triggers PR #4088's
|
|
max(continuation_count, current_count + 1) guard against its own sibling
|
|
call's prior write (current_count is already next_count from the first
|
|
call), bumping continuation_count to next_count + 1 a second time --
|
|
consuming 2 units of the continuation budget for a cycle that delivered
|
|
zero actual continuations. The fix must leave it at next_count (1).
|
|
"""
|
|
inner = InMemorySaver()
|
|
thread_id = "race-before-continuation-thread"
|
|
await _seed_goal_thread(inner, thread_id=thread_id, goal_text="Finish all tests")
|
|
checkpointer = _RaceAfterFirstContinuationCommitCheckpointer(inner, thread_id)
|
|
bridge = _CollectingBridge()
|
|
|
|
async def fake_evaluate_goal_completion(_goal, _messages, **_kwargs):
|
|
return GoalEvaluation(
|
|
satisfied=False,
|
|
blocker="goal_not_met_yet",
|
|
reason="More work remains.",
|
|
evidence_summary="Work remains.",
|
|
)
|
|
|
|
monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion)
|
|
|
|
continuation = await worker._prepare_goal_continuation_input(
|
|
accessor=_full_accessor(checkpointer),
|
|
bridge=bridge,
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id="run-race-before-continuation",
|
|
model_name="test-model",
|
|
app_config=None,
|
|
)
|
|
|
|
assert continuation is None
|
|
latest_goal = await read_thread_goal(inner, thread_id)
|
|
assert latest_goal is not None
|
|
# Without the fix this is 2 (double-bumped). It must be 1: one real
|
|
# continuation attempt was committed and then stood down, not two.
|
|
assert latest_goal["continuation_count"] == 1
|
|
assert latest_goal["last_evaluation"]["stand_down_reason"] == "thread_changed_before_continuation"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_goal_worker_stands_down_without_durable_assistant_receipt():
|
|
checkpointer = InMemorySaver()
|
|
thread_id = "no-receipt-thread"
|
|
await _seed_goal_thread(
|
|
checkpointer,
|
|
thread_id=thread_id,
|
|
goal_text="Finish all tests",
|
|
messages=[HumanMessage(content="Please finish this task.")],
|
|
)
|
|
bridge = _CollectingBridge()
|
|
|
|
continuation = await worker._prepare_goal_continuation_input(
|
|
accessor=_full_accessor(checkpointer),
|
|
bridge=bridge,
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id="run-7",
|
|
model_name="test-model",
|
|
app_config=None,
|
|
)
|
|
|
|
assert continuation is None
|
|
latest_goal = await read_thread_goal(checkpointer, thread_id)
|
|
assert latest_goal is not None
|
|
assert latest_goal["last_evaluation"]["blocker"] == "run_failed"
|
|
assert latest_goal["last_evaluation"]["stand_down_reason"] == "no_durable_end_of_turn"
|
|
|
|
|
|
def test_stand_down_reason_uses_documented_default_caps_when_missing():
|
|
"""_stand_down_reason must fall back to the same default caps as
|
|
should_continue_goal (8 / 2). A bare goal dict missing the cap fields must
|
|
not be reported as 'max reached' / 'no progress' when it has not actually
|
|
exhausted the documented defaults.
|
|
"""
|
|
bare_goal = {"objective": "x", "status": "active", "continuation_count": 0}
|
|
unmet = GoalEvaluation(satisfied=False, blocker="goal_not_met_yet", reason="", evidence_summary="")
|
|
|
|
assert worker._stand_down_reason(bare_goal, unmet, no_progress_count=0) is None
|
|
# And the two gate functions agree on the same bare goal.
|
|
from deerflow.runtime.goal import should_continue_goal
|
|
|
|
assert should_continue_goal(bare_goal, unmet, no_progress_count=0) is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_agent_does_not_stream_continuation_after_abort(monkeypatch):
|
|
class FakeAgent:
|
|
def __init__(self) -> None:
|
|
self.inputs = []
|
|
self.metadata = {}
|
|
self.checkpointer = None
|
|
self.store = None
|
|
self.interrupt_before_nodes = []
|
|
self.interrupt_after_nodes = []
|
|
|
|
def astream(self, input_payload, **_kwargs):
|
|
self.inputs.append(input_payload)
|
|
|
|
async def _gen():
|
|
yield {"messages": []}
|
|
|
|
return _gen()
|
|
|
|
class FakeRunManager:
|
|
async def try_start(self, _run_id):
|
|
record.status = RunStatus.running
|
|
return RunStartOutcome.started
|
|
|
|
async def set_status(self, _run_id, status, **_kwargs):
|
|
record.status = status
|
|
|
|
async def set_status_if_not_cancelled(self, _run_id, status, **kwargs):
|
|
await self.set_status(_run_id, status, **kwargs)
|
|
return None
|
|
|
|
async def update_model_name(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def update_run_completion(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def wait_for_prior_finalizing(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def set_finalizing(self, _run_id, finalizing):
|
|
record.finalizing = finalizing
|
|
|
|
class FakeBridge:
|
|
async def publish(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def publish_end(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def cleanup(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def fake_prepare(**kwargs):
|
|
kwargs["abort_event"].set()
|
|
return {"messages": [HumanMessage(content="continue", additional_kwargs={"hide_from_ui": True})]}
|
|
|
|
monkeypatch.setattr(worker, "_prepare_goal_continuation_input", fake_prepare)
|
|
|
|
fake_agent = FakeAgent()
|
|
record = RunRecord(
|
|
run_id="run-abort-loop",
|
|
thread_id="thread-abort-loop",
|
|
assistant_id="lead-agent",
|
|
status=RunStatus.pending,
|
|
on_disconnect=DisconnectMode.cancel,
|
|
model_name="test-model",
|
|
)
|
|
record.abort_event = asyncio.Event()
|
|
|
|
await worker.run_agent(
|
|
FakeBridge(),
|
|
FakeRunManager(),
|
|
record,
|
|
ctx=worker.RunContext(checkpointer=None),
|
|
agent_factory=lambda config: fake_agent,
|
|
graph_input={"messages": [HumanMessage(content="start")]},
|
|
config={"configurable": {"thread_id": "thread-abort-loop"}},
|
|
)
|
|
|
|
assert len(fake_agent.inputs) == 1
|
|
assert fake_agent.inputs[0] == {"messages": [HumanMessage(content="start")]}
|
|
assert record.status == RunStatus.interrupted
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_agent_reuses_goal_evaluator_model_for_goal_loop(monkeypatch):
|
|
class FakeAgent:
|
|
def __init__(self) -> None:
|
|
self.inputs = []
|
|
self.metadata = {}
|
|
self.checkpointer = None
|
|
self.store = None
|
|
self.interrupt_before_nodes = []
|
|
self.interrupt_after_nodes = []
|
|
|
|
def astream(self, input_payload, **_kwargs):
|
|
self.inputs.append(input_payload)
|
|
|
|
async def _gen():
|
|
yield {"messages": []}
|
|
|
|
return _gen()
|
|
|
|
class FakeRunManager:
|
|
async def try_start(self, _run_id):
|
|
record.status = RunStatus.running
|
|
return RunStartOutcome.started
|
|
|
|
async def set_status(self, _run_id, status, **_kwargs):
|
|
record.status = status
|
|
|
|
async def set_status_if_not_cancelled(self, _run_id, status, **kwargs):
|
|
await self.set_status(_run_id, status, **kwargs)
|
|
return None
|
|
|
|
async def update_model_name(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def update_run_completion(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def wait_for_prior_finalizing(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def set_finalizing(self, _run_id, finalizing):
|
|
record.finalizing = finalizing
|
|
|
|
class FakeBridge:
|
|
async def publish(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def publish_end(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def cleanup(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
evaluator_model = object()
|
|
create_calls = []
|
|
|
|
def fake_create_goal_evaluator_model(**kwargs):
|
|
create_calls.append(kwargs)
|
|
return evaluator_model
|
|
|
|
prepare_models = []
|
|
|
|
async def fake_prepare(**kwargs):
|
|
prepare_models.append(kwargs["evaluator_model_factory"]())
|
|
if len(prepare_models) == 1:
|
|
return {"messages": [HumanMessage(content="continue", additional_kwargs={"hide_from_ui": True})]}
|
|
return None
|
|
|
|
monkeypatch.setattr(worker, "create_goal_evaluator_model", fake_create_goal_evaluator_model)
|
|
monkeypatch.setattr(worker, "_prepare_goal_continuation_input", fake_prepare)
|
|
|
|
fake_agent = FakeAgent()
|
|
record = RunRecord(
|
|
run_id="run-model-cache",
|
|
thread_id="thread-model-cache",
|
|
assistant_id="lead-agent",
|
|
status=RunStatus.pending,
|
|
on_disconnect=DisconnectMode.cancel,
|
|
model_name="test-model",
|
|
)
|
|
record.abort_event = asyncio.Event()
|
|
|
|
await worker.run_agent(
|
|
FakeBridge(),
|
|
FakeRunManager(),
|
|
record,
|
|
ctx=worker.RunContext(checkpointer=None, app_config=object()),
|
|
agent_factory=lambda config: fake_agent,
|
|
graph_input={"messages": [HumanMessage(content="start")]},
|
|
config={"configurable": {"thread_id": "thread-model-cache"}},
|
|
)
|
|
|
|
assert len(fake_agent.inputs) == 2
|
|
assert prepare_models == [evaluator_model, evaluator_model]
|
|
assert len(create_calls) == 1
|
|
assert create_calls[0]["model_name"] == "test-model"
|
|
assert record.status == RunStatus.success
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persist_goal_evaluation_does_not_regress_continuation_count_on_race():
|
|
"""A racing continuation must not overwrite a higher count with a lower one.
|
|
|
|
Scenario: two goal continuations run concurrently. Continuation A reads
|
|
continuation_count=1, computes next=2. Continuation B reads the same
|
|
count=1, computes next=2, but acquires the lock first and writes count=2.
|
|
When A acquires the lock, the current_goal already has count=2. Without
|
|
the defensive guard, A would write count=2 again (stale computation),
|
|
effectively losing one continuation event. The guard must compute
|
|
``max(stale_next, current_count + 1)`` so A writes count=3.
|
|
"""
|
|
checkpointer = InMemorySaver()
|
|
thread_id = "race-count-thread"
|
|
await _seed_goal_thread(checkpointer, thread_id=thread_id, goal_text="Race test")
|
|
# Simulate a racing continuation: bump the persisted continuation_count to 2
|
|
# before calling _persist_goal_evaluation with a next_count computed from
|
|
# stale state (count=1 → next=2).
|
|
existing_goal = await read_thread_goal(checkpointer, thread_id)
|
|
assert existing_goal is not None
|
|
bumped_goal = attach_goal_evaluation(
|
|
existing_goal,
|
|
GoalEvaluation(satisfied=False, blocker="goal_not_met_yet", reason="racing", evidence_summary=""),
|
|
run_id="racing-run",
|
|
continuation_count=2, # racing continuation already bumped to 2
|
|
)
|
|
await write_thread_goal(checkpointer, thread_id, bumped_goal)
|
|
|
|
# Now call _persist_goal_evaluation with continuation_count=2 computed from
|
|
# stale state (old count was 1). The guard should detect current_count=2
|
|
# and write max(2, 2+1) = 3.
|
|
bridge = _CollectingBridge()
|
|
result = await worker._persist_goal_evaluation(
|
|
bridge=bridge,
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id="run-late",
|
|
goal=existing_goal, # stale goal with continuation_count=1
|
|
evaluation=GoalEvaluation(
|
|
satisfied=False,
|
|
blocker="goal_not_met_yet",
|
|
reason="More work remains.",
|
|
evidence_summary="Work remains.",
|
|
),
|
|
no_progress_count=1,
|
|
continuation_count=2, # computed from stale state: stale_count(1) + 1
|
|
)
|
|
|
|
assert result is not None
|
|
# Without the guard this would be 2 (stale computation wins). With the
|
|
# guard it must be 3 (current_count + 1 taken inside the lock).
|
|
assert result["continuation_count"] == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persist_goal_evaluation_no_race_uses_caller_count():
|
|
"""When no racing continuation exists, the caller's continuation_count is used."""
|
|
checkpointer = InMemorySaver()
|
|
thread_id = "no-race-thread"
|
|
await _seed_goal_thread(checkpointer, thread_id=thread_id, goal_text="No race test")
|
|
existing_goal = await read_thread_goal(checkpointer, thread_id)
|
|
assert existing_goal is not None
|
|
|
|
bridge = _CollectingBridge()
|
|
result = await worker._persist_goal_evaluation(
|
|
bridge=bridge,
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id="run-normal",
|
|
goal=existing_goal,
|
|
evaluation=GoalEvaluation(
|
|
satisfied=False,
|
|
blocker="goal_not_met_yet",
|
|
reason="More work.",
|
|
evidence_summary="Work.",
|
|
),
|
|
no_progress_count=1,
|
|
continuation_count=1, # 0 + 1 = 1
|
|
)
|
|
|
|
assert result is not None
|
|
assert result["continuation_count"] == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_agent_strips_branch_checkpoint_for_goal_continuation(monkeypatch):
|
|
class FakeAgent:
|
|
def __init__(self) -> None:
|
|
self.calls = []
|
|
self.metadata = {}
|
|
self.checkpointer = None
|
|
self.store = None
|
|
self.interrupt_before_nodes = []
|
|
self.interrupt_after_nodes = []
|
|
|
|
def astream(self, input_payload, **kwargs):
|
|
configurable = dict(kwargs["config"].get("configurable", {}))
|
|
self.calls.append((input_payload, configurable))
|
|
|
|
async def _gen():
|
|
yield {"messages": []}
|
|
|
|
return _gen()
|
|
|
|
class FakeRunManager:
|
|
async def try_start(self, _run_id):
|
|
record.status = RunStatus.running
|
|
return RunStartOutcome.started
|
|
|
|
async def set_status(self, _run_id, status, **_kwargs):
|
|
record.status = status
|
|
|
|
async def set_status_if_not_cancelled(self, _run_id, status, **kwargs):
|
|
await self.set_status(_run_id, status, **kwargs)
|
|
return None
|
|
|
|
async def update_model_name(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def update_run_completion(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def wait_for_prior_finalizing(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def set_finalizing(self, _run_id, finalizing):
|
|
record.finalizing = finalizing
|
|
|
|
class FakeBridge:
|
|
async def publish(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def publish_end(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def cleanup(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
async def fake_prepare(**_kwargs):
|
|
if len(fake_agent.calls) == 1:
|
|
return {"messages": [HumanMessage(content="continue", additional_kwargs={"hide_from_ui": True})]}
|
|
return None
|
|
|
|
monkeypatch.setattr(worker, "_prepare_goal_continuation_input", fake_prepare)
|
|
|
|
fake_agent = FakeAgent()
|
|
record = RunRecord(
|
|
run_id="run-branch-continuation",
|
|
thread_id="thread-branch-continuation",
|
|
assistant_id="lead-agent",
|
|
status=RunStatus.pending,
|
|
on_disconnect=DisconnectMode.cancel,
|
|
model_name="test-model",
|
|
)
|
|
record.abort_event = asyncio.Event()
|
|
|
|
await worker.run_agent(
|
|
FakeBridge(),
|
|
FakeRunManager(),
|
|
record,
|
|
ctx=worker.RunContext(checkpointer=None),
|
|
agent_factory=lambda config: fake_agent,
|
|
graph_input={"messages": [HumanMessage(content="start")]},
|
|
config={
|
|
"configurable": {
|
|
"thread_id": "thread-branch-continuation",
|
|
"checkpoint_ns": "branch",
|
|
"checkpoint_id": "old-checkpoint",
|
|
"checkpoint_map": {"": "old-checkpoint"},
|
|
}
|
|
},
|
|
)
|
|
|
|
assert len(fake_agent.calls) == 2
|
|
first_config = fake_agent.calls[0][1]
|
|
second_config = fake_agent.calls[1][1]
|
|
assert first_config["checkpoint_ns"] == "branch"
|
|
assert first_config["checkpoint_id"] == "old-checkpoint"
|
|
assert first_config["checkpoint_map"] == {"": "old-checkpoint"}
|
|
assert second_config["checkpoint_ns"] == ""
|
|
assert "checkpoint_id" not in second_config
|
|
assert "checkpoint_map" not in second_config
|
|
assert second_config["thread_id"] == "thread-branch-continuation"
|