deer-flow/backend/packages/harness/deerflow/subagents/batch_service.py

from __future__ import annotations

import asyncio
import logging
import socket
import uuid
from datetime import UTC, datetime
from typing import Any

from deerflow.config.app_config import AppConfig, get_app_config
from deerflow.config.subagent_batches_config import SubagentBatchesConfig
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
from deerflow.subagents.batch_acceptance import check_batch_acceptance
from deerflow.subagents.batch_runtime import BatchSubmitRequest
from deerflow.subagents.capacity import SubagentExecutionCapacity
from deerflow.subagents.config import SubagentConfig, resolve_subagent_model_name
from deerflow.subagents.executor import (
    SubagentExecutor,
    SubagentStatus,
    cleanup_background_task,
    get_background_task_result,
    request_cancel_background_task,
)
from deerflow.utils.assembly_io import run_assembly

logger = logging.getLogger(__name__)


def _usage(records: list[dict[str, Any]] | None) -> dict[str, int] | None:
    if not records:
        return None
    return {
        "input_tokens": sum(int(row.get("input_tokens") or 0) for row in records),
        "output_tokens": sum(int(row.get("output_tokens") or 0) for row in records),
        "total_tokens": sum(int(row.get("total_tokens") or 0) for row in records),
    }


class SubagentBatchService:
    """Lease, execute, and recover durable native-subagent batch items."""

    def __init__(
        self,
        *,
        repository,
        config: SubagentBatchesConfig,
        runtime_config: SubagentRuntimeConfig,
        app_config: AppConfig | None = None,
        execution_capacity: SubagentExecutionCapacity | None = None,
    ) -> None:
        self._repository = repository
        self._config = config
        self._runtime_config = runtime_config
        self._app_config = app_config
        self._execution_capacity = execution_capacity
        self._lease_owner = f"{socket.gethostname()}:{uuid.uuid4().hex}"
        self._stop = asyncio.Event()
        self._poller: asyncio.Task[None] | None = None
        self._executions: dict[str, asyncio.Task[None]] = {}
        self._execution_ids: dict[str, str] = {}
        self._item_batches: dict[str, str] = {}

    async def start(self) -> None:
        if self._poller is not None:
            return
        self._stop.clear()
        self._poller = asyncio.create_task(self._run(), name="subagent-batch-poller")

    async def stop(self) -> None:
        self._stop.set()
        poller = self._poller
        self._poller = None
        if poller is not None:
            poller.cancel()
            await asyncio.gather(poller, return_exceptions=True)
        execution_ids = list(self._execution_ids.values())
        for execution_id in execution_ids:
            request_cancel_background_task(execution_id)
        tasks = list(self._executions.values())
        for task in tasks:
            task.cancel()
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)
        self._executions.clear()
        self._execution_ids.clear()
        self._item_batches.clear()

    async def _run(self) -> None:
        while not self._stop.is_set():
            try:
                await self.run_once(now=datetime.now(UTC))
            except asyncio.CancelledError:
                raise
            except Exception:
                logger.exception("Subagent batch scheduler pass failed")
            try:
                await asyncio.wait_for(
                    self._stop.wait(),
                    timeout=self._config.poll_interval_seconds,
                )
            except TimeoutError:
                pass

    async def run_once(self, *, now: datetime) -> None:
        available = max(0, self._runtime_config.max_running - len(self._executions))
        if available <= 0:
            return
        items = await self._repository.claim_items(
            now=now,
            lease_owner=self._lease_owner,
            lease_seconds=self._config.lease_seconds,
            limit=available,
        )
        for item in items:
            item_id = item["id"]
            if item_id in self._executions:
                continue
            task = asyncio.create_task(
                self._execute_item(item),
                name=f"subagent-batch-item-{item_id}",
            )
            self._executions[item_id] = task
            task.add_done_callback(
                lambda _task, current_id=item_id: self._executions.pop(
                    current_id,
                    None,
                )
            )

    async def submit(self, request: BatchSubmitRequest) -> dict[str, Any]:
        total = len(request.items)
        if total < 1 or total > self._config.max_items_per_batch:
            raise ValueError(f"Batch item count must be between 1 and {self._config.max_items_per_batch}")
        max_live = request.max_live_items or self._config.default_max_live_items
        max_running = request.max_running_items or self._config.default_max_running_items
        if not 1 <= max_live <= self._config.max_live_items_per_batch:
            raise ValueError(f"max_live_items must be between 1 and {self._config.max_live_items_per_batch}")
        if not 1 <= max_running <= self._config.max_running_items_per_batch:
            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")
        return await self._repository.create_batch(
            batch_id=f"subagent-batch-{uuid.uuid4().hex}",
            user_id=request.user_id,
            thread_id=request.thread_id,
            run_id=request.run_id,
            tool_call_id=request.tool_call_id,
            submission_key=request.submission_key,
            title=request.title,
            subagent_type=request.subagent_type,
            items=request.items,
            max_live_items=max_live,
            max_running_items=max_running,
            max_attempts=self._config.max_attempts,
            execution_spec=request.execution_spec,
        )

    async def get_batch(
        self,
        *,
        batch_id: str,
        user_id: str,
    ) -> dict[str, Any] | None:
        return await self._repository.get_batch(batch_id, user_id=user_id)

    async def cancel_batch(
        self,
        *,
        batch_id: str,
        user_id: str,
    ) -> dict[str, Any] | None:
        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)
        # 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.
        return batch

    async def _execute_item(self, item: dict[str, Any]) -> None:
        item_id = item["id"]
        execution_id: str | None = None
        try:
            batch = item["batch"]
            self._item_batches[item_id] = batch["id"]
            spec = batch["execution_spec"]
            config = SubagentConfig(**spec["subagent_config"])
            app_config = self._app_config or get_app_config()
            from deerflow.tools import get_available_tools

            effective_model = resolve_subagent_model_name(
                config,
                spec.get("parent_model"),
                app_config=app_config,
            )
            # Assemble off-loop: tool assembly may block on MCP cache
            # initialization, which must not stall the calling event loop (issue #5172).
            tools = await run_assembly(
                get_available_tools,
                groups=spec.get("tool_groups"),
                model_name=effective_model,
                subagent_enabled=False,
                include_upload_tool=False,
                app_config=app_config,
            )
            # Revalidate durable state before launching: cancel_batch may have
            # terminalized this item (or its lease may have been lost) while
            # assembly blocked in the worker thread — the poll loop's checks
            # only start after execute_async(), so launching without this
            # check would run work the user already cancelled.
            lease = await self._repository.renew_item_lease(
                item_id,
                lease_owner=self._lease_owner,
                lease_seconds=self._config.lease_seconds,
                now=datetime.now(UTC),
            )
            if not lease["valid"]:
                logger.info(
                    "Durable batch item %s cancelled or lease lost during tool assembly; skipping launch",
                    item_id,
                )
                return
            executor = SubagentExecutor(
                config=config,
                tools=tools,
                app_config=app_config,
                parent_model=spec.get("parent_model"),
                thread_id=batch["thread_id"],
                user_id=batch["user_id"],
                user_role=spec.get("user_role"),
                oauth_provider=spec.get("oauth_provider"),
                oauth_id=spec.get("oauth_id"),
                run_id=batch.get("run_id"),
                channel_user_id=spec.get("channel_user_id"),
                is_internal=spec.get("is_internal") is True,
                authz_attributes=spec.get("authz_attributes"),
                execution_capacity=self._execution_capacity,
                acceptance_criteria=item.get("acceptance_criteria"),
            )
            prompt = f"Durable batch item key: {item['item_key']}\nThis item may be retried after a worker crash. Keep side effects idempotent and use the item key as the idempotency identity.\n\n{item['prompt']}"
            execution_id = executor.execute_async(prompt, task_id=item_id)
            self._execution_ids[item_id] = execution_id
            marked_running = False
            renew_every = max(1.0, self._config.lease_seconds / 3)
            status_poll_every = min(
                self._config.poll_interval_seconds,
                renew_every,
            )
            loop = asyncio.get_running_loop()
            next_renew_at = loop.time() + renew_every
            while True:
                result = get_background_task_result(execution_id)
                if result is None:
                    raise RuntimeError("Native subagent execution disappeared")
                if result.status is SubagentStatus.RUNNING and not marked_running:
                    marked_running = await self._repository.mark_item_running(
                        item_id,
                        lease_owner=self._lease_owner,
                        now=datetime.now(UTC),
                    )
                    if not marked_running:
                        request_cancel_background_task(execution_id)
                if result.status.is_terminal:
                    break
                now_monotonic = loop.time()
                if now_monotonic >= next_renew_at:
                    lease = await self._repository.renew_item_lease(
                        item_id,
                        lease_owner=self._lease_owner,
                        lease_seconds=self._config.lease_seconds,
                        now=datetime.now(UTC),
                    )
                    next_renew_at = loop.time() + renew_every
                    if not lease["valid"]:
                        request_cancel_background_task(execution_id)
                try:
                    until_renew = max(0.0, next_renew_at - loop.time())
                    await asyncio.wait_for(
                        self._stop.wait(),
                        timeout=min(status_poll_every, until_renew),
                    )
                    if self._stop.is_set():
                        raise asyncio.CancelledError
                except TimeoutError:
                    pass

            raw_result = result.result or ""
            if getattr(result, "admission_failure", False):
                await self._repository.requeue_item_after_admission_failure(
                    item_id,
                    lease_owner=self._lease_owner,
                    error=result.error,
                    now=datetime.now(UTC),
                )
                return
            truncated = len(raw_result) > self._config.max_result_chars
            stored_result = raw_result[: self._config.max_result_chars] if raw_result else None
            preview = raw_result[: self._config.result_preview_max_chars] if raw_result else None
            acceptance_verdict = None
            if result.status is SubagentStatus.COMPLETED and item.get("acceptance_criteria"):
                try:
                    valid, acceptance_verdict = await self._check_acceptance_with_lease(item, result, app_config)
                    if not valid:
                        return
                except Exception:
                    # Advisory like ordinary task acceptance: an unavailable
                    # checker must not discard useful work or trigger a retry.
                    logger.warning("Batch acceptance check failed; result remains unchecked (item_id=%s)", item_id, exc_info=True)
            await self._repository.finalize_item(
                item_id,
                lease_owner=self._lease_owner,
                succeeded=result.status is SubagentStatus.COMPLETED,
                result=stored_result,
                result_preview=preview,
                result_truncated=truncated,
                error=result.error,
                stop_reason=result.stop_reason,
                token_usage=_usage(result.token_usage_records),
                model_name=effective_model,
                completed_at=datetime.now(UTC),
                acceptance_verdict=acceptance_verdict,
            )
        except asyncio.CancelledError:
            if execution_id is not None:
                request_cancel_background_task(execution_id)
            # Do not finalize on process shutdown. The durable lease expires and
            # another worker reclaims the same stable item key.
            raise
        except Exception as exc:
            logger.exception(
                "Durable subagent batch item failed (item_id=%s)",
                item_id,
            )
            await self._repository.finalize_item(
                item_id,
                lease_owner=self._lease_owner,
                succeeded=False,
                result=None,
                result_preview=None,
                result_truncated=False,
                error=str(exc)[:4_000],
                stop_reason=None,
                token_usage=None,
                model_name=None,
                completed_at=datetime.now(UTC),
            )
        finally:
            self._execution_ids.pop(item_id, None)
            self._item_batches.pop(item_id, None)
            if execution_id is not None:
                cleanup_background_task(execution_id)

    async def _check_acceptance_with_lease(self, item, result, app_config):
        """Keep a completed execution leased until its advisory check drains."""

        async def renew():
            lease = await self._repository.renew_item_lease(
                item["id"],
                lease_owner=self._lease_owner,
                lease_seconds=self._config.lease_seconds,
                now=datetime.now(UTC),
            )
            return lease["valid"]

        if not await renew():
            return False, None
        check = asyncio.create_task(
            check_batch_acceptance(
                item["acceptance_criteria"],
                batch=item["batch"],
                app_config=app_config,
                bash_executions=getattr(result, "bash_executions", None),
            )
        )
        try:
            while True:
                done, _ = await asyncio.wait({check}, timeout=max(1.0, self._config.lease_seconds / 3))
                if done:
                    return True, check.result()
                if not await renew():
                    return False, None
        finally:
            if not check.done():
                check.cancel()
            # The checklist's sandbox offload drains before releasing its
            # holder, even when shutdown or a lost lease cancels this task.
            await asyncio.gather(check, return_exceptions=True)