deer-flow/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py

"""Explicit durable batch mode for many independent native-subagent items."""

from __future__ import annotations

import asyncio
import hashlib
import json
import re
from collections.abc import Callable
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
from pydantic import BaseModel, Field, StrictInt

from deerflow.authz.principal import normalize_authz_attributes
from deerflow.config.tool_output_config import ToolOutputConfig
from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_RUNTIME_KEY, execution_scope
from deerflow.mcp_scope import (
    THREAD_INCARNATION_CONTEXT_KEY,
    THREAD_INCARNATION_METADATA_GUARD_KEY,
    runtime_thread_incarnation,
)
from deerflow.runtime.user_context import resolve_runtime_user_id
from deerflow.subagents.batch_runtime import (
    BatchItemInput,
    BatchSubmitRequest,
    SubagentBatchSubmitter,
    get_subagent_batch_submitter,
)
from deerflow.subagents.registry import get_available_subagent_names, get_subagent_config
from deerflow.tools.sync import make_sync_tool_wrapper
from deerflow.tools.types import Runtime


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()
_explicit_batch_submitter: ContextVar[SubagentBatchSubmitter | None | object] = ContextVar(
    "deerflow_explicit_subagent_batch_submitter",
    default=_NO_EXPLICIT_BATCH_SUBMITTER,
)
_explicit_batch_app_config: ContextVar[Any | None] = ContextVar(
    "deerflow_explicit_subagent_batch_app_config",
    default=None,
)


def _batch_submitter() -> SubagentBatchSubmitter | None:
    explicit = _explicit_batch_submitter.get()
    if explicit is not _NO_EXPLICIT_BATCH_SUBMITTER:
        return cast(SubagentBatchSubmitter | None, explicit)
    return get_subagent_batch_submitter()


def _batch_app_config(runtime: Runtime) -> Any | None:
    explicit = _explicit_batch_app_config.get()
    if explicit is not None:
        return explicit
    context = runtime.context if runtime is not None and isinstance(runtime.context, dict) else {}
    return context.get("app_config")


def _bind_batch_tool(
    tool,
    submitter_provider: Callable[[], SubagentBatchSubmitter | None],
    app_config: Any | None,
):
    original_coroutine = tool.coroutine
    if original_coroutine is None:  # pragma: no cover - all batch tools are async
        raise RuntimeError(f"{tool.name} has no async implementation")

    async def bound_coroutine(**kwargs):
        submitter_token = _explicit_batch_submitter.set(submitter_provider())
        config_token = _explicit_batch_app_config.set(app_config)
        try:
            return await original_coroutine(**kwargs)
        finally:
            _explicit_batch_app_config.reset(config_token)
            _explicit_batch_submitter.reset(submitter_token)

    # The source tool may carry a sync func around the unbound coroutine (set
    # in place by _ensure_sync_invocable_tool) or none at all; either way the
    # copy's sync path must go through the bound coroutine.
    return tool.model_copy(
        update={
            "coroutine": bound_coroutine,
            "func": make_sync_tool_wrapper(bound_coroutine, tool.name),
        }
    )


def bind_batch_tools(
    submitter: SubagentBatchSubmitter | None = None,
    *,
    submitter_provider: Callable[[], SubagentBatchSubmitter | None] | None = None,
    app_config: Any | None = None,
):
    """Return batch tools bound to an explicit SDK runtime submitter.

    A provider preserves runtime lifecycle semantics for already-compiled
    graphs: after their owned worker stops, the tools report unavailable and
    never fall through to another application's process-global submitter.
    """

    if (submitter is None) == (submitter_provider is None):
        raise ValueError("Provide exactly one of submitter or submitter_provider")
    provider = submitter_provider if submitter_provider is not None else lambda: submitter

    return tuple(_bind_batch_tool(tool, provider, app_config) for tool in (batch_task, batch_status, cancel_batch, read_batch_result))


