mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-11-08 03:30:34 +00:00
593ba49458
* feat(subagents): continue native batches and honor item dependencies * docs: link native batch follow-up changelog entries * test(migrations): pin isolated history-seed upgrade to 0050 * fix(subagents): preserve follow-up scope, blockers and fair scheduling * test(memory): freeze capacity evaluation time * fix(subagents): guard follow-up admission and retain run controls * fix(subagents): enforce runs:cancel on tool cancellation of follow-up batches * style(tests): format test_batch_followup_host with ruff * fix(subagents): drop the admission checkpoint check and keep SDK cancel compatible The checkpoint check inside run admission waited on the checkpoint saver's lock while holding the SQLite write lock, which could make another chat's checkpoint write fail. Remove it; dispatch still checks the conversation before admission. A follow-up waiting for an unanswered question now keeps its claim until the lease expires instead of rereading the conversation every poll. Cancel passes the caller's run authority only to submitters that take it, so existing SDK submitters keep the two-argument call. Document that retrying a prerequisite by hand does not requeue its blocked items, and list 0052 as the migration head. --------- Co-authored-by: Totoro-qaq <279883115+Totoro-qaq@users.noreply.github.com> Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
413 lines
20 KiB
Python
413 lines
20 KiB
Python
"""Real AgentRuns admission and SQL fencing for native batch follow-ups."""
|
|
|
|
import asyncio
|
|
from dataclasses import replace
|
|
from datetime import UTC, datetime
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
from langchain_core.messages import AIMessage
|
|
from langgraph.graph import MessagesState, StateGraph
|
|
from test_batch_dependencies import repo as _repo_fixture
|
|
from test_extension_agent_runs import host as _host_fixture
|
|
from test_extension_agent_runs import runtime as _runtime_fixture
|
|
|
|
host = _host_fixture
|
|
runtime = _runtime_fixture
|
|
repo = _repo_fixture
|
|
|
|
|
|
def graph_factory(factory, answer, schema=MessagesState):
|
|
def build(*, config):
|
|
graph = StateGraph(schema)
|
|
graph.add_node("answer", answer)
|
|
graph.set_entry_point("answer")
|
|
graph.set_finish_point("answer")
|
|
return replace(factory(config=config), graph=graph.compile())
|
|
|
|
return build
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def batch_host(runtime, repo, monkeypatch, request):
|
|
from app.gateway import services
|
|
from app.subagent_batches.continuation import BatchContinuations
|
|
from deerflow.config.subagent_batches_config import SubagentBatchesConfig
|
|
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
|
|
from deerflow.persistence.engine import get_session_factory
|
|
from deerflow.persistence.run.sql import RunRepository
|
|
from deerflow.runtime import RunManager
|
|
from deerflow.subagents.batch_runtime import BatchSubmitRequest
|
|
from deerflow.subagents.batch_service import SubagentBatchService
|
|
|
|
settings = getattr(request, "param", {})
|
|
runs, state, request, _, _ = runtime
|
|
sf = get_session_factory()
|
|
state.run_store = RunRepository(sf)
|
|
state.run_manager = RunManager(store=state.run_store)
|
|
state.subagent_batch_repo = repo
|
|
entered, release = asyncio.Event(), asyncio.Event()
|
|
factory = services.resolve_agent_factory("lead_agent")
|
|
|
|
async def answer(state):
|
|
text = state["messages"][-1].content
|
|
if text == "source":
|
|
entered.set()
|
|
await release.wait()
|
|
return {"messages": [AIMessage(content="finished:" + text)]}
|
|
|
|
monkeypatch.setattr(services, "resolve_agent_factory", lambda _: graph_factory(factory, answer))
|
|
thread = await runs.create_thread(thread_id="native-batch-followup")
|
|
source = await runs.start(thread_id=thread, input={"messages": [{"role": "user", "content": "source"}]})
|
|
await asyncio.wait_for(entered.wait(), 5)
|
|
controller = BatchContinuations(request.app)
|
|
owner = str(request.state.user.id)
|
|
thread_row = await state.thread_store.get(thread, user_id=owner)
|
|
submission = BatchSubmitRequest(
|
|
owner,
|
|
thread,
|
|
source.run_id,
|
|
"call",
|
|
"submission",
|
|
"Process records",
|
|
"general-purpose",
|
|
[{"key": "a", "prompt": "Read A"}],
|
|
1,
|
|
1,
|
|
{"thread_incarnation": thread_row["incarnation"]},
|
|
continuation="Write the requested summary",
|
|
continuation_authority=runs,
|
|
continuation_context={"goal": None, "context": {"mode": "flash", "locale": "en-US", **settings}, "config": {"recursion_limit": 99}},
|
|
)
|
|
service = SubagentBatchService(
|
|
repository=repo,
|
|
config=SubagentBatchesConfig(enabled=True, max_attempts=1),
|
|
runtime_config=SubagentRuntimeConfig(),
|
|
prepare_continuation=controller.prepare,
|
|
dispatch_continuation=controller.dispatch,
|
|
cancel_continuation=controller.cancel,
|
|
)
|
|
state.subagent_batch_service = service
|
|
await service.submit(submission)
|
|
release.set()
|
|
assert (await runs.wait(thread_id=thread, run_id=source.run_id, timeout=5)).status == "success"
|
|
now = datetime.now(UTC)
|
|
item = (await repo.claim_items(now=now, lease_owner="item-worker", lease_seconds=120, limit=1))[0]
|
|
await repo.finalize_item(
|
|
item["id"], lease_owner="item-worker", succeeded=True, result="Source read", result_preview="Source read", result_truncated=False, error=None, stop_reason=None, token_usage=None, model_name="fixture", completed_at=now
|
|
)
|
|
claimed = (await repo.claim_continuations(now=datetime.now(UTC), lease_owner="continuation-worker", lease_seconds=120))[0]
|
|
try:
|
|
yield controller, claimed, repo, runs, state, request
|
|
finally:
|
|
release.set()
|
|
await state.run_manager.shutdown()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_followup_runs_once_and_reconciles_an_unknown_admission(batch_host):
|
|
controller, batch, repo, runs, state, _ = batch_host
|
|
first = await controller.dispatch(batch)
|
|
assert first["status"] == "dispatched"
|
|
assert (await runs.wait(thread_id=batch["thread_id"], run_id=first["run_id"], timeout=5)).status == "success"
|
|
# Simulate the batch worker losing the admission response before receipt persistence.
|
|
second = await controller.dispatch(batch)
|
|
assert second == first
|
|
record = await state.run_manager.get(first["run_id"])
|
|
assert record.kwargs["config"]["recursion_limit"] == 99
|
|
assert record.metadata["deerflow_origin"]["kind"] == "subagent_batch"
|
|
assert await repo.finish_continuation(batch["id"], lease_owner="continuation-worker", status="dispatched", run_id=first["run_id"])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("batch_host", [{"model_name": "parent", "model_selection_requested": "parent", "subagent_model_name": "child"}], indirect=True)
|
|
async def test_followup_preserves_explicit_model_controls(batch_host):
|
|
context = batch_host[1]["continuation"]["context"]
|
|
assert context["model_selection_requested"] == context["model_name"] == "parent"
|
|
assert context["subagent_model_name"] == "child"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("action", ["cancel", "supersede"])
|
|
async def test_cancellation_or_goal_change_fences_a_launch_already_in_flight(batch_host, monkeypatch, action):
|
|
from app.gateway import services
|
|
|
|
controller, batch, repo, _, state, _ = batch_host
|
|
reached, release = asyncio.Event(), asyncio.Event()
|
|
start = services.start_run
|
|
|
|
async def delayed(*args, **kwargs):
|
|
reached.set()
|
|
await release.wait()
|
|
return await start(*args, **kwargs)
|
|
|
|
monkeypatch.setattr(services, "start_run", delayed)
|
|
launch = asyncio.create_task(controller.dispatch(batch))
|
|
try:
|
|
await asyncio.wait_for(reached.wait(), 5)
|
|
if action == "cancel":
|
|
await repo.cancel_batch(batch["id"], user_id=batch["user_id"])
|
|
else:
|
|
from app.gateway.routers.threads import ThreadGoalRequest, set_thread_goal
|
|
from deerflow.runtime.user_context import reset_current_user, set_current_user
|
|
|
|
request = batch_host[-1]
|
|
token = set_current_user(request.state.user)
|
|
try:
|
|
await set_thread_goal(batch["thread_id"], ThreadGoalRequest(objective="The user's replacement goal"), request)
|
|
finally:
|
|
reset_current_user(token)
|
|
release.set()
|
|
assert await launch is None
|
|
assert await repo.continuation_run_id(batch["id"], user_id=batch["user_id"]) is None
|
|
finally:
|
|
release.set()
|
|
await asyncio.gather(launch, return_exceptions=True)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_revoked_original_account_cannot_launch_followup(batch_host):
|
|
controller, batch, repo, _, state, request = batch_host
|
|
state.agent_runs_host.load_user.return_value = request.state.user.model_copy(update={"token_version": request.state.user.token_version + 1})
|
|
assert await controller.dispatch(batch) == {"status": "suppressed"}
|
|
assert await repo.continuation_run_id(batch["id"], user_id=batch["user_id"]) is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"scope",
|
|
[
|
|
{"version": 1, "mode": "disabled"},
|
|
{"version": 1, "mode": "selected", "dataset_ids": ["approved"], "document_filters": [{"dataset_id": "approved", "document_ids": ["one"]}]},
|
|
],
|
|
)
|
|
async def test_resumed_native_batch_preserves_admitted_scope_in_followup(batch_host, monkeypatch, scope):
|
|
from types import SimpleNamespace
|
|
|
|
from langgraph.runtime import Runtime
|
|
from langgraph.types import interrupt
|
|
from test_batch_task_tool import tool_module
|
|
|
|
from app.gateway import services
|
|
from deerflow.config.app_config import AppConfig, set_app_config
|
|
from deerflow.config.subagent_batches_config import SubagentBatchesConfig
|
|
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
|
|
from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_KEY, KNOWLEDGE_SCOPE_RUNTIME_KEY
|
|
from deerflow.subagents.batch_service import SubagentBatchService
|
|
|
|
controller, _, repo, runs, state, request = batch_host
|
|
set_app_config(
|
|
AppConfig.model_validate(
|
|
{
|
|
"sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"},
|
|
"knowledge_base": {"enabled": True},
|
|
"tools": [{"name": "knowledge_search", "group": "knowledge", "use": "deerflow.community.ragflow.tools:knowledge_search_tool"}],
|
|
}
|
|
)
|
|
)
|
|
service = SubagentBatchService(repository=repo, config=SubagentBatchesConfig(enabled=True), runtime_config=SubagentRuntimeConfig(), prepare_continuation=controller.prepare, dispatch_continuation=controller.dispatch)
|
|
monkeypatch.setattr(tool_module, "get_subagent_batch_submitter", lambda: service)
|
|
observed = []
|
|
factory = services.resolve_agent_factory("lead_agent")
|
|
|
|
async def answer(values, runtime: Runtime, config):
|
|
if values["messages"][-1].content == "scoped work":
|
|
interrupt("Continue?")
|
|
record = await state.run_manager.get(runtime.context["run_id"])
|
|
assert record.kwargs["input"] is None
|
|
command = await tool_module.batch_task.coroutine(
|
|
title="Scoped batch",
|
|
items=[tool_module.BatchTaskItem(key="a", prompt="Read authorized records")],
|
|
subagent_type="general-purpose",
|
|
runtime=SimpleNamespace(context=runtime.context, state=values, config=config),
|
|
tool_call_id="scoped-call",
|
|
continuation="Present the result",
|
|
)
|
|
assert command.update["messages"][0].status != "error", command.update["messages"][0].content
|
|
observed.append(runtime.context.get(KNOWLEDGE_SCOPE_RUNTIME_KEY))
|
|
return {"messages": [AIMessage(content="Finished")]}
|
|
|
|
monkeypatch.setattr(services, "resolve_agent_factory", lambda _: graph_factory(factory, answer))
|
|
thread = await runs.create_thread(thread_id="scoped-resume")
|
|
first = await runs.start(thread_id=thread, input={"messages": [{"role": "user", "content": "scoped work", "additional_kwargs": {KNOWLEDGE_SCOPE_KEY: scope}}]})
|
|
assert (await runs.wait(thread_id=thread, run_id=first.run_id, timeout=5)).status == "success"
|
|
resumed = await runs.resume(thread_id=thread, resume="yes")
|
|
assert (await runs.wait(thread_id=thread, run_id=resumed.run_id, timeout=5)).status == "success"
|
|
assert observed == [scope]
|
|
now = datetime.now(UTC)
|
|
item = (await repo.claim_items(now=now, lease_owner="scoped-worker", lease_seconds=120, limit=1))[0]
|
|
await repo.finalize_item(item["id"], lease_owner="scoped-worker", succeeded=True, result="Done", result_preview="Done", result_truncated=False, error=None, stop_reason=None, token_usage=None, model_name="fixture", completed_at=now)
|
|
batch = (await repo.claim_continuations(now=datetime.now(UTC), lease_owner="scoped-followup", lease_seconds=120))[0]
|
|
assert batch["continuation"]["knowledge_scope"] == scope
|
|
launched = await controller.dispatch(batch)
|
|
assert (await runs.wait(thread_id=thread, run_id=launched["run_id"], timeout=5)).status == "success"
|
|
assert observed == [scope, scope]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_followup_waits_for_an_actual_answer_without_a_goal(batch_host, monkeypatch):
|
|
from langchain_core.messages import ToolMessage
|
|
|
|
from app.gateway import services
|
|
from deerflow.agents.human_input import extract_unanswered_human_input_request
|
|
|
|
controller, batch, repo, runs, state, _ = batch_host
|
|
factory = services.resolve_agent_factory("lead_agent")
|
|
question = {"kind": "human_input_request", "request_id": "recipient", "question": "Which recipient?", "source": "ask_clarification"}
|
|
|
|
async def answer(values):
|
|
if values["messages"][-1].content == "ask":
|
|
return {
|
|
"messages": [
|
|
AIMessage(content="", tool_calls=[{"name": "ask_clarification", "args": {}, "id": "ask"}]),
|
|
ToolMessage(content="Which recipient?", tool_call_id="ask", name="ask_clarification", artifact={"human_input": question}),
|
|
]
|
|
}
|
|
return {"messages": [AIMessage(content="Understood")]}
|
|
|
|
monkeypatch.setattr(services, "resolve_agent_factory", lambda _: graph_factory(factory, answer))
|
|
|
|
async def ask():
|
|
asked = await runs.start(thread_id=batch["thread_id"], input={"messages": [{"role": "user", "content": "ask"}]})
|
|
assert (await runs.wait(thread_id=batch["thread_id"], run_id=asked.run_id, timeout=5)).status == "success"
|
|
assert (await state.run_manager.get(asked.run_id)).human_input_request == question
|
|
|
|
await ask()
|
|
assert await controller.dispatch(batch) == {"status": "waiting"}
|
|
# Waiting keeps the claim until its lease expires instead of every poll.
|
|
await state.subagent_batch_service._continue_batch(batch, "continuation-worker")
|
|
assert await repo.claim_continuations(now=datetime.now(UTC), lease_owner="poller", lease_seconds=120) == []
|
|
assert await repo.continuation_run_id(batch["id"], user_id=batch["user_id"]) is None
|
|
pending = await runs.get_state(thread_id=batch["thread_id"])
|
|
assert extract_unanswered_human_input_request(pending["values"]["messages"]) == question
|
|
answered = await runs.start(thread_id=batch["thread_id"], input={"messages": [{"role": "user", "content": "Use the team inbox"}]})
|
|
assert (await runs.wait(thread_id=batch["thread_id"], run_id=answered.run_id, timeout=5)).status == "success"
|
|
launched = await controller.dispatch(batch)
|
|
assert (await runs.wait(thread_id=batch["thread_id"], run_id=launched["run_id"], timeout=5)).status == "success"
|
|
assert await controller.dispatch(batch) == launched
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(("status", "reason", "cancelled"), [("error", None, True), ("success", "token_capped", True), ("success", "loop_capped", True), ("success", "subagent_limit_capped", False)])
|
|
async def test_source_completion_preserves_existing_hard_stop_policy(batch_host, status, reason, cancelled):
|
|
from types import SimpleNamespace
|
|
|
|
from deerflow.config.subagent_batches_config import SubagentBatchesConfig
|
|
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
|
|
from deerflow.subagents.batch_service import SubagentBatchService
|
|
|
|
controller, batch, repo, _, state, _ = batch_host
|
|
state.subagent_batch_service = SubagentBatchService(repository=repo, config=SubagentBatchesConfig(), runtime_config=SubagentRuntimeConfig(), cancel_continuation=controller.cancel)
|
|
await controller.source_finished(SimpleNamespace(run_id=batch["run_id"], thread_id=batch["thread_id"], status=status, stop_reason=reason))
|
|
current = await repo.get_batch(batch["id"], user_id=batch["user_id"])
|
|
assert (current["continuation_status"] == "cancelled") is cancelled
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_worker_wakes_the_same_waiting_goal_and_accounts_once(batch_host, monkeypatch):
|
|
from app.gateway import services
|
|
from deerflow.agents.thread_state import ThreadState
|
|
from deerflow.runtime.goal import build_goal_state, write_thread_goal
|
|
|
|
controller, batch, _, runs, state, _ = batch_host
|
|
goal = build_goal_state("Finish the requested work", max_continuations=2)
|
|
goal["last_evaluation"] = {"blocker": "external_wait", "stand_down_reason": "blocked:external_wait"}
|
|
await write_thread_goal(state.checkpointer, batch["thread_id"], goal)
|
|
batch["continuation"]["goal"] = {key: goal[key] for key in ("created_at", "status", "objective")}
|
|
observed = []
|
|
factory = services.resolve_agent_factory("lead_agent")
|
|
|
|
async def answer(values):
|
|
observed.append(values["goal"]["continuation_count"])
|
|
return {"messages": [AIMessage(content="The requested work is finished.")]}
|
|
|
|
async def evaluated(*args, **kwargs):
|
|
return {"satisfied": True, "blocker": "none", "reason": "The controlled work finished"}
|
|
|
|
monkeypatch.setattr(services, "resolve_agent_factory", lambda _: graph_factory(factory, answer, ThreadState))
|
|
monkeypatch.setattr("deerflow.runtime.runs.worker.create_goal_evaluator_model", lambda **_: None)
|
|
monkeypatch.setattr("deerflow.runtime.runs.worker.evaluate_goal_completion", evaluated)
|
|
launched = await controller.dispatch(batch)
|
|
assert launched["status"] == "dispatched"
|
|
assert (await runs.wait(thread_id=batch["thread_id"], run_id=launched["run_id"], timeout=5)).status == "success"
|
|
result = await runs.get_state(thread_id=batch["thread_id"])
|
|
assert result["values"]["goal_outcome"]["status"] == "achieved"
|
|
assert result["values"]["goal_outcome"]["continuation_count"] == 1
|
|
assert observed == [1]
|
|
assert await controller.dispatch(batch) == launched
|
|
assert observed == [1]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_batch_tool_requires_runs_cancel_for_followups(batch_host, monkeypatch):
|
|
from types import SimpleNamespace
|
|
|
|
from deerflow_extension_api.agent_runs import AGENT_RUNS_CONTEXT_KEY, AgentRunError
|
|
from test_batch_task_tool import tool_module
|
|
|
|
from app.gateway import services
|
|
from deerflow.subagents.batch_runtime import BatchSubmitRequest
|
|
|
|
controller, batch, repo, runs, state, request = batch_host
|
|
monkeypatch.setattr(tool_module, "get_subagent_batch_submitter", lambda: state.subagent_batch_service)
|
|
request.state.auth.permissions = ["threads:write", "threads:read", "runs:create", "runs:read"]
|
|
no_cancel_runs = state.agent_runs_host.bind(request)
|
|
request.state.auth.permissions.append("runs:cancel")
|
|
|
|
def tool_runtime(authority):
|
|
context = {"user_id": batch["user_id"], "thread_id": batch["thread_id"]}
|
|
if authority is not None:
|
|
context[AGENT_RUNS_CONTEXT_KEY] = authority
|
|
return SimpleNamespace(context=context)
|
|
|
|
ordinary = await state.subagent_batch_service.submit(
|
|
BatchSubmitRequest(
|
|
batch["user_id"],
|
|
batch["thread_id"],
|
|
batch["run_id"],
|
|
"ordinary-call",
|
|
"ordinary-submission",
|
|
"Ordinary batch",
|
|
"general-purpose",
|
|
[{"key": "a", "prompt": "Read A"}],
|
|
1,
|
|
1,
|
|
{},
|
|
)
|
|
)
|
|
assert await tool_module.cancel_batch.coroutine(runtime=tool_runtime(no_cancel_runs), batch_id=ordinary["id"]) == (f"Batch {ordinary['id']} cancellation requested.")
|
|
|
|
for authority in (no_cancel_runs, None):
|
|
with pytest.raises(AgentRunError) as exc_info:
|
|
await tool_module.cancel_batch.coroutine(runtime=tool_runtime(authority), batch_id=batch["id"])
|
|
assert exc_info.value.status_code == 403
|
|
assert (await repo.get_batch(batch["id"], user_id=batch["user_id"]))["continuation_status"] == "dispatching"
|
|
|
|
entered, release = asyncio.Event(), asyncio.Event()
|
|
factory = services.resolve_agent_factory("lead_agent")
|
|
|
|
async def hold(_values):
|
|
entered.set()
|
|
await release.wait()
|
|
return {"messages": [AIMessage(content="finished followup")]}
|
|
|
|
monkeypatch.setattr(services, "resolve_agent_factory", lambda _: graph_factory(factory, hold))
|
|
try:
|
|
launched = await controller.dispatch(batch)
|
|
assert launched["status"] == "dispatched"
|
|
assert await repo.finish_continuation(batch["id"], lease_owner="continuation-worker", status="dispatched", run_id=launched["run_id"])
|
|
await asyncio.wait_for(entered.wait(), 5)
|
|
with pytest.raises(AgentRunError) as exc_info:
|
|
await tool_module.cancel_batch.coroutine(runtime=tool_runtime(no_cancel_runs), batch_id=batch["id"])
|
|
assert exc_info.value.status_code == 403
|
|
record = await state.run_manager.get(launched["run_id"])
|
|
assert not record.abort_event.is_set()
|
|
assert (await repo.get_batch(batch["id"], user_id=batch["user_id"]))["continuation_status"] == "dispatched"
|
|
|
|
result = await tool_module.cancel_batch.coroutine(runtime=tool_runtime(runs), batch_id=batch["id"])
|
|
assert result == f"Batch {batch['id']} cancellation requested."
|
|
assert (await runs.wait(thread_id=batch["thread_id"], run_id=launched["run_id"], timeout=5)).status == "interrupted"
|
|
assert (await repo.get_batch(batch["id"], user_id=batch["user_id"]))["continuation_status"] == "cancelled"
|
|
finally:
|
|
release.set()
|