diff --git a/CHANGELOG.md b/CHANGELOG.md index 70116de2a6..752dd403d6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,11 @@ This section accumulates work toward the **2.3.0** milestone ### Added +- **subagents:** Durable batch items can wait for explicitly named prerequisites. + An optional, once-only lead-Agent follow-up completes a previously requested + summary or delivery after the batch ends, with persisted admission, cancellation + and authorization checks. Active Goals retain continuation ownership. ([#7839](https://github.com/bytedance/deer-flow/pull/7839)) + - **chat:** Running Web chats offer Queue and Steer in the composer. Queue waits for the next turn; Steer lets the main Agent handle updates while native subagents pause and resume. Unconfirmed messages can be retried diff --git a/CHANGELOG_zh.md b/CHANGELOG_zh.md index 5e582d7c65..20e1fc6dee 100644 --- a/CHANGELOG_zh.md +++ b/CHANGELOG_zh.md @@ -14,6 +14,9 @@ ### 新增 +- **子智能体:** 后台批次可声明条目前置关系,并可登记一次主 Agent 接续,在批次结束后完成 + 用户已要求的汇总或交付。接续保留持久化启动记录、取消和授权检查;启用 Goal 时仍由 Goal 负责续跑。([#7839](https://github.com/bytedance/deer-flow/pull/7839)) + - **对话:** 网页端在任务运行时,可在输入框选择 Queue(排队执行)或 Steer(调整当前任务)。 Queue 等当前任务结束后再执行下一条消息;Steer 让主 Agent 先处理新要求,原生子 Agent 可暂停后恢复。 Gateway 重启后,尚未确认生效的消息可以手动重试。([#7729]) diff --git a/README.md b/README.md index 8d04730fd4..f2c227586b 100644 --- a/README.md +++ b/README.md @@ -2929,6 +2929,13 @@ Native subagent hard timeouts publish `TIMED_OUT` at the execution deadline whil Durable `batch_task` workers use one app-owned plugin snapshot for tool assembly and execution. Recovered tasks adopt the new worker's plugin snapshot after a Gateway restart; plugin objects are never stored in durable task records. +Batch items can declare `depends_on` using other keys in the same batch. Ready +items run without waiting for unrelated work; waiting items consume no execution +slot. For a user-requested summary or delivery after background work, the optional +`continuation` argument registers one lead-Agent follow-up. The existing Goal +keeps ownership when active. See [batch dependencies and follow-up](backend/docs/BATCH_CONTINUATION.md) +for authorization, cancellation, recovery and result semantics. + Ordinary `task` delegation and explicit durable `batch_task` execution share the startup-scoped `subagent_runtime` process capacity. Batch mode keeps large independent item sets in SQL with separate total, live, and running limits, restart recovery, bounded results, and a thread-scoped Web UI panel. The panel pages through bounded previews on demand and opens a selected item’s stored report and captured evidence for review. Citations open the matching saved excerpt without contacting the provider, and an open result refreshes until the item finishes and when the item list observes a later retry. Closed result dialogs do not fetch full details. Bulk stored results remain available through the owner-scoped JSONL export, while internal execution and authorization context never enters owner-facing responses. If the batch worker is later stopped or disabled, threads with persisted batches retain read-only item inspection and JSONL export; execution controls remain disabled until the worker is running again. See `config.example.yaml` and [the implementation contract](docs/plans/2026-08-24-subagent-batch-capacity-implementation.md) for limits and recovery semantics. While durable batch tools are available, a later owner turn can explicitly inspect selected results with `read_batch_result(batch_id, position=0, offset=0, max_chars=4000)`. The reader is restricted to the current user and thread, including in explicitly bound SDK graphs. It returns one window of a JSON document containing the stored report, execution/error state and separate acceptance criteria/verdict, without starting or waiting for work. Concatenate `content` windows before parsing: pass `next_offset` and the returned `revision` as `expected_revision` until `next_offset` is null. A changed item returns `restart_required`; discard prior windows and start at zero. Use `next_position` for a selected next item within `batch_status.total_items`. The maximum window is 8,192 characters; metadata shares the window budget, and JSON response framing adds overhead. `result_truncated` means storage already capped the report; continuation cannot recover discarded content. A succeeded item is not an accepted task, and unchecked criteria remain UNVERIFIED. No automatic result ingestion, waiting loop or UI change is added. diff --git a/README_zh.md b/README_zh.md index 18fda15983..dfa4893f78 100644 --- a/README_zh.md +++ b/README_zh.md @@ -854,6 +854,11 @@ Web UI 会在输入框上方展示当前激活的 goal,说明自动续跑为 ### Sub-Agents +显式后台 `batch_task` 可用 `depends_on` 声明同批条目的前置关系,前置工作完成后即可继续, +无需等待无关任务。用户已要求后续汇总或交付时,可用可选的 `continuation` 登记一次主 Agent 接续; +启用 Goal 时仍由 Goal 负责续跑。执行、取消、授权和恢复范围见 +[批次依赖与完成接续说明](backend/docs/BATCH_CONTINUATION.md)。 + Sub-agent 是一种执行优化,而不是遇到复杂任务时的默认选择。 lead agent 只会在委派具有明确净收益时动态拉起 sub-agents,例如真正缩短耗时的并行工作、专业能力收益或上下文隔离收益。存在跨 Agent 依赖或重叠副作用的工作不会并行分派;当专业能力或上下文隔离收益明显占优时,一条有界的顺序任务链仍可交给一个 sub-agent 完成。lead agent 会使用能取得收益的最少 sub-agents,并在每一批完成后重新评估,而不会仅仅因为任务规模大或步骤多就继续拆分。每个 sub-agent 都有自己独立的上下文、工具和终止条件,返回结构化结果后由 lead agent 验证并汇总成完整输出。 diff --git a/backend/app/gateway/app.py b/backend/app/gateway/app.py index a098929046..80a50a501f 100644 --- a/backend/app/gateway/app.py +++ b/backend/app/gateway/app.py @@ -923,11 +923,18 @@ async def _gateway_lifespan(app: FastAPI) -> AsyncGenerator[None, None]: if subagent_batches_config.enabled and batch_repo is None: raise RuntimeError("subagent_batches.enabled requires database.backend sqlite or postgres") if batch_repo is not None: + from app.subagent_batches.continuation import BatchContinuations + + continuations = BatchContinuations(app) + app.state.batch_continuations = continuations batch_service = SubagentBatchService( repository=batch_repo, config=subagent_batches_config, runtime_config=subagent_runtime_config, extensions=getattr(app.state, "extensions", None), + prepare_continuation=continuations.prepare, + dispatch_continuation=continuations.dispatch, + cancel_continuation=continuations.cancel, ) app.state.subagent_batch_service = batch_service if subagent_batches_config.enabled: diff --git a/backend/app/gateway/deps.py b/backend/app/gateway/deps.py index 1d2e6e3e0d..371998e6ee 100644 --- a/backend/app/gateway/deps.py +++ b/backend/app/gateway/deps.py @@ -1245,13 +1245,15 @@ def _selected_checkpoint_ids(config: dict[str, Any]) -> frozenset[str]: return frozenset(selected) -def _run_completion_hook(*, scheduled_task_service: Any | None): +def _run_completion_hook(*, scheduled_task_service: Any | None, batch_continuations: Any | None = None): """Notify the scheduler after ordinary run finalization.""" async def handle(record: Any) -> None: if scheduled_task_service is not None: await scheduled_task_service.handle_run_completion(record) + if batch_continuations is not None: + await batch_continuations.source_finished(record) return handle @@ -1286,6 +1288,7 @@ def get_run_context(request: Request) -> RunContext: on_run_admitted=_run_admission_hook(checkpointer=checkpointer), on_run_completed=_run_completion_hook( scheduled_task_service=scheduled_task_service, + batch_continuations=getattr(request.app.state, "batch_continuations", None), ), ) diff --git a/backend/app/gateway/extension_agent_runs.py b/backend/app/gateway/extension_agent_runs.py index 29d0785209..3f43188bfe 100644 --- a/backend/app/gateway/extension_agent_runs.py +++ b/backend/app/gateway/extension_agent_runs.py @@ -73,15 +73,32 @@ class GatewayAgentRunsHost: return None return _BoundAgentRuns(self, str(user.id), user.token_version, source, frozenset(auth.permissions)) + def capture_batch_authority(self, runs, *, user_id: str) -> dict: + """Capture only a native batch's existing host grant; never browser input.""" + if not isinstance(runs, _BoundAgentRuns) or runs._host is not self or runs._user_id != user_id or not self.active: + raise AgentRunError(403, "Batch follow-up requires its originating Agent run authority") + return {"token_version": runs._token_version, "source": runs._source, "permissions": sorted(runs._permissions)} + + async def restore_batch_authority(self, *, user_id: str, grant: dict): + """Restore a persisted native batch grant with its original permission ceiling.""" + source, version, permissions = grant.get("source"), grant.get("token_version"), grant.get("permissions") + if source not in (AUTH_SOURCE_SESSION, AUTH_SOURCE_AUTH_DISABLED) or type(version) is not int or not isinstance(permissions, list) or any(not isinstance(value, str) for value in permissions): + raise AgentRunError(403, "Invalid batch follow-up authority") + user = get_auth_disabled_user() if source == AUTH_SOURCE_AUTH_DISABLED and is_auth_disabled() else await self.load_user(user_id) if source == AUTH_SOURCE_SESSION and self.load_user is not None else None + if user is None or str(user.id) != user_id or user.token_version != version or user.needs_setup or getattr(user, "disabled", False): + raise AgentRunError(403, "Batch follow-up authorization revoked") + return _BoundAgentRuns(self, user_id, version, source, frozenset(permissions), origin=make_origin("subagent_batch")) + class _BoundAgentRuns: - def __init__(self, host, user_id: str, token_version: int, source: str, permissions: frozenset[str], plugin_namespace: str | None = None): + def __init__(self, host, user_id: str, token_version: int, source: str, permissions: frozenset[str], plugin_namespace: str | None = None, origin: dict | None = None): self._host = host self._user_id = user_id self._token_version = token_version self._source = source self._permissions = permissions self._plugin_namespace = plugin_namespace + self._origin = origin def for_plugin(self, namespace: str): if not isinstance(namespace, str) or not namespace or len(namespace) > 96 or ":" in namespace: @@ -113,7 +130,7 @@ class _BoundAgentRuns: request.state.auth = AuthContext(user, list(permissions)) # Host-set, never HTTP input: threads and runs created through this # handle are marked as extension work (activity feed, unread state). - request.state.run_origin = make_origin("extension", namespace=self._plugin_namespace) + request.state.run_origin = self._origin or make_origin("extension", namespace=self._plugin_namespace) token = set_current_user(user) try: if thread_id is not None: @@ -279,6 +296,28 @@ class _BoundAgentRuns: async def start(self, *, thread_id, input, context=None, idempotency_key=None, execution_fence=None, artifact_imports=None) -> AgentRun: return await self._start(thread_id=thread_id, input=input, context=context, idempotency_key=idempotency_key, execution_fence=execution_fence, artifact_imports=artifact_imports) + async def start_batch_followup(self, *, thread_id, batch_id, input, context, config, execution_fence, goal_wakeup=None): + """Native single follow-up; the existing SQL admission fences cancellation.""" + from app.gateway.run_models import RunCreateRequest + from app.gateway.services import resolve_thread_assistant_id, start_run + from deerflow.subagents.batch_runtime import BATCH_CONTINUATION_NAMESPACE, batch_continuation_key + + if self._origin != make_origin("subagent_batch") or execution_fence.namespace != BATCH_CONTINUATION_NAMESPACE or execution_fence.resource_id != batch_id: + raise AgentRunError(403, "Invalid native batch continuation") + async with self._request("runs:create", thread_id) as request: + request.state.extension_execution_fence = execution_fence + assistant_id = await resolve_thread_assistant_id(request, thread_id, fail_closed=True) + body = RunCreateRequest( + assistant_id=assistant_id, + input=_json_object(input), + context=_json_object(context), + config=_json_object(config), + metadata={"batch_goal_wakeup": _json_object(goal_wakeup)} if goal_wakeup is not None else None, + on_disconnect="continue", + ) + record = await start_run(body, thread_id, request, idempotency_key=batch_continuation_key(batch_id), require_existing_thread=True) + return _view(record) + async def resume(self, *, thread_id, resume, idempotency_key=None) -> AgentRun: if resume is None: raise AgentRunError(422, "An explicit interrupt response is required") diff --git a/backend/app/gateway/routers/subagent_batches.py b/backend/app/gateway/routers/subagent_batches.py index a1d44952a8..cafdd47c13 100644 --- a/backend/app/gateway/routers/subagent_batches.py +++ b/backend/app/gateway/routers/subagent_batches.py @@ -8,7 +8,7 @@ from collections.abc import AsyncIterator from fastapi import APIRouter, HTTPException, Path, Query, Request from fastapi.responses import StreamingResponse -from app.gateway.authz import require_permission +from app.gateway.authz import require_cancel_permission_if, require_permission from app.gateway.deps import ( get_current_user, get_subagent_batch_repo, @@ -110,6 +110,7 @@ async def cancel_batch(thread_id: ThreadId, batch_id: str, request: Request) -> if not getattr(request.app.state, "subagent_batches_available", False): raise HTTPException(status_code=503, detail="Subagent batch worker is not running") _repo, user_id, _batch = await _owned_batch(request, thread_id, batch_id) + require_cancel_permission_if(request, _batch.get("continuation_status") is not None) result = await get_subagent_batch_service(request).cancel_batch(batch_id=batch_id, user_id=user_id) if result is None: raise HTTPException(status_code=404, detail="Subagent batch not found") diff --git a/backend/app/gateway/routers/threads.py b/backend/app/gateway/routers/threads.py index 586c5b9b80..175602bd82 100644 --- a/backend/app/gateway/routers/threads.py +++ b/backend/app/gateway/routers/threads.py @@ -805,6 +805,9 @@ async def _delete_thread_data_with_reservation(thread_id: str, request: Request) # persisted runs/events/feedback and the thread_meta row all belong to the # same owner, so they must not resolve their scope independently. user_id = get_effective_user_id() + batch_repo = getattr(request.app.state, "subagent_batch_repo", None) + if batch_repo is not None: + await batch_repo.suppress_thread_continuations(thread_id) # Legacy IDs may predate the canonical filesystem-safe contract. They can # still be removed from metadata/checkpoint stores, but must never be @@ -1512,6 +1515,9 @@ async def set_thread_goal(thread_id: ThreadId, body: ThreadGoalRequest, request: async with reserve_checkpoint_write(request, thread_id, user_id=get_effective_user_id()): await _ensure_thread_for_goal(thread_id, request) await write_thread_goal(checkpointer, thread_id, goal, as_node="goal", create_if_missing=True) + batch_repo = getattr(request.app.state, "subagent_batch_repo", None) + if batch_repo is not None: + await batch_repo.suppress_thread_continuations(thread_id) except ValueError as exc: raise HTTPException(status_code=422, detail=str(exc)) from exc except ConflictError: @@ -1532,6 +1538,9 @@ async def clear_thread_goal(thread_id: ThreadId, request: Request) -> ThreadGoal try: async with reserve_checkpoint_write(request, thread_id, user_id=get_effective_user_id()): await write_thread_goal(checkpointer, thread_id, None, as_node="goal") + batch_repo = getattr(request.app.state, "subagent_batch_repo", None) + if batch_repo is not None: + await batch_repo.suppress_thread_continuations(thread_id) except ConflictError: raise HTTPException(status_code=409, detail="Thread has a run in flight. Clear the goal after the run finishes.") from None except LookupError: diff --git a/backend/app/gateway/services.py b/backend/app/gateway/services.py index 0a6bb14fb0..a99bb0bbf5 100644 --- a/backend/app/gateway/services.py +++ b/backend/app/gateway/services.py @@ -688,7 +688,7 @@ def _normalize_external_goal(value: Any, *, location: str) -> dict[str, Any] | N for field in ("continuation_count", "max_continuations", "no_progress_count", "max_no_progress_continuations"): if field in value and (type(value[field]) is not int or value[field] < 0): raise HTTPException(status_code=400, detail=f"{location}.{field} must be a non-negative integer") - return {**value, "objective": objective} + return {**{key: item for key, item in value.items() if key != "last_batch_wakeup_run_id"}, "objective": objective} def strip_server_owned_state_metadata(values: Mapping[str, Any]) -> dict[str, Any]: diff --git a/backend/app/subagent_batches/continuation.py b/backend/app/subagent_batches/continuation.py new file mode 100644 index 0000000000..53df67a18a --- /dev/null +++ b/backend/app/subagent_batches/continuation.py @@ -0,0 +1,139 @@ +"""Native batch follow-ups through the existing revocable AgentRuns host.""" + +import copy +import json + +from deerflow_extension_api.agent_runs import AgentRunError, ExecutionFence + +from deerflow.agents.human_input import extract_unanswered_human_input_request +from deerflow.mcp_scope import THREAD_INCARNATION_CONTEXT_KEY +from deerflow.runtime.goal import goal_waits_for_external_result, read_thread_goal +from deerflow.runtime.runs.schemas import GOAL_RUN_HARD_STOP_REASONS + +_TRUSTED = object() + + +class BatchContinuations: + def __init__(self, app): + self.app = app + + async def prepare(self, request): + from app.gateway.services import _CONTEXT_CONFIGURABLE_KEYS, _MODEL_SELECTION_CONTEXT_KEYS + + host = self.app.state.agent_runs_host + grant = host.capture_batch_authority(request.continuation_authority, user_id=request.user_id) + source = await request.continuation_authority.get(thread_id=request.thread_id, run_id=request.run_id) + if source.status not in ("pending", "running"): + raise AgentRunError(409, "The originating Run has already stopped") + record = await self.app.state.run_manager.get(request.run_id, user_id=request.user_id) + if record is None or record.user_id != request.user_id or record.thread_id != request.thread_id: + raise AgentRunError(403, "Batch follow-up must belong to its originating Run") + captured = request.continuation_context or {} + if (record.metadata or {}).get("deerflow_origin", {}).get("kind") == "subagent_batch" and (captured.get("goal") or {}).get("status") != "active": + raise AgentRunError(409, "A batch follow-up cannot register another automatic follow-up") + context = {key: value for key, value in captured.get("context", {}).items() if key in _CONTEXT_CONFIGURABLE_KEYS | _MODEL_SELECTION_CONTEXT_KEYS and key not in ("agent_name", "is_bootstrap")} + if isinstance(captured.get("context", {}).get("locale"), str): + context["locale"] = captured["context"]["locale"] + scope = request.execution_spec.get("knowledge_scope") + return { + "prompt": request.continuation, + "grant": grant, + "thread_incarnation": request.execution_spec.get(THREAD_INCARNATION_CONTEXT_KEY), + "goal": copy.deepcopy(captured.get("goal")), + "context": copy.deepcopy(context), + "config": copy.deepcopy(captured.get("config", {})), + "knowledge_scope": copy.deepcopy(scope), + } + + async def dispatch(self, batch): + continuation = batch["continuation"] + try: + runs = await self.app.state.agent_runs_host.restore_batch_authority(user_id=batch["user_id"], grant=continuation["grant"]) + existing_id = await self.app.state.subagent_batch_repo.continuation_run_id(batch["id"], user_id=batch["user_id"]) + if existing_id is not None: + existing = await runs.get(thread_id=batch["thread_id"], run_id=existing_id) + return {"status": "dispatched", "run_id": existing.run_id} + source = await runs.get(thread_id=batch["thread_id"], run_id=batch["run_id"]) + if source.status in ("pending", "running"): + return None + if source.status != "success" or source.stop_reason in GOAL_RUN_HARD_STOP_REASONS: + return {"status": "suppressed"} + thread = await self.app.state.thread_store.get(batch["thread_id"], user_id=batch["user_id"]) + if thread is None or thread.get("incarnation") != continuation.get("thread_incarnation"): + return {"status": "suppressed"} + state = await runs.get_state(thread_id=batch["thread_id"]) + values = state.get("values", {}) + if extract_unanswered_human_input_request(values.get("messages")) is not None: + return {"status": "waiting"} + goal = values.get("goal") + goal_identity = {key: goal.get(key) for key in ("created_at", "status", "objective")} if isinstance(goal, dict) else None + expected_goal = continuation.get("goal") + if goal_identity != expected_goal: + return {"status": "suppressed"} + goal_wakeup = None + if expected_goal is not None: + latest = await self.app.state.subagent_batch_repo.latest_run_state(batch["thread_id"], user_id=batch["user_id"]) + if latest and latest["status"] in ("pending", "running"): + return None + if latest and (latest["status"] != "success" or latest.get("stop_reason") in GOAL_RUN_HARD_STOP_REASONS): + return {"status": "suppressed"} + if not goal_waits_for_external_result(goal, expected_goal): + return None if not (goal.get("last_evaluation") or {}).get("stand_down_reason") else {"status": "suppressed"} + goal_wakeup = expected_goal + public = await self.app.state.subagent_batch_repo.get_batch(batch["id"], user_id=batch["user_id"]) + if public is None or public.get("continuation_status") == "cancelled" or public["status"] == "cancelled": + return {"status": "suppressed"} + if public["status"] not in ("completed", "failed"): + return None + # Stable input is essential after an unknown admission outcome. + # Counts/results stay in the existing owner-scoped result reader. + content = "The requested background batch has finished. Continue the agreed follow-up using its stored results; preserve any failures or missing information.\n" + json.dumps( + {"batch_id": batch["id"], "follow_up": continuation["prompt"]}, ensure_ascii=False + ) + from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_KEY + + metadata = {"hide_from_ui": True, "untrusted_input": True} + if continuation.get("knowledge_scope") is not None: + metadata[KNOWLEDGE_SCOPE_KEY] = continuation["knowledge_scope"] + result = await runs.start_batch_followup( + thread_id=batch["thread_id"], + batch_id=batch["id"], + input={"messages": [{"role": "user", "content": content, "additional_kwargs": metadata}]}, + context=continuation.get("context"), + config=continuation.get("config"), + execution_fence=ExecutionFence(**batch["execution_fence"]), + goal_wakeup=goal_wakeup, + ) + return {"status": "dispatched", "run_id": result.run_id} + except AgentRunError as exc: + if exc.status_code == 409: + return None + if exc.status_code in (400, 401, 403, 404, 422): + return {"status": "suppressed"} + raise + + async def cancel(self, batch_id, user_id, *, thread_id=None, authority=_TRUSTED, authorize_only=False): + if authority is not _TRUSTED: + self.app.state.agent_runs_host.capture_batch_authority(authority, user_id=user_id) + async with authority._request("runs:cancel", thread_id): + if authorize_only: + return + elif authorize_only: + return + run_id = await self.app.state.subagent_batch_repo.continuation_run_id(batch_id, user_id=user_id) + if run_id is None: + return + record = await self.app.state.run_manager.get(run_id, user_id=user_id) + if record is not None: + await self.app.state.run_manager.cancel(run_id) + + async def source_finished(self, record): + if record.stop_reason == "goal_wakeup_superseded": + return + if record.status == "success" and record.stop_reason not in GOAL_RUN_HARD_STOP_REASONS: + return + for batch in await self.app.state.subagent_batch_repo.list_continuation_batches_for_run(record.run_id): + await self.app.state.subagent_batch_service.cancel_batch(batch_id=batch["id"], user_id=batch["user_id"]) + goal = await read_thread_goal(self.app.state.checkpointer, record.thread_id) + if goal and goal.get("status") == "active" and goal.get("created_at"): + await self.app.state.subagent_batch_repo.suppress_thread_continuations(record.thread_id, goal_created_at=goal["created_at"]) diff --git a/backend/docs/BATCH_CONTINUATION.md b/backend/docs/BATCH_CONTINUATION.md new file mode 100644 index 0000000000..a7fc8d696a --- /dev/null +++ b/backend/docs/BATCH_CONTINUATION.md @@ -0,0 +1,64 @@ +# Native batch dependencies and follow-up + +Optional fields extend the existing SQL-backed batch service. Omitting them keeps +independent batch execution. Items remain read-only or idempotent under the existing +retry/recovery policy. + +## Dependencies + +Use `depends_on` to name prerequisite item keys in the same batch. Unknown keys, +self-references and cycles are rejected before saving. Dependent items consume no +worker slot until every prerequisite succeeds without a guardrail stop reason. +Retryable attempts keep dependents waiting; exhausted failures, cancellation or +terminal capped results fail unstarted descendants. Independent items continue. +Retrying a failed prerequisite by hand does not requeue the descendants it already +failed; retry each blocked item as well. + +Item prompts must name their input/output paths. Ordering does not certify factual +accuracy or make advisory acceptance mandatory. Use existing result readers; full +reports are not copied into every dependent's context. + +## Requested follow-up + +`continuation` describes remaining work the user already requested. After terminal +batch results are saved, the lead performs that work in the original conversation +through ordinary Run admission. Preserve the originating model selection, reasoning, +limits and admitted retrieval scope. Busy conversations wait. While a question is +unanswered the claim is kept until its lease expires, so the conversation is not +reread on every poll. + +A stable idempotency key reconciles lost admission replies without another Run. +Pending claims rotate so busy conversations do not starve others. Manual item retries +hold an unadmitted follow-up but do not re-arm one already dispatched. Without a Goal, +a follow-up cannot request another automatic follow-up. + +An active Goal wakes only from its own `external_wait`, consuming one existing +continuation allowance, once even after recovery. Other blockers, no-progress limits +and hard stops remain effective. Goal replacement, clearing and thread deletion +suppress pending work; a failed, cancelled or hard-stopped source Run cannot restart. + +## Authority, cancellation and recovery + +Follow-up needs a session-authorized Gateway AgentRuns capability or local auth-disabled +mode. Internal/PAT callers and unsupported embedded runtimes retain ordinary batches. +Persist only the original permission ceiling, account token version and auth source; +no cookie, API key or Python capability. Restore checks current account/permissions, +conversation ownership and incarnation. Private continuation data is not returned. + +Batch cancellation revokes the existing SQL execution fence and stops an admitted +follow-up, including after a lost launch reply. HTTP and tool cancellation of a batch +with a follow-up also require `runs:cancel`. Trusted internal cleanup retains its existing +path. Claim ownership/expiry lives only in the existing execution-lease table. +Completed item results remain available. The Web panel discovers follow-up Runs +through existing observers; refresh never launches work. Chat Stop remains available. + +## Upgrade and verification + +Migration `0052_batch_continuation` adds dependency edges and nullable follow-up data, +status and Run ID. Existing rows need no backfill. Upgrade all batch workers before +using dependencies. Do not downgrade with opted-in work outstanding; downgrade drops +new metadata while retaining items and reports. + +Focused regressions cover dependency/retry ordering, SQL claims and recovery, +permission revocation, admission races, Goal limits and page reconnection. PostgreSQL +coverage uses `TEST_POSTGRES_URI`; browser backend timing is simulated. diff --git a/backend/packages/harness/deerflow/agents/goal_state.py b/backend/packages/harness/deerflow/agents/goal_state.py index 776bc409a6..3981cc3a17 100644 --- a/backend/packages/harness/deerflow/agents/goal_state.py +++ b/backend/packages/harness/deerflow/agents/goal_state.py @@ -30,6 +30,7 @@ class GoalState(TypedDict): no_progress_count: int max_no_progress_continuations: int last_evaluation: NotRequired[dict[str, Any]] + last_batch_wakeup_run_id: NotRequired[str] class GoalOutcomeState(TypedDict): diff --git a/backend/packages/harness/deerflow/agents/lead_agent/prompt.py b/backend/packages/harness/deerflow/agents/lead_agent/prompt.py index 328225c319..633cbb5521 100644 --- a/backend/packages/harness/deerflow/agents/lead_agent/prompt.py +++ b/backend/packages/harness/deerflow/agents/lead_agent/prompt.py @@ -576,12 +576,13 @@ A single subagent is justified only by material specialist or context-isolation durable_batch_guidance = """ ## Explicit durable batch mode -`batch_task` is a separate execution mode for a large collection of independent, -idempotent or read-only items. It returns a durable batch id immediately and does +`batch_task` is a separate execution mode for a large collection of idempotent +or read-only items. It returns a durable batch id immediately and does not consume the ordinary `task` per-run total. Never infer batch mode from item count and never emulate it by repeatedly calling `task`. -- Every item must be self-contained and must not depend on another item's output. +- Keep simultaneously runnable items independent. Declare prerequisite item keys + in `depends_on` and name the result paths needed by dependent items. - Give every item a stable unique key; retries reuse that key as idempotency identity. - Set total, live-window, and running concurrency separately. A high total never implies that all items become live or run at once. diff --git a/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md b/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md index c088c56ba6..e2b80b6d48 100644 --- a/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md +++ b/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md @@ -40,7 +40,7 @@ The empty-DB path keeps using `create_all` because `Base.metadata` is the only a `0039_user_disabled` → `0040_scheduled_group_notice` → `0041_im_notice_preferences` → `0042_channel_leases` → `0043_runtime_im_notices` → `0044_channel_runtime_configs` → -`0045_completed_run_evidence` → `0046_skill_mutations` → `0047_mcp_oauth_tokens` → `0048_extension_execution_leases`, then `0049_run_inputs` and `0049_run_producer_fencing` → `0050_history_seed_authority`, joined by `0051_merge_run_inputs_fencing` (head). +`0045_completed_run_evidence` → `0046_skill_mutations` → `0047_mcp_oauth_tokens` → `0048_extension_execution_leases`, then `0049_run_inputs` and `0049_run_producer_fencing` → `0050_history_seed_authority`, joined by `0051_merge_run_inputs_fencing` → `0052_batch_continuation` (head). The preference revision uses an owner/key table with a cascading users FK; project documents use an owner-scoped shelf. MCP lease tokens, batch evidence, and run-idempotency identity add nullable columns without historical backfill or diff --git a/backend/packages/harness/deerflow/persistence/migrations/versions/0052_batch_continuation.py b/backend/packages/harness/deerflow/persistence/migrations/versions/0052_batch_continuation.py new file mode 100644 index 0000000000..9b08fea084 --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/migrations/versions/0052_batch_continuation.py @@ -0,0 +1,44 @@ +"""Optional native batch dependencies and completion follow-up. + +Revision ID: 0052_batch_continuation +Revises: 0051_merge_run_inputs_fencing +""" + +import sqlalchemy as sa +from alembic import op + +revision = "0052_batch_continuation" +down_revision = "0051_merge_run_inputs_fencing" +branch_labels = None +depends_on = None + + +def upgrade(): + from deerflow.persistence.migrations._helpers import safe_add_column + + for column in ( + sa.Column("continuation_json", sa.JSON(), nullable=True), + sa.Column("continuation_status", sa.String(24), nullable=True), + sa.Column("continuation_run_id", sa.String(64), nullable=True), + ): + safe_add_column("subagent_batches", column) + if "ix_subagent_batches_continuation_status" not in {index["name"] for index in sa.inspect(op.get_bind()).get_indexes("subagent_batches")}: + op.create_index("ix_subagent_batches_continuation_status", "subagent_batches", ["continuation_status"]) + if "subagent_batch_dependencies" not in sa.inspect(op.get_bind()).get_table_names(): + op.create_table( + "subagent_batch_dependencies", + sa.Column("item_id", sa.String(64), sa.ForeignKey("subagent_batch_items.id", ondelete="CASCADE"), primary_key=True), + sa.Column("prerequisite_id", sa.String(64), sa.ForeignKey("subagent_batch_items.id", ondelete="CASCADE"), primary_key=True), + ) + op.create_index("ix_subagent_batch_dependencies_prerequisite", "subagent_batch_dependencies", ["prerequisite_id"]) + + +def downgrade(): + from deerflow.persistence.migrations._helpers import safe_drop_column + + if "subagent_batch_dependencies" in sa.inspect(op.get_bind()).get_table_names(): + op.drop_table("subagent_batch_dependencies") + if "ix_subagent_batches_continuation_status" in {index["name"] for index in sa.inspect(op.get_bind()).get_indexes("subagent_batches")}: + op.drop_index("ix_subagent_batches_continuation_status", table_name="subagent_batches") + for name in ("continuation_json", "continuation_status", "continuation_run_id"): + safe_drop_column("subagent_batches", name) diff --git a/backend/packages/harness/deerflow/persistence/subagent_batches/model.py b/backend/packages/harness/deerflow/persistence/subagent_batches/model.py index 7aeed073bc..b264e424ab 100644 --- a/backend/packages/harness/deerflow/persistence/subagent_batches/model.py +++ b/backend/packages/harness/deerflow/persistence/subagent_batches/model.py @@ -26,6 +26,9 @@ class SubagentBatchRow(Base): max_running_items: Mapped[int] = mapped_column(Integer) max_attempts: Mapped[int] = mapped_column(Integer) execution_spec: Mapped[dict[str, Any]] = mapped_column(JSON) + continuation_json: Mapped[dict[str, Any] | None] = mapped_column(JSON, nullable=True) + continuation_status: Mapped[str | None] = mapped_column(String(24), nullable=True, index=True) + continuation_run_id: Mapped[str | None] = mapped_column(String(64), nullable=True) created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) @@ -72,3 +75,12 @@ class SubagentBatchItemRow(Base): UniqueConstraint("batch_id", "position", name="uq_subagent_batch_items_position"), Index("ix_subagent_batch_items_claim", "status", "lease_expires_at", "batch_id"), ) + + +class SubagentBatchDependencyRow(Base): + __tablename__ = "subagent_batch_dependencies" + + item_id: Mapped[str] = mapped_column(ForeignKey("subagent_batch_items.id", ondelete="CASCADE"), primary_key=True) + prerequisite_id: Mapped[str] = mapped_column(ForeignKey("subagent_batch_items.id", ondelete="CASCADE"), primary_key=True) + + __table_args__ = (Index("ix_subagent_batch_dependencies_prerequisite", "prerequisite_id"),) diff --git a/backend/packages/harness/deerflow/persistence/subagent_batches/sql.py b/backend/packages/harness/deerflow/persistence/subagent_batches/sql.py index 1ab94b35e5..791563ada0 100644 --- a/backend/packages/harness/deerflow/persistence/subagent_batches/sql.py +++ b/backend/packages/harness/deerflow/persistence/subagent_batches/sql.py @@ -8,11 +8,14 @@ from typing import Any from sqlalchemy import func, or_, select, update from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker -from sqlalchemy.orm import defer +from sqlalchemy.orm import aliased, defer -from deerflow.persistence.subagent_batches.model import SubagentBatchItemRow, SubagentBatchRow +from deerflow.persistence.extension_leases.model import ExtensionExecutionLeaseRow +from deerflow.persistence.run.model import RunRow +from deerflow.persistence.subagent_batches.model import SubagentBatchDependencyRow, SubagentBatchItemRow, SubagentBatchRow from deerflow.subagents.acceptance_checks import AcceptanceVerdict, validate_acceptance_verdict -from deerflow.subagents.batch_runtime import BatchItemInput +from deerflow.subagents.batch_dependencies import batch_dependencies +from deerflow.subagents.batch_runtime import BATCH_CONTINUATION_NAMESPACE, BatchItemInput, batch_continuation_key from deerflow.subagents.report_contract import normalize_acceptance_criteria from deerflow.utils.time import coerce_iso @@ -70,6 +73,8 @@ class SubagentBatchRepository: for key in _BATCH_TIMESTAMP_FIELDS: if data.get(key) is not None: data[key] = coerce_iso(data[key]) + if row.continuation_status is not None: + data.update(continuation_status=row.continuation_status, continuation_run_id=row.continuation_run_id) return data @staticmethod @@ -111,7 +116,9 @@ class SubagentBatchRepository: max_running_items: int, max_attempts: int, execution_spec: dict[str, Any], + continuation: dict[str, Any] | None = None, ) -> dict[str, Any]: + dependencies = batch_dependencies(items) now = datetime.now(UTC) batch = SubagentBatchRow( id=batch_id, @@ -128,6 +135,8 @@ class SubagentBatchRepository: max_running_items=max_running_items, max_attempts=max_attempts, execution_spec=execution_spec, + continuation_json=continuation, + continuation_status="pending" if continuation is not None else None, created_at=now, updated_at=now, ) @@ -157,6 +166,10 @@ class SubagentBatchRepository: # key can fail here before commit. await session.flush() session.add_all(rows) + if dependencies: + await session.flush() + identities = {row.item_key: row.id for row in rows} + session.add_all(SubagentBatchDependencyRow(item_id=identities[key], prerequisite_id=identities[dependency]) for key, required in dependencies.items() for dependency in required) await session.commit() except IntegrityError: await session.rollback() @@ -241,6 +254,164 @@ class SubagentBatchRepository: values.append(value) return values + @staticmethod + def _dependencies_ready(): + prerequisite = aliased(SubagentBatchItemRow) + return ( + ~select(SubagentBatchDependencyRow.item_id) + .join(prerequisite, prerequisite.id == SubagentBatchDependencyRow.prerequisite_id) + .where( + SubagentBatchDependencyRow.item_id == SubagentBatchItemRow.id, + or_(prerequisite.status != "succeeded", prerequisite.stop_reason.is_not(None)), + ) + .exists() + ) + + async def _fail_blocked_dependencies(self, session, batch, *, now): + prerequisite = aliased(SubagentBatchItemRow) + blocked = ( + select(SubagentBatchDependencyRow.item_id) + .join(prerequisite, prerequisite.id == SubagentBatchDependencyRow.prerequisite_id) + .where( + prerequisite.batch_id == batch.id, + or_(prerequisite.status.in_(("failed", "cancelled")), (prerequisite.status == "succeeded") & prerequisite.stop_reason.is_not(None)), + ) + .cte("blocked_batch_dependencies", recursive=True) + ) + blocked = blocked.union(select(SubagentBatchDependencyRow.item_id).join(blocked, SubagentBatchDependencyRow.prerequisite_id == blocked.c.item_id)) + if (await session.execute(select(blocked.c.item_id).limit(1))).first() is None: + return + changed = await session.execute( + update(SubagentBatchItemRow) + .where( + SubagentBatchItemRow.batch_id == batch.id, + SubagentBatchItemRow.status.in_(("pending", "queued")), + SubagentBatchItemRow.id.in_(select(blocked.c.item_id)), + ) + .values(status="failed", error="A prerequisite did not complete successfully", completed_at=now, updated_at=now) + .returning(SubagentBatchItemRow.id) + .execution_options(synchronize_session=False) + ) + if changed.first() is not None: + await self._refresh_batch_status(session, batch, now=now) + + async def claim_continuations(self, *, now: datetime, lease_owner: str, lease_seconds: int, limit: int = 8, exclude_batch_ids: tuple[str, ...] = ()) -> list[dict]: + """Claim one follow-up after terminal results commit; retries retain the batch ID.""" + live_claim = ( + select(ExtensionExecutionLeaseRow.resource_id) + .where( + ExtensionExecutionLeaseRow.namespace == BATCH_CONTINUATION_NAMESPACE, + ExtensionExecutionLeaseRow.resource_id == SubagentBatchRow.id, + ExtensionExecutionLeaseRow.expires_at >= now.timestamp(), + ) + .exists() + ) + ready = or_(SubagentBatchRow.continuation_status == "pending", (SubagentBatchRow.continuation_status == "dispatching") & ~live_claim) + claimed = [] + async with self._sf() as session: + rows = ( + ( + await session.execute( + select(SubagentBatchRow) + .where(SubagentBatchRow.status.in_(("completed", "failed")), SubagentBatchRow.id.not_in(exclude_batch_ids), ready) + .order_by(SubagentBatchRow.updated_at, SubagentBatchRow.completed_at, SubagentBatchRow.id) + .limit(limit) + .with_for_update(skip_locked=True) + ) + ) + .scalars() + .all() + ) + for batch in rows: + changed = await session.execute( + update(SubagentBatchRow).where(SubagentBatchRow.id == batch.id, SubagentBatchRow.status.in_(("completed", "failed")), ready).values(continuation_status="dispatching").execution_options(synchronize_session=False) + ) + if changed.rowcount == 1: + await session.refresh(batch) + fence = await session.get(ExtensionExecutionLeaseRow, (BATCH_CONTINUATION_NAMESPACE, batch.id), with_for_update=True) + expires = (now + timedelta(seconds=lease_seconds)).timestamp() + if fence is None: + fence = ExtensionExecutionLeaseRow(namespace=BATCH_CONTINUATION_NAMESPACE, resource_id=batch.id, user_id=batch.user_id, holder_id=lease_owner, epoch=1, expires_at=expires) + session.add(fence) + else: + fence.holder_id, fence.epoch, fence.expires_at = lease_owner, fence.epoch + 1, expires + claimed.append( + { + **self._execution_batch_dict(batch), + "status": batch.status, + "continuation": batch.continuation_json, + "execution_fence": {"namespace": BATCH_CONTINUATION_NAMESPACE, "resource_id": batch.id, "holder_id": lease_owner, "epoch": fence.epoch}, + } + ) + await session.commit() + return claimed + + async def finish_continuation(self, batch_id: str, *, lease_owner: str, status: str, run_id: str | None = None) -> bool: + if status not in ("pending", "dispatched", "suppressed"): + raise ValueError("Invalid continuation outcome") + async with self._sf() as session: + changed = await session.execute( + update(SubagentBatchRow) + .where( + SubagentBatchRow.id == batch_id, + SubagentBatchRow.continuation_status == "dispatching", + select(ExtensionExecutionLeaseRow.resource_id) + .where( + ExtensionExecutionLeaseRow.namespace == BATCH_CONTINUATION_NAMESPACE, + ExtensionExecutionLeaseRow.resource_id == batch_id, + ExtensionExecutionLeaseRow.holder_id == lease_owner, + ) + .exists(), + ) + .values(continuation_status=status, continuation_run_id=run_id, updated_at=datetime.now(UTC)) + ) + await session.commit() + return changed.rowcount == 1 + + async def list_continuation_batches_for_run(self, run_id: str) -> list[dict]: + async with self._sf() as session: + rows = (await session.execute(select(SubagentBatchRow).where(SubagentBatchRow.run_id == run_id, SubagentBatchRow.continuation_status.in_(("pending", "dispatching"))))).scalars().all() + return [self._execution_batch_dict(row) for row in rows] + + async def continuation_run_id(self, batch_id: str, *, user_id: str) -> str | None: + async with self._sf() as session: + return await session.scalar(select(RunRow.run_id).where(RunRow.user_id == user_id, RunRow.idempotency_key == batch_continuation_key(batch_id))) + + async def suppress_thread_continuations(self, thread_id: str, *, goal_created_at: str | None = None) -> None: + """Called while the host owns the thread's Goal/deletion operation.""" + async with self._sf() as session: + batches = select(SubagentBatchRow.id).where(SubagentBatchRow.thread_id == thread_id, SubagentBatchRow.continuation_status.in_(("pending", "dispatching"))) + if goal_created_at is not None: + batches = batches.where(SubagentBatchRow.continuation_json["goal"]["created_at"].as_string() == goal_created_at, SubagentBatchRow.continuation_json["goal"]["status"].as_string() == "active") + # Match claim/cancel lock order: batch first, then execution fence. + ids = list((await session.scalars(batches.order_by(SubagentBatchRow.completed_at, SubagentBatchRow.id).with_for_update())).all()) + if not ids: + return + ids = list( + ( + await session.scalars( + update(SubagentBatchRow).where(SubagentBatchRow.id.in_(ids), SubagentBatchRow.continuation_status.in_(("pending", "dispatching"))).values(continuation_status="suppressed").returning(SubagentBatchRow.id) + ) + ).all() + ) + await session.execute( + update(ExtensionExecutionLeaseRow).where(ExtensionExecutionLeaseRow.namespace == BATCH_CONTINUATION_NAMESPACE, ExtensionExecutionLeaseRow.resource_id.in_(ids)).values(expires_at=0, epoch=ExtensionExecutionLeaseRow.epoch + 1) + ) + await session.commit() + + async def latest_run_state(self, thread_id: str, *, user_id: str) -> dict | None: + async with self._sf() as session: + row = ( + ( + await session.execute( + select(RunRow.status, RunRow.stop_reason).where(RunRow.thread_id == thread_id, RunRow.user_id == user_id, RunRow.operation_kind == "run").order_by(RunRow.created_at.desc(), RunRow.run_id.desc()).limit(1) + ) + ) + .mappings() + .one_or_none() + ) + return dict(row) if row is not None else None + async def claim_items( self, *, @@ -325,6 +496,8 @@ class SubagentBatchRepository: if batch.status == "paused": continue + await self._fail_blocked_dependencies(session, batch, now=now) + counts = await self._counts(session, batch.id) live = counts["queued"] + counts["leased"] + counts["running"] promote_count = max(0, batch.max_live_items - live) @@ -336,6 +509,7 @@ class SubagentBatchRepository: .where( SubagentBatchItemRow.batch_id == batch.id, SubagentBatchItemRow.status == "pending", + self._dependencies_ready(), ) .order_by(SubagentBatchItemRow.position) .limit(promote_count) @@ -349,6 +523,7 @@ class SubagentBatchRepository: .where( SubagentBatchItemRow.id == item.id, SubagentBatchItemRow.status == "pending", + self._dependencies_ready(), SubagentBatchItemRow.cancel_requested_at.is_(None), SubagentBatchItemRow.batch_id.in_(runnable_batch), ) @@ -370,6 +545,7 @@ class SubagentBatchRepository: .where( SubagentBatchItemRow.batch_id == batch.id, SubagentBatchItemRow.status == "queued", + self._dependencies_ready(), SubagentBatchItemRow.cancel_requested_at.is_(None), SubagentBatchItemRow.id.not_in(exclude_item_ids), ) @@ -389,6 +565,7 @@ class SubagentBatchRepository: .where( SubagentBatchItemRow.id == item.id, SubagentBatchItemRow.status == "queued", + self._dependencies_ready(), SubagentBatchItemRow.cancel_requested_at.is_(None), SubagentBatchItemRow.batch_id.in_(runnable_batch), ) @@ -576,6 +753,14 @@ class SubagentBatchRepository: batch = await session.get(SubagentBatchRow, batch_id, with_for_update=True) if batch is None or batch.user_id != user_id: return None + if action == "cancel" and batch.continuation_status is not None: + await session.execute(update(SubagentBatchRow).where(SubagentBatchRow.id == batch.id, SubagentBatchRow.user_id == user_id).values(continuation_status="cancelled").execution_options(synchronize_session=False)) + await session.refresh(batch) + await session.execute( + update(ExtensionExecutionLeaseRow) + .where(ExtensionExecutionLeaseRow.namespace == BATCH_CONTINUATION_NAMESPACE, ExtensionExecutionLeaseRow.resource_id == batch.id) + .values(expires_at=0, epoch=ExtensionExecutionLeaseRow.epoch + 1) + ) if action == "pause" and batch.status in ("queued", "running"): batch.status = "paused" elif action == "resume" and batch.status == "paused": @@ -635,6 +820,14 @@ class SubagentBatchRepository: # stopped. Returning without commit discards the item reset. if batch.status == "cancelled": return None + if batch.continuation_status in ("pending", "dispatching"): + batch.continuation_status = "pending" + await session.flush() + await session.execute( + update(ExtensionExecutionLeaseRow) + .where(ExtensionExecutionLeaseRow.namespace == BATCH_CONTINUATION_NAMESPACE, ExtensionExecutionLeaseRow.resource_id == batch.id) + .values(expires_at=0, epoch=ExtensionExecutionLeaseRow.epoch + 1) + ) if batch.status != "paused": batch.status = "queued" batch.completed_at = None diff --git a/backend/packages/harness/deerflow/runtime/goal.py b/backend/packages/harness/deerflow/runtime/goal.py index 0247577aed..e20e28163b 100644 --- a/backend/packages/harness/deerflow/runtime/goal.py +++ b/backend/packages/harness/deerflow/runtime/goal.py @@ -617,6 +617,11 @@ def should_continue_goal(goal: GoalState, evaluation: GoalEvaluation, *, no_prog return False if evaluation["blocker"] not in CONTINUABLE_GOAL_BLOCKERS: return False + return goal_has_continuation_budget(goal, no_progress_count=no_progress_count) + + +def goal_has_continuation_budget(goal: GoalState, *, no_progress_count: int | None = None) -> bool: + """Shared limits for the native loop and a confirmed external-result wake-up.""" if int(goal.get("continuation_count", 0)) >= int(goal.get("max_continuations", DEFAULT_MAX_GOAL_CONTINUATIONS)): return False current_no_progress = int(goal.get("no_progress_count", 0) if no_progress_count is None else no_progress_count) @@ -624,6 +629,39 @@ def should_continue_goal(goal: GoalState, evaluation: GoalEvaluation, *, no_prog return current_no_progress < max_no_progress +def goal_waits_for_external_result(goal: GoalState | None, expected: Mapping[str, Any]) -> bool: + if not goal or goal.get("status") != "active" or any(goal.get(key) != expected.get(key) for key in ("created_at", "objective")): + return False + evaluation = goal.get("last_evaluation") or {} + return evaluation.get("blocker") == "external_wait" and evaluation.get("stand_down_reason") == "blocked:external_wait" and goal_has_continuation_budget(goal) + + +async def wake_goal_after_external_wait(checkpointer: Any, thread_id: str, run_id: str, expected: Mapping[str, Any]) -> bool: + """Consume one existing Goal continuation after exclusive Run admission. + + The Run ID makes recovery of the admitted Run idempotent. Other blockers, + changed goals and the existing continuation/no-progress limits stay stopped. + """ + if checkpointer is None or not isinstance(expected, Mapping): + return False + async with goal_thread_lock(thread_id): + checkpoint = await _call_checkpointer_method(checkpointer, "aget_tuple", "get_tuple", {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}) + if checkpoint is None: + return False + goal = copy.deepcopy(checkpoint.checkpoint.get("channel_values", {}).get("goal")) + if not isinstance(goal, dict) or goal.get("status") != "active" or any(goal.get(key) != expected.get(key) for key in ("created_at", "objective")): + return False + if goal.get("last_batch_wakeup_run_id") == run_id: + return True + if not goal_waits_for_external_result(goal, expected): + return False + goal["continuation_count"] = int(goal.get("continuation_count", 0)) + 1 + goal["last_batch_wakeup_run_id"] = run_id + goal["updated_at"] = now_iso() + await write_thread_goal(checkpointer, thread_id, goal, as_node="batch_completion", expected_checkpoint_id=_checkpoint_id_from_tuple(checkpoint)) + return True + + def latest_visible_assistant_signature(messages: list[Any]) -> str: """Return a stable signature of the latest visible assistant evidence. diff --git a/backend/packages/harness/deerflow/runtime/run_origin.py b/backend/packages/harness/deerflow/runtime/run_origin.py index 7624e0de66..5cc685671a 100644 --- a/backend/packages/harness/deerflow/runtime/run_origin.py +++ b/backend/packages/harness/deerflow/runtime/run_origin.py @@ -22,7 +22,7 @@ import re from collections.abc import Mapping DEERFLOW_ORIGIN_KEY = "deerflow_origin" -ORIGIN_KINDS: frozenset[str] = frozenset({"schedule", "im_channel", "github", "extension", "mcp_notification"}) +ORIGIN_KINDS: frozenset[str] = frozenset({"schedule", "im_channel", "github", "extension", "mcp_notification", "subagent_batch"}) _PROVIDER_RE = re.compile(r"[a-z0-9_-]{1,32}") _MAX_NAMESPACE_CHARS = 96 diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 2c82a91ede..9aec954f02 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -91,10 +91,12 @@ from deerflow.runtime.goal import ( read_thread_goal, should_continue_goal, visible_conversation_signature, + wake_goal_after_external_wait, write_thread_goal, ) from deerflow.runtime.goal_progress import observe_goal_operations from deerflow.runtime.keyed_lock import AsyncKeyedLockTable +from deerflow.runtime.run_origin import origin_kind_of from deerflow.runtime.runs.schemas import GOAL_RUN_HARD_STOP_REASONS from deerflow.runtime.runs.stream_cleanup import close_agent_stream from deerflow.runtime.serialization import serialize @@ -1232,6 +1234,11 @@ async def run_agent( return started = True + goal_wakeup = (record.metadata or {}).get("batch_goal_wakeup") if origin_kind_of(record.metadata) == "subagent_batch" else None + if goal_wakeup is not None: + if not await wake_goal_after_external_wait(checkpointer, thread_id, run_id, goal_wakeup): + await run_manager.set_status(run_id, RunStatus.interrupted, stop_reason="goal_wakeup_superseded") + return if ctx.verify_artifacts is not None: await ctx.verify_artifacts() diff --git a/backend/packages/harness/deerflow/subagents/AGENTS.md b/backend/packages/harness/deerflow/subagents/AGENTS.md index 7190a6c38b..e1368c9d6a 100644 --- a/backend/packages/harness/deerflow/subagents/AGENTS.md +++ b/backend/packages/harness/deerflow/subagents/AGENTS.md @@ -2,6 +2,9 @@ Drain cleanup; preserve terminal outcomes on cancel. +Optional batch dependencies and one-shot follow-up reuse batch claims and AgentRuns. +See [behavior and ownership](../../../../docs/BATCH_CONTINUATION.md): keep one execution lease and check conversation state before admission; never persist live capabilities. + Accumulate validated ToolMessage knowledge across values frames in `SubagentResult.knowledge_sources`: at most 100 sources, deduplicated by evidence ID. Publish before terminal state; retain through compaction. Ordinary `task` forwards diff --git a/backend/packages/harness/deerflow/subagents/batch_dependencies.py b/backend/packages/harness/deerflow/subagents/batch_dependencies.py new file mode 100644 index 0000000000..02e8744eda --- /dev/null +++ b/backend/packages/harness/deerflow/subagents/batch_dependencies.py @@ -0,0 +1,24 @@ +"""Optional dependencies between items of one existing native batch.""" + +from graphlib import CycleError, TopologicalSorter + +from deerflow.subagents.batch_runtime import BatchItemInput + + +def batch_dependencies(items: list[BatchItemInput]) -> dict[str, tuple[str, ...]]: + if not any(item.get("depends_on") for item in items): + return {} + keys = {item["key"] for item in items} + if len(keys) != len(items): + raise ValueError("Batch item keys must be unique") + graph = {} + for item in items: + dependencies = item.get("depends_on", []) + if not isinstance(dependencies, (list, tuple)) or any(not isinstance(key, str) or key not in keys or key == item["key"] for key in dependencies): + raise ValueError("Dependencies must name other items in this batch") + graph[item["key"]] = tuple(dict.fromkeys(dependencies)) + try: + TopologicalSorter(graph).prepare() + except CycleError: + raise ValueError("Batch dependencies must not contain a cycle") from None + return graph diff --git a/backend/packages/harness/deerflow/subagents/batch_runtime.py b/backend/packages/harness/deerflow/subagents/batch_runtime.py index dfd2225c01..9ef39731b2 100644 --- a/backend/packages/harness/deerflow/subagents/batch_runtime.py +++ b/backend/packages/harness/deerflow/subagents/batch_runtime.py @@ -6,11 +6,18 @@ import threading from dataclasses import dataclass from typing import Any, NotRequired, Protocol, TypedDict +BATCH_CONTINUATION_NAMESPACE = "native-subagent-batches" + + +def batch_continuation_key(batch_id: str) -> str: + return "native-batch-followup:" + batch_id + class BatchItemInput(TypedDict): key: str prompt: str acceptance_criteria: NotRequired[list[str] | None] + depends_on: NotRequired[list[str]] @dataclass(frozen=True) @@ -26,6 +33,9 @@ class BatchSubmitRequest: max_live_items: int | None max_running_items: int | None execution_spec: dict[str, Any] + continuation: str | None = None + continuation_authority: Any = None + continuation_context: dict[str, Any] | None = None class SubagentBatchSubmitter(Protocol): diff --git a/backend/packages/harness/deerflow/subagents/batch_service.py b/backend/packages/harness/deerflow/subagents/batch_service.py index d2e4b30f45..c798627019 100644 --- a/backend/packages/harness/deerflow/subagents/batch_service.py +++ b/backend/packages/harness/deerflow/subagents/batch_service.py @@ -134,9 +134,15 @@ def _log_release_failure(pending: Future[None], *, execution_id: str) -> None: logger.error("Registry release failed for batch execution %s", execution_id, exc_info=exc) +_TRUSTED = object() + + class SubagentBatchService: """Lease, execute, and recover durable native-subagent batch items.""" + # cancel_batch enforces runs:cancel for follow-ups with the caller's authority. + accepts_continuation_authority = True + # Also used by focused shutdown tests constructing an instance via __new__. _stopping: bool = False _stop_drains: int = 0 @@ -151,6 +157,9 @@ class SubagentBatchService: app_config: AppConfig | None = None, execution_capacity: SubagentExecutionCapacity | None = None, extensions: LoadedExtensions | None = None, + prepare_continuation=None, + dispatch_continuation=None, + cancel_continuation=None, ) -> None: self._repository = repository self._config = config @@ -169,6 +178,10 @@ class SubagentBatchService: self._executions: dict[str, asyncio.Task[None]] = {} self._execution_ids: dict[str, str] = {} self._item_batches: dict[str, str] = {} + self._prepare_continuation = prepare_continuation + self._dispatch_continuation = dispatch_continuation + self._cancel_continuation = cancel_continuation + self._continuations: dict[str, asyncio.Task] = {} async def start(self) -> None: if self._stopping: @@ -199,7 +212,7 @@ class SubagentBatchService: execution_ids = list(self._execution_ids.values()) for execution_id in execution_ids: request_cancel_background_task(execution_id) - tasks = list(self._executions.values()) + tasks = [*self._executions.values(), *getattr(self, "_continuations", {}).values()] if poller is not None: poller.cancel() @@ -227,6 +240,7 @@ class SubagentBatchService: self._executions.clear() self._execution_ids.clear() self._item_batches.clear() + getattr(self, "_continuations", {}).clear() if self._poller is poller: self._poller = None self._stop_drains = 0 @@ -249,6 +263,12 @@ class SubagentBatchService: pass async def run_once(self, *, now: datetime) -> None: + if self._dispatch_continuation is not None: + owner = f"{self._lease_owner[:80]}:{uuid.uuid4().hex}" + for batch in await self._repository.claim_continuations(now=now, lease_owner=owner, lease_seconds=self._config.lease_seconds, exclude_batch_ids=tuple(self._continuations)): + task = asyncio.create_task(self._continue_batch(batch, owner), name=f"subagent-batch-followup-{batch['id']}") + self._continuations[batch["id"]] = task + task.add_done_callback(lambda _, key=batch["id"]: self._continuations.pop(key, None)) available = max(0, self._runtime_config.max_running - len(self._executions)) if available <= 0: return @@ -292,6 +312,11 @@ class SubagentBatchService: raise ValueError(f"max_running_items must be between 1 and {self._config.max_running_items_per_batch}") if max_running > max_live: raise ValueError("max_running_items must not exceed max_live_items") + continuation = None + if request.continuation: + if self._prepare_continuation is None or self._dispatch_continuation is None: + raise ValueError("Automatic follow-up requires an authorized Gateway runtime") + continuation = await self._prepare_continuation(request) return await self._repository.create_batch( batch_id=f"subagent-batch-{uuid.uuid4().hex}", user_id=request.user_id, @@ -306,8 +331,30 @@ class SubagentBatchService: max_running_items=max_running, max_attempts=self._config.max_attempts, execution_spec=request.execution_spec, + **({"continuation": continuation} if continuation is not None else {}), ) + async def _continue_batch(self, batch, owner): + try: + outcome = await self._dispatch_continuation(batch) + if outcome and outcome.get("status") == "waiting": + # A question can stay open for hours: keep the claim until its + # lease expires instead of rereading the conversation every poll. + return + status = outcome.get("status", "dispatched") if outcome else "pending" + run_id = outcome.get("run_id") if outcome else None + saved = await self._repository.finish_continuation(batch["id"], lease_owner=owner, status=status, run_id=run_id) + if not saved and run_id and self._cancel_continuation is not None: + current = await self._repository.get_batch(batch["id"], user_id=batch["user_id"]) + if current is None or current.get("continuation_status") in ("cancelled", "suppressed"): + await self._cancel_continuation(batch["id"], batch["user_id"]) + except asyncio.CancelledError: + raise + except Exception: + # The lease will expire. Re-admission uses the same immutable key, + # so an unknown launch outcome never creates a second follow-up. + logger.warning("Batch follow-up remains unconfirmed: %s", batch["id"], exc_info=True) + async def get_batch( self, *, @@ -380,13 +427,30 @@ class SubagentBatchService: *, batch_id: str, user_id: str, + continuation_authority: Any = _TRUSTED, ) -> dict[str, Any] | None: + if continuation_authority is not _TRUSTED: + current = await self._repository.get_batch(batch_id, user_id=user_id) + if current is None: + return None + if current.get("continuation_status") is not None: + if self._cancel_continuation is None: + raise ValueError("Automatic follow-up requires an authorized Gateway runtime") + await self._cancel_continuation( + batch_id, + user_id, + thread_id=current["thread_id"], + authority=continuation_authority, + authorize_only=True, + ) batch = await self._repository.cancel_batch(batch_id, user_id=user_id) if batch is None: return None for item_id, execution_id in list(self._execution_ids.items()): if self._item_batches.get(item_id) == batch_id: request_cancel_background_task(execution_id) + if batch.get("continuation_status") == "cancelled" and self._cancel_continuation is not None: + await self._cancel_continuation(batch_id, user_id) # Normal ids are not prefixed; the renew loop observes the durable # cancellation within lease_seconds/3. Keeping cancellation durable is # what lets another worker own the HTTP control request safely. diff --git a/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py b/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py index d622b7efd9..5329c17d80 100644 --- a/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py +++ b/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py @@ -11,6 +11,7 @@ from contextvars import ContextVar from dataclasses import asdict, replace from typing import Annotated, Any, cast +from deerflow_extension_api.agent_runs import AGENT_RUNS_CONTEXT_KEY from langchain.tools import InjectedToolCallId, tool from langchain_core.messages import ToolMessage from langgraph.types import Command @@ -40,6 +41,7 @@ class BatchTaskItem(BaseModel): key: str = Field(min_length=1, max_length=128) prompt: str = Field(min_length=1, max_length=100_000) acceptance_criteria: list[str] | None = Field(default=None, description="Optional completion requirements checked separately from execution status, using the same bounded checklist as task.") + depends_on: list[str] = Field(default_factory=list, description="Keys of prerequisite items in this batch. Supply the required result paths in this item's prompt; no result text is copied automatically.") _NO_EXPLICIT_BATCH_SUBMITTER = object() @@ -160,11 +162,14 @@ async def batch_task( tool_call_id: Annotated[str, InjectedToolCallId], max_live_items: int | None = None, max_running_items: int | None = None, + continuation: str | None = None, ) -> Command: - """Submit many independent items to DeerFlow's explicit durable batch mode. + """Submit durable native-subagent items with optional explicit prerequisites. - Use this only when every item is independent, idempotent or read-only, and - can be completed without another item's output. This tool returns a batch + Use this for idempotent or read-only items. Declare depends_on when an item + needs another item's completed output; waiting items use no worker slot. + A failed or capped prerequisite prevents its dependents from starting. + This tool returns a batch identifier immediately; it never inserts thousands of results into the lead agent context. Use ``batch_status`` for a compact progress snapshot. @@ -185,6 +190,9 @@ async def batch_task( subagent_type: Native subagent definition used for every item. max_live_items: Optional queued-plus-running item window; when set it must be >= 1. max_running_items: Optional per-batch real execution concurrency; when set it must be >= 1. + continuation: Optional remaining work for the lead Agent after this batch ends, such as combining and presenting results. + Use only for follow-up already requested by the user. Omit for progress monitoring. + An active Goal resumes only from external waiting and within its existing limits. """ submitter = _batch_submitter() if submitter is None: @@ -195,6 +203,8 @@ async def batch_task( ) if not items: return _result(tool_call_id, content="A batch must contain at least one item.", error=True) + if continuation is not None and (not continuation.strip() or len(continuation) > 4000): + return _result(tool_call_id, content="The follow-up must contain 1 to 4000 characters.", error=True) keys = [item.key for item in items] if len(set(keys)) != len(keys): return _result(tool_call_id, content="Batch item keys must be unique.", error=True) @@ -223,6 +233,17 @@ async def batch_task( if not thread_id: return _result(tool_call_id, content="Durable batches require a thread_id.", error=True) user_id = resolve_runtime_user_id(runtime) + continuation_context = None + if continuation: + state = getattr(runtime, "state", None) + goal = state.get("goal") if isinstance(state, dict) else None + continuation_context = { + "goal": {key: goal.get(key) for key in ("created_at", "status", "objective")} if isinstance(goal, dict) else None, + "context": {**(runtime.config.get("configurable") or {}), **context}, + "config": {"recursion_limit": runtime.config["recursion_limit"]} if "recursion_limit" in runtime.config else {}, + } + if metadata.get("model_name"): + continuation_context["context"]["model_name"] = metadata["model_name"] run_id = context.get("run_id") submission_key = f"{run_id or thread_id}:{tool_call_id}" execution_spec = { @@ -255,6 +276,9 @@ async def batch_task( max_live_items=max_live_items, max_running_items=max_running_items, execution_spec=execution_spec, + continuation=continuation, + continuation_authority=context.get(AGENT_RUNS_CONTEXT_KEY) if continuation else None, + continuation_context=continuation_context if continuation else None, ) ) except Exception as exc: @@ -262,7 +286,11 @@ async def batch_task( return _result( tool_call_id, batch=batch, - content=(f"Batch {batch['id']} accepted with {batch['total_items']} items. It is running independently and survives Gateway restarts. Use batch_status for progress; do not launch ordinary task calls for these items."), + content=( + f"Batch {batch['id']} accepted with {batch['total_items']} items. It is running independently and survives Gateway restarts. " + + ("The requested follow-up is scheduled after completion; this turn can finish without polling. " if continuation else "Use batch_status for progress. ") + + "Do not launch ordinary task calls for these items." + ), ) @@ -303,7 +331,10 @@ async def cancel_batch(runtime: Runtime, batch_id: str) -> str: submitter = _batch_submitter() if submitter is None: return "Durable subagent batches are unavailable." - batch = await submitter.cancel_batch(batch_id=batch_id, user_id=resolve_runtime_user_id(runtime)) + context = runtime.context if runtime is not None and isinstance(runtime.context, dict) else {} + # Only submitters that run follow-ups take the caller's run authority. + authority = {"continuation_authority": context.get(AGENT_RUNS_CONTEXT_KEY)} if getattr(submitter, "accepts_continuation_authority", False) else {} + batch = await submitter.cancel_batch(batch_id=batch_id, user_id=resolve_runtime_user_id(runtime), **authority) if batch is None: return "Batch not found." return f"Batch {batch_id} cancellation requested." diff --git a/backend/tests/test_batch_continuation.py b/backend/tests/test_batch_continuation.py new file mode 100644 index 0000000000..c6ee4bebb8 --- /dev/null +++ b/backend/tests/test_batch_continuation.py @@ -0,0 +1,184 @@ +"""One durable follow-up belongs to one explicitly opted-in native batch.""" + +from datetime import UTC, datetime, timedelta + +import pytest +from test_batch_dependencies import create, finish +from test_batch_dependencies import repo as _repo_fixture + +from deerflow.persistence.engine import get_session_factory +from deerflow.persistence.subagent_batches import SubagentBatchRepository + +repo = _repo_fixture + + +async def completed_batch(repo, *, succeeded=True): + await create( + repo, + [{"key": "a", "prompt": "Read A"}], + title="Summarize records", + max_live_items=1, + max_running_items=1, + continuation={"prompt": "Write the requested summary", "grant": {"token_version": 1}}, + ) + now = datetime.now(UTC) + assert await repo.claim_continuations(now=now, lease_owner="first", lease_seconds=30) == [] + item = (await repo.claim_items(now=now, lease_owner="worker", lease_seconds=120, limit=1))[0] + await finish(repo, item, succeeded=succeeded, result="Read A", result_preview="Read A", error=None, model_name="local", completed_at=now) + return now + + +@pytest.mark.asyncio +async def test_continuation_has_one_claim_and_recovers_with_the_same_identity(repo): + now = await completed_batch(repo) + first = await repo.claim_continuations(now=now, lease_owner="first", lease_seconds=30) + assert len(first) == 1 + assert await repo.claim_continuations(now=now, lease_owner="second", lease_seconds=30) == [] + reopened = SubagentBatchRepository(get_session_factory()) + recovered = await reopened.claim_continuations(now=now + timedelta(seconds=31), lease_owner="second", lease_seconds=30) + assert recovered[0]["id"] == first[0]["id"] + assert recovered[0]["continuation"] == first[0]["continuation"] + assert not await repo.finish_continuation("batch", lease_owner="first", status="dispatched", run_id="stale") + assert await repo.finish_continuation("batch", lease_owner="second", status="dispatched", run_id="follow-up") + assert await repo.claim_continuations(now=now + timedelta(minutes=2), lease_owner="third", lease_seconds=30) == [] + public = await repo.get_batch("batch", user_id="owner") + assert public["continuation_status"] == "dispatched" and public["continuation_run_id"] == "follow-up" + assert "continuation" not in public and "token_version" not in str(public) + + +@pytest.mark.asyncio +async def test_cancelling_a_completed_batch_suppresses_an_unlaunched_followup(repo): + now = await completed_batch(repo) + await repo.cancel_batch("batch", user_id="owner") + public = await repo.get_batch("batch", user_id="owner") + assert public["status"] == "completed" + assert public["continuation_status"] == "cancelled" + assert await repo.claim_continuations(now=now, lease_owner="worker", lease_seconds=30) == [] + + +@pytest.mark.asyncio +async def test_cancel_wins_over_an_inflight_followup_receipt(repo): + now = await completed_batch(repo) + assert len(await repo.claim_continuations(now=now, lease_owner="first", lease_seconds=30)) == 1 + await repo.cancel_batch("batch", user_id="owner") + assert not await repo.finish_continuation("batch", lease_owner="first", status="dispatched", run_id="late") + assert (await repo.get_batch("batch", user_id="owner"))["continuation_status"] == "cancelled" + + +@pytest.mark.asyncio +async def test_deferred_continuations_yield_to_other_conversations_after_reopen(repo): + for index in range(17): + await create( + repo, + [{"key": "a", "prompt": "Read"}], + batch_id=f"batch-{index:02}", + user_id=f"owner-{index}", + thread_id=f"thread-{index}", + run_id=f"run-{index}", + submission_key=f"submit-{index}", + max_live_items=1, + max_running_items=1, + continuation={"prompt": "Summarize"}, + ) + items = await repo.claim_items(now=datetime.now(UTC), lease_owner="worker", lease_seconds=120, limit=17) + assert len(items) == 17 + for item in items: + await finish(repo, item, result="Done", result_preview="Done", model_name="local") + seen = set() + for index in range(3): + reopened = SubagentBatchRepository(get_session_factory()) + owner = f"followup-{index}" + claimed = await reopened.claim_continuations(now=datetime.now(UTC), lease_owner=owner, lease_seconds=30, limit=8) + seen.update(batch["id"] for batch in claimed) + for batch in claimed: + assert await reopened.finish_continuation(batch["id"], lease_owner=owner, status="pending") + assert seen == {f"batch-{index:02}" for index in range(17)} + + +@pytest.mark.asyncio +async def test_manual_item_retry_defers_an_unadmitted_followup(repo): + now = await completed_batch(repo, succeeded=False) + await repo.claim_continuations(now=now, lease_owner="before-retry", lease_seconds=30) + item = (await repo.list_items("batch", user_id="owner"))[0] + assert await repo.retry_item("batch", item["id"], user_id="owner") + assert not await repo.finish_continuation("batch", lease_owner="before-retry", status="dispatched", run_id="late") + assert await repo.claim_continuations(now=now, lease_owner="after-retry", lease_seconds=30) == [] + assert (await repo.get_batch("batch", user_id="owner"))["continuation_status"] == "pending" + + +@pytest.mark.asyncio +async def test_goal_suppression_and_claim_do_not_deadlock(repo, monkeypatch): + import asyncio + + from sqlalchemy import text + from sqlalchemy.ext.asyncio import AsyncSession + + from deerflow.persistence.engine import get_engine + + if get_engine().dialect.name != "postgresql": + pytest.skip("PostgreSQL row-lock interleaving") + now = await completed_batch(repo) + await repo.claim_continuations(now=now, lease_owner="initial", lease_seconds=30) + await repo.finish_continuation("batch", lease_owner="initial", status="pending") + locked, release = asyncio.Event(), asyncio.Event() + original = AsyncSession.execute + claim_task = suppression = None + suppression_pid = None + + async def execute(session, statement, *args, **kwargs): + nonlocal suppression_pid + if asyncio.current_task() is suppression and suppression_pid is None: + suppression_pid = (await original(session, text("SELECT pg_backend_pid()"))).scalar_one() + result = await original(session, statement, *args, **kwargs) + if asyncio.current_task() is claim_task and getattr(statement, "is_update", False) and statement.table.name == "subagent_batches": + locked.set() + await release.wait() + return result + + monkeypatch.setattr(AsyncSession, "execute", execute) + claim_task = asyncio.create_task(repo.claim_continuations(now=now, lease_owner="claim", lease_seconds=30)) + try: + await asyncio.wait_for(locked.wait(), 3) + suppression = asyncio.create_task(repo.suppress_thread_continuations("thread")) + waiting = False + for _ in range(200): + async with repo._sf() as session: + waiting = await session.scalar(text("SELECT count(*) FROM pg_stat_activity WHERE pid=:pid AND wait_event_type='Lock'"), {"pid": suppression_pid}) + if waiting: + break + await asyncio.sleep(0.01) + assert waiting, "suppression must reach the competing row lock" + release.set() + await asyncio.wait_for(asyncio.gather(claim_task, suppression), 5) + assert (await repo.get_batch("batch", user_id="owner"))["continuation_status"] == "suppressed" + finally: + release.set() + for task in (claim_task, suppression): + if task is not None and not task.done(): + task.cancel() + await asyncio.gather(*(task for task in (claim_task, suppression) if task is not None), return_exceptions=True) + + +@pytest.mark.asyncio +async def test_schema_upgrade_and_rollback_preserve_existing_batch_work(repo): + import asyncio + + from alembic import command + from sqlalchemy import text + + from deerflow.persistence import bootstrap + from deerflow.persistence.engine import get_engine + + await completed_batch(repo) + engine = get_engine() + async with engine.connect() as connection: + schema = await connection.scalar(text("SELECT current_schema()")) if engine.dialect.name == "postgresql" else "" + config = bootstrap._get_alembic_config(engine, postgres_schema=schema) + await asyncio.to_thread(command.downgrade, config, "0051_merge_run_inputs_fencing") + async with engine.connect() as connection: + assert (await connection.execute(text("SELECT title FROM subagent_batches WHERE id='batch'"))).scalar_one() == "Summarize records" + assert (await connection.execute(text("SELECT status FROM subagent_batch_items WHERE batch_id='batch'"))).scalar_one() == "succeeded" + await asyncio.to_thread(command.upgrade, config, "head") + batch = await repo.get_batch("batch", user_id="owner") + assert batch["counts"]["succeeded"] == 1 + assert "continuation_status" not in batch diff --git a/backend/tests/test_batch_dependencies.py b/backend/tests/test_batch_dependencies.py new file mode 100644 index 0000000000..98f9f52520 --- /dev/null +++ b/backend/tests/test_batch_dependencies.py @@ -0,0 +1,150 @@ +"""Native batch dependencies use committed outcomes without occupying worker slots.""" + +import os +from datetime import UTC, datetime +from uuid import uuid4 + +import pytest +import pytest_asyncio +from sqlalchemy.schema import DropSchema + +from deerflow.config.database_config import DatabaseConfig +from deerflow.persistence.engine import close_engine, get_engine, get_session_factory, init_engine_from_config +from deerflow.persistence.subagent_batches import SubagentBatchRepository + + +@pytest_asyncio.fixture(params=["sqlite", "postgres"]) +async def repo(tmp_path, request): + config = DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)) + if request.param == "postgres": + uri = os.environ.get("TEST_POSTGRES_URI") + if not uri: + pytest.skip("requires TEST_POSTGRES_URI") + config = DatabaseConfig(backend="postgres", postgres_url=uri, postgres_schema="native_batch_" + uuid4().hex) + await init_engine_from_config(config) + try: + yield SubagentBatchRepository(get_session_factory()) + finally: + if config.backend == "postgres": + async with get_engine().begin() as connection: + await connection.execute(DropSchema(config.postgres_schema, cascade=True)) + await close_engine() + + +async def create(repo, items, **overrides): + values = dict( + batch_id="batch", + user_id="owner", + thread_id="thread", + run_id="run", + tool_call_id="call", + submission_key="run:call", + title="Compare source files", + subagent_type="general-purpose", + items=items, + max_live_items=3, + max_running_items=3, + max_attempts=1, + execution_spec={}, + ) + return await repo.create_batch(**(values | overrides)) + + +async def claim(repo): + return await repo.claim_items(now=datetime.now(UTC), lease_owner="worker", lease_seconds=120, limit=3) + + +async def finish(repo, item, *, succeeded=True, stop_reason=None, **overrides): + values = dict( + lease_owner="worker", + succeeded=succeeded, + result="A useful result" if succeeded else None, + result_preview="A useful result" if succeeded else None, + result_truncated=False, + error=None if succeeded else "Source unavailable", + stop_reason=stop_reason, + token_usage=None, + model_name="local-test", + completed_at=datetime.now(UTC), + ) + return await repo.finalize_item(item["id"], **(values | overrides)) + + +@pytest.mark.asyncio +async def test_ready_dependency_starts_before_unrelated_item_finishes_and_survives_reopen(repo): + await create( + repo, + [ + {"key": "d", "prompt": "Combine A and B", "depends_on": ["a", "b"]}, + {"key": "a", "prompt": "Read sales"}, + {"key": "b", "prompt": "Read risks"}, + {"key": "c", "prompt": "Wait for independent appendix"}, + ], + ) + running = {item["item_key"]: item for item in await claim(repo)} + assert set(running) == {"a", "b", "c"} + await finish(repo, running["a"]) + assert await claim(repo) == [] + await finish(repo, running["b"]) + reopened = SubagentBatchRepository(get_session_factory()) + ready = await claim(reopened) + assert [item["item_key"] for item in ready] == ["d"] + rows = await reopened.list_items("batch", user_id="owner") + assert next(item for item in rows if item["item_key"] == "c")["status"] == "leased" + assert next(item for item in rows if item["item_key"] == "d")["attempt"] == 1 + + +@pytest.mark.asyncio +async def test_failed_prerequisite_blocks_descendants_but_keeps_independent_work(repo): + await create( + repo, + [ + {"key": "a", "prompt": "Missing source"}, + {"key": "d", "prompt": "Requires A", "depends_on": ["a"]}, + {"key": "e", "prompt": "Requires D", "depends_on": ["d"]}, + {"key": "c", "prompt": "Independent work"}, + ], + ) + running = {item["item_key"]: item for item in await claim(repo)} + assert set(running) == {"a", "c"} + await finish(repo, running["a"], succeeded=False) + assert await claim(repo) == [] + await finish(repo, running["c"]) + batch = await repo.get_batch("batch", user_id="owner") + assert batch["counts"]["failed"] == 3 + assert batch["counts"]["succeeded"] == 1 + rows = await repo.list_items("batch", user_id="owner") + assert all(item["attempt"] == 0 for item in rows if item["item_key"] in {"d", "e"}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("dependencies", [["missing"], ["a"], ["b"]]) +async def test_invalid_dependencies_are_rejected_before_any_batch_is_saved(repo, dependencies): + items = [{"key": "a", "prompt": "A", "depends_on": dependencies}, {"key": "b", "prompt": "B", "depends_on": ["a"]}] + with pytest.raises(ValueError): + await create(repo, items) + assert await repo.get_batch("batch", user_id="owner") is None + + +@pytest.mark.asyncio +async def test_cancelled_batch_does_not_release_waiting_dependencies(repo): + await create(repo, [{"key": "a", "prompt": "A"}, {"key": "b", "prompt": "B", "depends_on": ["a"]}]) + running = await claim(repo) + assert len(running) == 1 + await repo.cancel_batch("batch", user_id="owner") + await finish(repo, running[0]) + assert await claim(repo) == [] + assert (await repo.get_batch("batch", user_id="owner"))["status"] == "cancelled" + + +@pytest.mark.asyncio +async def test_capped_attempt_can_retry_before_releasing_its_dependent(repo): + await create(repo, [{"key": "a", "prompt": "A"}, {"key": "b", "prompt": "Use A", "depends_on": ["a"]}], max_attempts=3) + first = (await claim(repo))[0] + await finish(repo, first, succeeded=False, stop_reason="turn_capped", error="No usable partial result") + retry = (await claim(repo))[0] + assert retry["item_key"] == "a" and retry["attempt"] == 2 + await finish(repo, retry) + ready = await claim(repo) + assert [item["item_key"] for item in ready] == ["b"] + assert ready[0]["attempt"] == 1 diff --git a/backend/tests/test_batch_followup_host.py b/backend/tests/test_batch_followup_host.py new file mode 100644 index 0000000000..ca122f36c6 --- /dev/null +++ b/backend/tests/test_batch_followup_host.py @@ -0,0 +1,412 @@ +"""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() diff --git a/backend/tests/test_batch_goal_wakeup.py b/backend/tests/test_batch_goal_wakeup.py new file mode 100644 index 0000000000..e9e0af5032 --- /dev/null +++ b/backend/tests/test_batch_goal_wakeup.py @@ -0,0 +1,59 @@ +"""A confirmed result can resolve external waiting without resetting Goal limits.""" + +import pytest +from langgraph.checkpoint.memory import InMemorySaver + +from deerflow.runtime.goal import build_goal_state, read_thread_goal, wake_goal_after_external_wait, write_thread_goal + + +async def waiting_goal(**changes): + goal = build_goal_state("Combine the completed batch", max_continuations=2) + goal["last_evaluation"] = {"satisfied": False, "blocker": "external_wait", "stand_down_reason": "blocked:external_wait"} + goal.update(changes) + saver = InMemorySaver() + await write_thread_goal(saver, "thread", goal, create_if_missing=True) + return saver, goal + + +@pytest.mark.asyncio +async def test_external_result_consumes_one_continuation_and_run_recovery_does_not_charge_it_twice(): + saver, goal = await waiting_goal() + assert await wake_goal_after_external_wait(saver, "thread", "wake-run", goal) + assert await wake_goal_after_external_wait(saver, "thread", "wake-run", goal) + current = await read_thread_goal(saver, "thread") + assert current["continuation_count"] == 1 + assert current["no_progress_count"] == 0 + assert current["last_batch_wakeup_run_id"] == "wake-run" + assert current["last_evaluation"] == goal["last_evaluation"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "changes", + [ + {"continuation_count": 2}, + {"no_progress_count": 2}, + {"last_evaluation": {"blocker": "needs_user_input", "stand_down_reason": "blocked:needs_user_input"}}, + {"last_evaluation": {"blocker": "external_wait", "stand_down_reason": "token_capped"}}, + ], +) +async def test_external_result_does_not_bypass_existing_limits_or_other_blockers(changes): + saver, goal = await waiting_goal(**changes) + assert not await wake_goal_after_external_wait(saver, "thread", "wake-run", goal) + assert await read_thread_goal(saver, "thread") == goal + + +@pytest.mark.asyncio +async def test_changed_goal_cannot_be_resumed_by_an_old_batch_event(): + saver, goal = await waiting_goal() + replacement = build_goal_state("A different task") + await write_thread_goal(saver, "thread", replacement) + assert not await wake_goal_after_external_wait(saver, "thread", "wake-run", goal) + assert await read_thread_goal(saver, "thread") == replacement + + +def test_external_goal_input_cannot_forge_the_wakeup_receipt(): + from app.gateway.services import normalize_input + + value = normalize_input({"goal": {"objective": "A task", "last_batch_wakeup_run_id": "forged"}}) + assert "last_batch_wakeup_run_id" not in value["goal"] diff --git a/backend/tests/test_batch_task_tool.py b/backend/tests/test_batch_task_tool.py index b69909547d..0319167482 100644 --- a/backend/tests/test_batch_task_tool.py +++ b/backend/tests/test_batch_task_tool.py @@ -164,6 +164,17 @@ async def test_batch_task_rejects_duplicate_item_keys_without_submitting(monkeyp submitter.submit.assert_not_awaited() +@pytest.mark.asyncio +async def test_cancel_batch_keeps_the_existing_submitter_signature(monkeypatch) -> None: + class ExistingSubmitter: + async def cancel_batch(self, *, batch_id, user_id): + return {"id": batch_id} + + monkeypatch.setattr(tool_module, "get_subagent_batch_submitter", lambda: ExistingSubmitter()) + result = await tool_module.cancel_batch.coroutine(runtime=_runtime(), batch_id="batch-1") + assert result == "Batch batch-1 cancellation requested." + + @pytest.mark.asyncio async def test_bound_batch_tools_use_the_explicit_submitter(monkeypatch) -> None: explicit = AsyncMock() @@ -283,6 +294,31 @@ def _explicit_submitter() -> AsyncMock: return explicit +@pytest.mark.asyncio +async def test_followup_captures_configurable_settings_with_context_precedence(monkeypatch): + _mock_subagent_catalog(monkeypatch) + submitter = _explicit_submitter() + monkeypatch.setattr(tool_module, "get_subagent_batch_submitter", lambda: submitter) + runtime = _runtime() + runtime.config["configurable"].update(thinking_enabled=False, reasoning_effort="high", max_concurrent_subagents=2, max_total_subagents=4) + runtime.context["reasoning_effort"] = "low" + await tool_module.batch_task.coroutine( + title="Meeting records", + items=[BatchTaskItem(key="a", prompt="Read")], + subagent_type="general-purpose", + continuation="Summarize", + runtime=runtime, + tool_call_id="batch-call", + ) + captured = submitter.submit.call_args.args[0].continuation_context["context"] + assert {key: captured[key] for key in ("thinking_enabled", "reasoning_effort", "max_concurrent_subagents", "max_total_subagents")} == { + "thinking_enabled": False, + "reasoning_effort": "low", + "max_concurrent_subagents": 2, + "max_total_subagents": 4, + } + + def test_bound_batch_tools_sync_path_uses_the_explicit_submitter(monkeypatch) -> None: """Sync invocation of a bound copy must not bypass the explicit submitter. diff --git a/backend/tests/test_migration_0047_mcp_oauth_tokens.py b/backend/tests/test_migration_0047_mcp_oauth_tokens.py index c928a5401c..c1c5ea6fb2 100644 --- a/backend/tests/test_migration_0047_mcp_oauth_tokens.py +++ b/backend/tests/test_migration_0047_mcp_oauth_tokens.py @@ -44,9 +44,10 @@ pytestmark = pytest.mark.asyncio async def test_0047_remains_in_the_single_migration_chain(): - assert bootstrap._get_head_revision() == "0051_merge_run_inputs_fencing" + assert bootstrap._get_head_revision() == "0052_batch_continuation" script = ScriptDirectory(str(bootstrap._MIGRATIONS_DIR)) - assert script.get_heads() == ["0051_merge_run_inputs_fencing"] + assert script.get_heads() == ["0052_batch_continuation"] + assert script.get_revision("0052_batch_continuation").down_revision == "0051_merge_run_inputs_fencing" assert set(script.get_revision("0051_merge_run_inputs_fencing").down_revision) == {"0049_run_inputs", "0050_history_seed_authority"} assert script.get_revision("0050_history_seed_authority").down_revision == "0049_run_producer_fencing" assert script.get_revision("0049_run_producer_fencing").down_revision == "0048_extension_execution_leases" diff --git a/backend/tests/test_migration_0050_history_seed_authority.py b/backend/tests/test_migration_0050_history_seed_authority.py index eb29dbbcbc..0ab03d17c8 100644 --- a/backend/tests/test_migration_0050_history_seed_authority.py +++ b/backend/tests/test_migration_0050_history_seed_authority.py @@ -36,7 +36,9 @@ async def test_existing_strict_database_refreshes_persisted_guard(seed_db, leadi authority = history_seed_authority(reservation, db.thread, seed_kind="checkpoint") with producer_scope(authority), pytest.raises(DBAPIError): await db.events.put_batch(batch) - await asyncio.to_thread(command.upgrade, cfg, "head") + # This fixture contains only the Run/history tables. Exercise 0050, + # not later migrations that require unrelated application tables. + await asyncio.to_thread(command.upgrade, cfg, "0050_history_seed_authority") with producer_scope(authority): rows = await db.events.put_batch(batch) assert [row["seq"] for row in rows] == ([1, 2, 3] if leading_ai else [1, 2, 3, 4]) diff --git a/backend/tests/test_migration_0051_run_inputs_fencing.py b/backend/tests/test_migration_0051_run_inputs_fencing.py index 0e0ce54fd2..a4b4896559 100644 --- a/backend/tests/test_migration_0051_run_inputs_fencing.py +++ b/backend/tests/test_migration_0051_run_inputs_fencing.py @@ -52,7 +52,7 @@ async def test_both_migration_branches_upgrade_to_one_head(tmp_path, backend, pr ) await asyncio.to_thread(command.upgrade, cfg, "head") async with engine.connect() as conn: - assert (await conn.execute(text("SELECT version_num FROM alembic_version"))).scalars().all() == ["0051_merge_run_inputs_fencing"] + assert (await conn.execute(text("SELECT version_num FROM alembic_version"))).scalars().all() == ["0052_batch_continuation"] columns = await conn.run_sync(lambda sync: {c["name"] for c in inspect(sync).get_columns("runs")}) assert {"run_inputs_json", "run_incarnation", "producer_epoch", "producer_state"} <= columns assert (await conn.execute(text("SELECT status,run_inputs_json,run_incarnation,producer_epoch,producer_state FROM runs WHERE run_id='old'"))).one() == ("success", None, None, None, None) diff --git a/backend/tests/test_subagent_batches_router.py b/backend/tests/test_subagent_batches_router.py index 8298e3feda..d306072eee 100644 --- a/backend/tests/test_subagent_batches_router.py +++ b/backend/tests/test_subagent_batches_router.py @@ -177,16 +177,28 @@ async def test_selected_item_route_returns_service_projection_and_scopes_thread( @pytest.mark.asyncio -async def test_cancel_requires_running_worker_and_exact_owner(monkeypatch) -> None: +@pytest.mark.parametrize(("continuation", "can_cancel"), [(None, False), ("pending", False), ("dispatched", False), ("dispatched", True)]) +async def test_cancel_requires_running_worker_and_exact_owner(monkeypatch, continuation, can_cancel) -> None: + from app.gateway.authz import AuthContext + repo = Repository() + repo.batch["continuation_status"] = continuation service = AsyncMock() service.cancel_batch.return_value = _batch(status="cancelled") monkeypatch.setattr(subagent_batches, "get_current_user", AsyncMock(return_value="user-1")) + request = _request(repo, service=service) + request.state = SimpleNamespace(auth_source="pat", auth=AuthContext(None, ["threads:write", *(["runs:cancel"] if can_cancel else [])])) + if continuation and not can_cancel: + with pytest.raises(HTTPException) as denied: + await subagent_batches.cancel_batch.__wrapped__(thread_id="thread-1", batch_id="batch-1", request=request) + assert denied.value.status_code == 403 + service.cancel_batch.assert_not_awaited() + return result = await subagent_batches.cancel_batch.__wrapped__( thread_id="thread-1", batch_id="batch-1", - request=_request(repo, service=service), + request=request, ) assert result["status"] == "cancelled" service.cancel_batch.assert_awaited_once_with(batch_id="batch-1", user_id="user-1") diff --git a/backend/tests/test_subagent_routing_prompt.py b/backend/tests/test_subagent_routing_prompt.py index 954df4dfe5..4064ba72d2 100644 --- a/backend/tests/test_subagent_routing_prompt.py +++ b/backend/tests/test_subagent_routing_prompt.py @@ -26,6 +26,15 @@ def test_routing_requires_clear_net_benefit(monkeypatch) -> None: assert "side-effect risk" in section +def test_durable_batch_guidance_allows_only_explicit_prerequisites(monkeypatch) -> None: + monkeypatch.setattr(prompt_module, "get_available_subagent_names", lambda: ["general-purpose"]) + section = prompt_module._build_subagent_section(3, batch_enabled=True) + assert "depends_on" in section + assert "must not depend on another item's output" not in section + assert "simultaneously runnable items independent" in section + assert "depends_on" not in prompt_module._build_subagent_section(3, batch_enabled=False) + + def test_hard_vetoes_apply_to_parallel_dispatch_not_single_agent_chains(monkeypatch) -> None: section = _build_section(monkeypatch) diff --git a/contracts/thread_origin_contract.json b/contracts/thread_origin_contract.json index 3d99a7df13..fa2be03d2b 100644 --- a/contracts/thread_origin_contract.json +++ b/contracts/thread_origin_contract.json @@ -7,6 +7,7 @@ "github", "im_channel", "mcp_notification", - "schedule" + "schedule", + "subagent_batch" ] } diff --git a/frontend/src/components/workspace/thread-subagent-batches.tsx b/frontend/src/components/workspace/thread-subagent-batches.tsx index d2284460b5..b8cf7eee22 100644 --- a/frontend/src/components/workspace/thread-subagent-batches.tsx +++ b/frontend/src/components/workspace/thread-subagent-batches.tsx @@ -209,6 +209,15 @@ function BatchCard({ )} + {batch.continuation_status && + batch.continuation_status !== "dispatched" && ( +
+ {batch.continuation_status === "pending" || + batch.continuation_status === "dispatching" + ? t.subagentBatches.followUpPending + : t.subagentBatches.followUpStopped} +
+ )}