def _result(tool_call_id: str, *, content: str, batch: dict[str, Any] | None = None, error: bool = False) -> Command:
    metadata: dict[str, Any] = {"subagent_batch_error": error}
    if batch is not None:
        metadata.update(
            {
                "subagent_batch_id": batch["id"],
                "subagent_batch_status": batch["status"],
                "subagent_batch_total_items": batch["total_items"],
            }
        )
    return Command(
        update={
            "messages": [
                ToolMessage(
                    content=content,
                    tool_call_id=tool_call_id,
                    name="batch_task",
                    status="error" if error else "success",
                    additional_kwargs=metadata,
                )
            ]
        }
    )


def _merge_skill_allowlists(parent: list[str] | None, child: list[str] | None) -> list[str] | None:
    if parent is None:
        return child
    if child is None:
        return list(parent)
    allowed = set(parent)
    return [name for name in child if name in allowed]


@tool("batch_task", parse_docstring=True)
async def batch_task(
    runtime: Runtime,
    title: str,
    items: list[BatchTaskItem],
    subagent_type: str,
    tool_call_id: Annotated[str, InjectedToolCallId],
    max_live_items: int | None = None,
    max_running_items: int | None = None,
    continuation: str | None = None,
) -> Command:
    """Submit durable native-subagent items with optional explicit prerequisites.

    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.

    Each item may carry ``acceptance_criteria`` using the same deterministic
    checks as ``task``: ``file:<path> exists|non-empty``, ``file_written:<path>``,
    or ``tests_passed:<command>`` against recorded execution evidence.
    Other conditions are UNVERIFIED. Item queries and JSONL exports include the
    separate acceptance verdict; ``succeeded`` only means execution completed.
    Retain useful results, repair unmet conditions, and verify consequential
    unknowns or preserve uncertainty. Acceptance never triggers automatic retries.
    JSON deliverables can opt into ``file:<path> json-valid``: validate complete UTF-8
    JSON syntax up to 50,000 bytes, rejecting NaN/Infinity. Oversize files or incomplete
    reads return UNVERIFIED. No schema or business validation is performed.

    Args:
        title: Short batch name shown to the user.
        items: Stable item keys, self-contained prompts, and optional per-item acceptance_criteria.
        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:
        return _result(
            tool_call_id,
            content="Durable subagent batches are unavailable. Enable subagent_batches with a SQL database and restart Gateway.",
            error=True,
        )
    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)

    context = runtime.context if runtime is not None and isinstance(runtime.context, dict) else {}
    if context.get(THREAD_INCARNATION_METADATA_GUARD_KEY) is True:
        runtime_thread_incarnation(runtime)
    metadata = runtime.config.get("metadata", {}) if runtime is not None else {}
    app_config = _batch_app_config(runtime)
    allowed_subagents = metadata.get("allowed_subagents")
    available = get_available_subagent_names(app_config=app_config, allowed_subagents=allowed_subagents)
    config = get_subagent_config(subagent_type, app_config=app_config)
    if config is None or subagent_type not in available:
        names = ", ".join(available) if available else "none"
        return _result(
            tool_call_id,
            content=f"Unknown or disallowed subagent type {subagent_type!r}. Available: {names}",
            error=True,
        )

    parent_skills = metadata.get("available_skills")
    if parent_skills is not None:
        config = replace(config, skills=_merge_skill_allowlists(list(parent_skills), config.skills))

    thread_id = context.get("thread_id") or runtime.config.get("configurable", {}).get("thread_id")
    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 = {
        "subagent_config": {**asdict(config), "prompt_overlay": config.prompt_overlay.model_dump()},
        "parent_model": metadata.get("model_name"),
        "tool_groups": metadata.get("tool_groups"),
        "mcp_plugins": metadata.get("mcp_plugins"),
        "user_role": context.get("user_role"),
        "oauth_provider": context.get("oauth_provider"),
        "oauth_id": context.get("oauth_id"),
        "channel_user_id": context.get("channel_user_id"),
        "is_internal": context.get("is_internal") is True,
        "authz_attributes": normalize_authz_attributes(context.get("authz_attributes")),
    }
    if THREAD_INCARNATION_CONTEXT_KEY in context:
        execution_spec[THREAD_INCARNATION_CONTEXT_KEY] = context[THREAD_INCARNATION_CONTEXT_KEY]
    if KNOWLEDGE_SCOPE_RUNTIME_KEY in context:
        execution_spec["knowledge_scope"] = execution_scope(context[KNOWLEDGE_SCOPE_RUNTIME_KEY])
    try:
        batch = await submitter.submit(
            BatchSubmitRequest(
                user_id=user_id,
                thread_id=str(thread_id),
                run_id=str(run_id) if run_id else None,
                tool_call_id=tool_call_id,
                submission_key=submission_key,
                title=title.strip()[:256] or "Subagent batch",
                subagent_type=subagent_type,
                items=[cast(BatchItemInput, item.model_dump(exclude_none=True)) for item in items],
                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:
        return _result(tool_call_id, content=f"Batch submission failed: {exc}", error=True)
    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. "
            + ("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."
        ),
    )


@tool("batch_status", parse_docstring=True)
async def batch_status(runtime: Runtime, batch_id: str) -> str:
    """Return a compact durable batch progress snapshot.

    Counts describe execution status, not acceptance. Inspect item queries or
    JSONL exports for the recorded acceptance criteria and verdicts.

    Args:
        batch_id: Server batch identifier returned by ``batch_task``.
    """
    submitter = _batch_submitter()
    if submitter is None:
        return "Durable subagent batches are unavailable."
    batch = await submitter.get_batch(batch_id=batch_id, user_id=resolve_runtime_user_id(runtime))
    if batch is None:
        return "Batch not found."
    return json.dumps(
        {
            "batch_id": batch["id"],
            "status": batch["status"],
            "total_items": batch["total_items"],
            "counts": batch["counts"],
        },
        ensure_ascii=False,
    )


@tool("cancel_batch", parse_docstring=True)
async def cancel_batch(runtime: Runtime, batch_id: str) -> str:
    """Cancel pending and running work in one durable subagent batch.

    Args:
        batch_id: Server batch identifier returned by ``batch_task``.
    """
    submitter = _batch_submitter()
    if submitter is None:
        return "Durable subagent batches are unavailable."
    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."


def _batch_result_window(item: dict[str, Any], offset: int, max_chars: int, expected_revision: str | None, response_limit: int) -> str:
    # Criteria, errors and reports are all untrusted data; bound the entire
    # document window instead of allowing metadata to bypass the read limit.
    document = json.dumps(item, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
    revision = hashlib.sha256(document.encode("utf-8")).hexdigest()
    if expected_revision is not None and revision != expected_revision:
        return json.dumps({"status": "restart_required"})
    if offset > len(document):
        return json.dumps({"status": "invalid_request"})
    end = min(offset + max_chars, len(document))

    def render(stop: int) -> str:
        return json.dumps(
            {
                "status": "ok",
                "position": item["position"],
                "next_position": item["position"] + 1,
                "revision": revision,
                "offset": offset,
                "next_offset": stop if stop < len(document) else None,
                "total_chars": len(document),
                "content": document[offset:stop],
            },
            ensure_ascii=False,
        ).replace("<", "\\u003c")

    # Account for the complete escaped response, not just the document slice.
    # Keep continuation inline under the host's externalization/fallback caps.
    response = render(end)
    if len(response) > response_limit:
        if len(render(min(offset + 1, len(document)))) > response_limit:
            return json.dumps({"status": "budget_too_small"})
        low, high = offset, end
        while low < high:
            middle = (low + high + 1) // 2
            if len(render(middle)) <= response_limit:
                low = middle
            else:
                high = middle - 1
        response = render(low)
    return response


def _batch_result_response_limit(runtime: Runtime) -> int:
    # Keep this inline cap aligned with tool_output_budget_middleware's
    # _effective_trigger / _budget_content, consumed by _patch_tool_message.
    # Changes there must also update this cap and the reader budget regressions.
    config = getattr(_batch_app_config(runtime), "tool_output", None)
    if not isinstance(config, ToolOutputConfig):
        config = ToolOutputConfig()
    if not config.enabled or "read_batch_result" in config.exempt_tools:
        return 10_000
    limits = [limit for limit in (config.tool_overrides.get("read_batch_result", config.externalize_min_chars), config.fallback_max_chars) if limit > 0]
    return min(10_000, *limits) if limits else 10_000


@tool("read_batch_result", parse_docstring=True)
async def read_batch_result(runtime: Runtime, batch_id: str, position: StrictInt = 0, offset: StrictInt = 0, max_chars: StrictInt = 4000, expected_revision: str | None = None) -> str:
    """Read one durable batch item as bounded, untrusted data in this thread.

    Use only for explicitly requested result inspection or synthesis. Never
    wait or poll for completion, and do not ingest every item automatically.
    The JSON response contains a window of a JSON document with the stored
    report, execution state and separate acceptance criteria/verdict. Concatenate
    content windows before parsing; a window alone need not be valid JSON.
    Follow next_offset with the returned revision as expected_revision. On
    restart_required discard earlier windows and start at zero. End of one
    document is next_offset=null; next_position selects another item, up to the
    total_items reported by batch_status. Stored result_truncated means the
    worker already capped the report; continuation cannot recover that suffix.
    Execution succeeded is not acceptance, and unchecked criteria are UNVERIFIED.
    Windows may be shorter to fit the host's output budget. budget_too_small
    means the operator must raise the read_batch_result output budget; stop reading.

    Args:
        batch_id: Server identifier returned by batch_task.
        position: Stable zero-based item position, not a status-filtered index.
        offset: Character offset in the serialized result document.
        max_chars: Window character limit, between 1 and 8192 (default 4000).
        expected_revision: Previous revision required when offset is nonzero.
    """
    # 100_000 is SubagentBatchesConfig.max_items_per_batch's schema ceiling.
    # Keep these bounds in lockstep; use the ceiling rather than today's configured
    # cap so lowering the submission limit does not hide existing batch items.
    if any(type(value) is not int for value in (position, offset, max_chars)) or not 0 <= position < 100_000 or offset < 0 or not 1 <= max_chars <= 8192:
        return json.dumps({"status": "invalid_request"})
    if expected_revision is not None and (not isinstance(expected_revision, str) or re.fullmatch(r"[a-f0-9]{64}", expected_revision) is None):
        return json.dumps({"status": "invalid_request"})
    if offset and expected_revision is None:
        return json.dumps({"status": "invalid_request"})
    context = runtime.context if isinstance(runtime.context, dict) else {}
    configurable = (runtime.config or {}).get("configurable", {})
    thread_id = context.get("thread_id") or configurable.get("thread_id")
    if not isinstance(thread_id, str) or not thread_id:
        return json.dumps({"status": "invalid_request"})
    submitter = _batch_submitter()
    if submitter is None or not callable(getattr(submitter, "read_batch_item", None)):
        return json.dumps({"status": "unavailable"})
    item = await submitter.read_batch_item(batch_id=batch_id, user_id=resolve_runtime_user_id(runtime), thread_id=thread_id, position=position)
    if item is None:
        return json.dumps({"status": "not_found"})
    # Reports can reach the configured million-character cap; serializing and
    # hashing the snapshot must not run on the streaming event loop.
    return await asyncio.to_thread(_batch_result_window, item, offset, max_chars, expected_revision, _batch_result_response_limit(runtime))