deer-flow/backend/app/subagent_batches/continuation.py

"""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"])