deer-flow/backend/packages/harness/deerflow/agents/middlewares/tool_progress_middleware.py

"""Middleware for task-level tool call progress tracking with a state machine.

Implements RFC #3177: structured tool result signals drive a per-(thread, tool)
state machine that detects stagnation and repetition, injects hints early
(WARNED), and hard-blocks the tool when it has stopped producing value (BLOCKED).

Architecture:
  ToolProgressMiddleware (outer)
    └── handler → ToolErrorHandlingMiddleware (inner) → actual tool
                                                              ↓
  ToolProgressMiddleware reads deerflow_tool_meta from the normalized result

State machine transitions per (thread_id, tool_name):
  ACTIVE → WARNED (at stagnation_threshold problems)
  Any problem-free call resets consecutive_problems=0 and reverts to ACTIVE.

  Whether WARNED can escalate to BLOCKED depends on recoverable_by_model:
  - recoverable_by_model=True  (no_results, not_found, permission, Jaccard-duplicate success):
      WARNED is terminal. The model received a hint and is expected to change strategy;
      blocking would prevent a legitimate retry with different parameters.
  - recoverable_by_model=False, action≠stop (transient, rate_limited):
      WARNED → BLOCKED after warn_escalation_count more problems. The model cannot fix
      these by retrying the same tool, so hard-blocking conserves API calls.
  - recoverable_by_model=False, action=stop (auth, config, internal):
      Immediately BLOCKED on the first occurrence — no retry can help.

Division of labor with LoopDetectionMiddleware (middleware position 23):
  ToolProgressMiddleware (position 10) is a result-quality guard — it fires
  after a tool executes, inspects what came back, and blocks *specific tools*
  that have stopped producing new information.

  LoopDetectionMiddleware is a call-pattern guard — it fires after the model
  responds (before tools execute), inspects the tool_calls signature in the
  AIMessage, and forces the *whole turn* to stop when the model keeps issuing
  the same calls regardless of results.

  They are complementary, not competing:
  - ToolProgressMiddleware is fine-grained (per-tool BLOCK, other tools normal).
  - LoopDetectionMiddleware is coarse-grained (strips all tool_calls, ends turn).
  - Both can inject HumanMessage hints in the same model call without conflict;
    the model sees both sets of hints and can reason about them.
  - If LoopDetectionMiddleware hard-stops (strips tool_calls), no wrap_tool_call
    is issued so ToolProgressMiddleware never fires — there is no double-stop.
  - If ToolProgressMiddleware BLOCKs a tool (returns an error ToolMessage),
    the model still makes a tool call that LoopDetectionMiddleware tracks; both
    continue to operate on their own independent state.
"""

from __future__ import annotations

import logging
import re
import threading
from collections import OrderedDict, defaultdict
from collections.abc import Awaitable, Callable, Sequence
from dataclasses import dataclass, field, replace
from typing import TYPE_CHECKING, Literal, override

from langchain.agents import AgentState
from langchain.agents.middleware import AgentMiddleware
from langchain.agents.middleware.types import ModelCallResult, ModelRequest, ModelResponse
from langchain_core.messages import HumanMessage, ToolMessage
from langgraph.prebuilt.tool_node import ToolCallRequest
from langgraph.runtime import Runtime
from langgraph.types import Command

from deerflow.agents.middlewares.tool_result_meta import TOOL_META_KEY, ToolResultMeta

if TYPE_CHECKING:
    from deerflow.config.tool_progress_config import ToolProgressConfig

logger = logging.getLogger(__name__)

_MAX_PENDING_PER_RUN = 3
# Jaccard word-set computation is capped to avoid O(n) regex work on very large tool results.
_MAX_CONTENT_FOR_WORDSET = 8192


# ---------------------------------------------------------------------------
# State data structures


@dataclass(slots=True)
class ToolPhaseState:
    """Per (thread_id, tool_name) tracking state."""

    phase: Literal["active", "warned", "blocked"] = "active"
    consecutive_problems: int = 0
    block_reason: str | None = None
    # Immutable tuple so that dataclasses.replace() calls that omit recent_word_sets
    # (problem paths) cannot accidentally share a mutable list between the old and new
    # state objects and cause silent cross-state corruption via .append().
    recent_word_sets: tuple[frozenset[str], ...] = field(default_factory=tuple)


# ---------------------------------------------------------------------------
# Content helpers


def word_set(content: str) -> frozenset[str]:
    """Extract lowercase words of length >= 3 for Jaccard similarity.

    Content is capped at _MAX_CONTENT_FOR_WORDSET chars to bound memory and CPU cost on
    large tool results (e.g. web pages).  Tail content beyond the cap is omitted from the
    set, which is acceptable because duplicate-detection is a heuristic, not a guarantee.
    """
    return frozenset(re.findall(r"\b\w{3,}\b", content[:_MAX_CONTENT_FOR_WORDSET].lower()))


def is_near_duplicate(
    current: frozenset[str],
    recent: Sequence[frozenset[str]],
    threshold: float,
    min_words: int,
) -> bool:
    """Return True if current is similar to any of the last 3 recent word sets."""
    if len(current) < min_words:
        return False
    for prev in recent[-3:]:
        if len(prev) < min_words:
            continue
        union = len(current | prev)
        if union == 0:
            continue
        if len(current & prev) / union >= threshold:
            return True
    return False


def _message_content_str(msg: ToolMessage) -> str:
    return msg.content if isinstance(msg.content, str) else ""


def _parse_tool_meta(meta_dict: object) -> ToolResultMeta | None:
    """Safely deserialize a ToolResultMeta from a raw dict; returns None on schema mismatch."""
    if not isinstance(meta_dict, dict):
        return None
    try:
        return ToolResultMeta(**meta_dict)
    except TypeError:
        logger.warning("Unexpected tool meta schema, skipping progress tracking: %s", meta_dict)
        return None


# ---------------------------------------------------------------------------
# Hint / block reason formatting


def _format_hint(meta: ToolResultMeta) -> str:
    action_map = {
        "rewrite_query": "Try rephrasing your search query with different keywords or approach.",
        "try_alternative": "Consider using a different tool or strategy.",
        "summarize": "Consider summarizing your current findings and moving forward.",
        "stop": "Do not retry this operation — it is not recoverable.",
        # Near-duplicate success results: recommended_next_action is "continue" by default,
        # but the model should still change strategy to avoid re-fetching the same content.
        "continue": "Try rephrasing your query or using a different search term.",
    }
    base = {
        "no_results": "[PROGRESS HINT] Your search returned no results.",
        "not_found": "[PROGRESS HINT] The resource was not found repeatedly.",
        "rate_limited": "[PROGRESS HINT] The tool is being rate-limited.",
        "transient": "[PROGRESS HINT] The tool encountered repeated transient failures.",
        "partial_success": "[PROGRESS HINT] The tool has returned incomplete results multiple times.",
        # Jaccard near-duplicate success: the tool is returning the same content repeatedly.
        "success": "[PROGRESS HINT] The tool is returning duplicate results.",
    }.get(
        meta.error_type or meta.status,
        "[PROGRESS HINT] The tool is not producing new information.",
    )
    suffix = action_map.get(meta.recommended_next_action, "")
    return f"{base} {suffix}".strip()


def _block_reason(meta: ToolResultMeta) -> str:
    return {
        "no_results": "Repeated no-results — rewrite your query or try a different tool.",
        "not_found": "Repeated not-found — rewrite your query or try a different resource.",
        "rate_limited": "Repeated rate-limiting — summarize current findings and proceed.",
        "transient": "Repeated transient failures — try a different approach.",
        "auth": "Authentication failure — this tool cannot be used.",
        "config": "Tool is not configured — this tool cannot be used.",
        "internal": "Repeated internal errors — this tool is unavailable.",
    }.get(
        meta.error_type or "",
        "Tool has not produced new information after multiple attempts — summarize and move on.",
    )


# ---------------------------------------------------------------------------
# Middleware


class ToolProgressMiddleware(AgentMiddleware[AgentState]):
    """State-machine-based tool stagnation guard (RFC #3177)."""

    def __init__(
        self,
        *,
        stagnation_threshold: int = 3,
        warn_escalation_count: int = 2,
        inject_assessment: bool = True,
        jaccard_threshold: float = 0.8,
        min_words: int = 10,
        exempt_tools: set[str] | None = None,
        max_tracked_threads: int = 100,
    ) -> None:
        self._stagnation_threshold = stagnation_threshold
        self._warn_escalation = warn_escalation_count
        self._inject_assessment = inject_assessment
        self._jaccard_threshold = jaccard_threshold
        self._min_words = min_words
        self._exempt_tools: set[str] = exempt_tools if exempt_tools is not None else {"ask_clarification", "write_todos", "present_files", "task"}
        self._max_tracked_threads = max_tracked_threads

        # threading.Lock (not asyncio.Lock): critical sections are short in-memory dict
        # ops with no I/O, so event-loop stall risk is negligible.  asyncio.Lock would
        # not protect the sync wrap_tool_call path used by subagent executor thread
        # pools — two separate locks would be required instead.  This matches the
        # existing LoopDetectionMiddleware pattern; see module docstring for details.
        self._lock = threading.Lock()
        # LRU-evicting store: thread_id → {tool_name → ToolPhaseState}
        self._phase_states: OrderedDict[str, dict[str, ToolPhaseState]] = OrderedDict()
        # Pending hint queue: (thread_id, run_id) → [hint texts]
        self._pending: dict[tuple[str, str], list[str]] = defaultdict(list)

    @classmethod
    def from_config(cls, config: ToolProgressConfig) -> ToolProgressMiddleware:
        return cls(
            stagnation_threshold=config.stagnation_threshold,
            warn_escalation_count=config.warn_escalation_count,
            inject_assessment=config.inject_assessment,
            jaccard_threshold=config.jaccard_similarity_threshold,
            min_words=config.min_word_count_for_similarity,
            exempt_tools=set(config.exempt_tools),
            max_tracked_threads=config.max_tracked_threads,
        )

    # ------------------------------------------------------------------
    # Runtime helpers

    @staticmethod
    def _thread_id(runtime: Runtime) -> str:
        tid = runtime.context.get("thread_id") if runtime.context else None
        return str(tid) if tid else "default"

    @staticmethod
    def _run_id(runtime: Runtime) -> str:
        rid = runtime.context.get("run_id") if runtime.context else None
        return str(rid) if rid else "default"

    def _pending_key(self, runtime: Runtime) -> tuple[str, str]:
        return self._thread_id(runtime), self._run_id(runtime)

    # ------------------------------------------------------------------
    # State store (caller holds lock)

    def _get_state(self, thread_id: str, tool_name: str) -> ToolPhaseState:
        if thread_id not in self._phase_states:
            self._phase_states[thread_id] = {}
            while len(self._phase_states) > self._max_tracked_threads:
                evicted_thread, _ = self._phase_states.popitem(last=False)
                # Evict pending hints for the evicted thread to prevent unbounded growth.
                for key in [k for k in self._pending if k[0] == evicted_thread]:
                    del self._pending[key]
        self._phase_states.move_to_end(thread_id)
        return self._phase_states[thread_id].get(tool_name, ToolPhaseState())

    def _set_state(self, thread_id: str, tool_name: str, state: ToolPhaseState) -> None:
        self._phase_states[thread_id][tool_name] = state

    def _get_block_reason(self, runtime: Runtime, tool_name: str) -> str | None:
        thread_id = self._thread_id(runtime)
        with self._lock:
            thread_tools = self._phase_states.get(thread_id)
            if thread_tools is None:
                return None
            # Read-only check: do NOT call move_to_end here. Bumping recency on the read path
            # would keep blocked threads permanently warm in the LRU, preventing healthy active
            # threads from occupying those slots. Recency is updated only on _get_state writes.
            tool_state = thread_tools.get(tool_name)
            return tool_state.block_reason if tool_state is not None and tool_state.phase == "blocked" else None

    def _make_blocked_message(self, request: ToolCallRequest, tool_name: str, block_reason: str) -> ToolMessage:
        return ToolMessage(
            content=f"[TOOL_BLOCKED] {block_reason}",
            tool_call_id=str(request.tool_call.get("id", "")),
            name=tool_name,
            status="error",
            additional_kwargs={
                TOOL_META_KEY: {
                    "status": "error",
                    "error_type": "blocked_by_progress_guard",
                    "recoverable_by_model": True,
                    "recommended_next_action": "summarize",
                    "source": "progress_middleware",
                }
            },
        )

    def _update_state_from_result(
        self,
        result: ToolMessage | Command,
        tool_name: str,
        runtime: Runtime,
    ) -> ToolMessage | Command:
        """Update the state machine from a tool result; queue hints if warranted."""
        if not isinstance(result, ToolMessage):
            return result
        meta = _parse_tool_meta((result.additional_kwargs or {}).get(TOOL_META_KEY))
        if meta is None:
            if tool_name not in self._exempt_tools:
                logger.warning(
                    "tool_progress: deerflow_tool_meta missing for non-exempt tool %s — verify ToolProgressMiddleware is outer of ToolErrorHandlingMiddleware",
                    tool_name,
                )
            return result
        content = _message_content_str(result)
        thread_id = self._thread_id(runtime)
        with self._lock:
            state = self._get_state(thread_id, tool_name)
            new_state, hint = self._assess_and_transition(state, meta, content)
            self._set_state(thread_id, tool_name, new_state)
        if new_state.phase != state.phase:
            if new_state.phase == "blocked":
                logger.warning(
                    "tool_progress: %s/%s -> BLOCKED: %s",
                    thread_id,
                    tool_name,
                    new_state.block_reason,
                )
            elif new_state.phase == "warned":
                logger.info(
                    "tool_progress: %s/%s -> WARNED (consecutive_problems=%d)",
                    thread_id,
                    tool_name,
                    new_state.consecutive_problems,
                )
            elif new_state.phase == "active":
                logger.info(
                    "tool_progress: %s/%s -> ACTIVE (reset after good result)",
                    thread_id,
                    tool_name,
                )
        if hint and self._inject_assessment:
            self._queue_assessment(runtime, hint)
        return result

    # ------------------------------------------------------------------
    # State machine

    def _assess_and_transition(
        self,
        state: ToolPhaseState,
        meta: ToolResultMeta,
        content: str,
    ) -> tuple[ToolPhaseState, str | None]:
        """Return (new_state, hint_text_or_None).

        The outer wrap_tool_call gate intercepts already-blocked states before
        the handler is called, so this function is normally reached only for
        active/warned states. If a blocked state arrives (e.g., concurrent
        transition), the function returns it unchanged — no counter inflation,
        no phase regression.
        """
        # Guard: blocked is a terminal state; nothing should change it here.
        # (In normal flow this branch is unreachable because wrap_tool_call
        # intercepts blocked tools before calling the handler.  The check exists
        # to make concurrent-race semantics well-defined and prevent a
        # recoverable-error result from silently demoting the phase back to warned.)
        if state.phase == "blocked":
            return state, None

        # Count this call as a problem before branching so all exit paths leave
        # consecutive_problems in a consistent state (never 0 when the tool has failed).
        new_count = state.consecutive_problems + 1

        # Immediately block on unrecoverable stop signals (auth, config, internal).
        if not meta.recoverable_by_model and meta.recommended_next_action == "stop":
            return replace(
                state,
                phase="blocked",
                consecutive_problems=new_count,
                block_reason=_block_reason(meta),
            ), None

        # Compute word_set only for success results: error/partial_success are problems by
        # definition and never reach the Jaccard check, so the O(n) regex is wasted on them.
        ws = word_set(content) if meta.status == "success" else frozenset()
        is_problem = meta.status in ("error", "partial_success") or (meta.status == "success" and is_near_duplicate(ws, state.recent_word_sets, self._jaccard_threshold, self._min_words))

        if not is_problem:
            # Good result: reset consecutive count, return to active.
            new_recent = (*state.recent_word_sets, ws)[-3:]
            return replace(state, consecutive_problems=0, phase="active", recent_word_sets=new_recent), None

        hint: str | None = None

        if new_count >= self._stagnation_threshold + self._warn_escalation:
            if meta.recoverable_by_model:
                # Model can fix this by changing strategy — keep warned, re-inject hint.
                # BLOCKED would prevent a legitimate retry with different parameters.
                hint = _format_hint(meta)
                new_state = replace(state, consecutive_problems=new_count, phase="warned")
            else:
                # Model cannot fix this by retrying — block the tool.
                reason = _block_reason(meta)
                new_state = replace(state, consecutive_problems=new_count, phase="blocked", block_reason=reason)
        elif new_count >= self._stagnation_threshold:
            hint = _format_hint(meta)
            new_state = replace(state, consecutive_problems=new_count, phase="warned")
        else:
            new_state = replace(state, consecutive_problems=new_count)

        return new_state, hint

    # ------------------------------------------------------------------
    # Pending queue helpers

    def _queue_assessment(self, runtime: Runtime, text: str) -> None:
        key = self._pending_key(runtime)
        thread_id = key[0]
        with self._lock:
            # Guard against creating a phantom _pending entry for a thread that was just
            # evicted from _phase_states by the LRU.  Such entries can never be cleaned up
            # by the eviction loop (which only walks _phase_states) and accumulate silently.
            if thread_id not in self._phase_states:
                return
            queue = self._pending[key]
            if len(queue) < _MAX_PENDING_PER_RUN:
                queue.append(text)

    def _drain_pending(self, runtime: Runtime) -> list[str]:
        key = self._pending_key(runtime)
        with self._lock:
            return self._pending.pop(key, [])

    def _clear_stale_pending(self, runtime: Runtime) -> None:
        thread_id, current_run = self._pending_key(runtime)
        with self._lock:
            for key in list(self._pending):
                if key[0] == thread_id and key[1] != current_run:
                    del self._pending[key]

    def _reset_run_states(self, runtime: Runtime) -> None:
        """Reset all per-run tool state for the thread at the start of a new agent run.

        Every tool's consecutive_problems counter and recent_word_sets Jaccard window are
        cleared unconditionally so state from a previous run never bleeds into the next:
        - BLOCKED/WARNED tools are reset to ACTIVE (they re-block immediately if the root
          cause persists, and the model has no memory of the prior-run hint).
        - ACTIVE tools with non-zero consecutive_problems or non-empty recent_word_sets from
          the previous run are also cleared so a single first-call problem in the new run
          cannot falsely trip WARNED against stale context from a run the model no longer sees.

        **Cross-run scoping vs LoopDetectionMiddleware**: this per-run reset is an intentional
        policy choice, not an oversight.  Errors like ``rate_limited`` and ``transient`` are
        time-bound: their root cause may resolve between user turns, so carrying a stale
        counter forward risks a false-positive BLOCKED on calls that would now succeed.
        LoopDetectionMiddleware takes the opposite stance — it retains ``_history`` across
        runs (only clearing other-run *pending* warnings at ``before_agent``), because
        call-pattern loops are time-invariant: a model that keeps issuing the same tool_calls
        regardless of results does so regardless of when the run started.  The two middlewares
        therefore guard different failure modes (result quality vs. call pattern) and their
        cross-run scoping policies intentionally differ as a consequence.
        """
        thread_id = self._thread_id(runtime)
        with self._lock:
            thread_tools = self._phase_states.get(thread_id)
            if thread_tools is None:
                return
            for tool_name, tool_state in list(thread_tools.items()):
                thread_tools[tool_name] = replace(
                    tool_state,
                    phase="active",
                    consecutive_problems=0,
                    block_reason=None,
                    recent_word_sets=(),
                )

    # ------------------------------------------------------------------
    # wrap_tool_call

    @override
    def wrap_tool_call(
        self,
        request: ToolCallRequest,
        handler: Callable[[ToolCallRequest], ToolMessage | Command],
    ) -> ToolMessage | Command:
        tool_name = str(request.tool_call.get("name", ""))
        if not tool_name or tool_name in self._exempt_tools:
            return handler(request)
        runtime = getattr(request, "runtime", None)
        if runtime is None:
            return handler(request)
        block_reason = self._get_block_reason(runtime, tool_name)
        if block_reason:
            logger.info(
                "tool_progress: %s/%s call intercepted (blocked): %s",
                self._thread_id(runtime),
                tool_name,
                block_reason,
            )
            return self._make_blocked_message(request, tool_name, block_reason)
        return self._update_state_from_result(handler(request), tool_name, runtime)

    @override
    async def awrap_tool_call(
        self,
        request: ToolCallRequest,
        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
    ) -> ToolMessage | Command:
        tool_name = str(request.tool_call.get("name", ""))
        if not tool_name or tool_name in self._exempt_tools:
            return await handler(request)
        runtime = getattr(request, "runtime", None)
        if runtime is None:
            return await handler(request)
        block_reason = self._get_block_reason(runtime, tool_name)
        if block_reason:
            logger.info(
                "tool_progress: %s/%s call intercepted (blocked): %s",
                self._thread_id(runtime),
                tool_name,
                block_reason,
            )
            return self._make_blocked_message(request, tool_name, block_reason)
        return self._update_state_from_result(await handler(request), tool_name, runtime)

    # ------------------------------------------------------------------
    # wrap_model_call: drain pending hints and inject before model sees messages

    def _augment_request(self, request: ModelRequest) -> ModelRequest:
        hints = self._drain_pending(request.runtime)
        if not hints:
            return request
        deduped = list(dict.fromkeys(hints))
        logger.debug(
            "tool_progress: injecting %d hint(s) for %s/%s",
            len(deduped),
            *self._pending_key(request.runtime),
        )
        new_messages = [
            *request.messages,
            HumanMessage(content="\n\n".join(deduped), name="progress_hint"),
        ]
        return request.override(messages=new_messages)

    @override
    def wrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> ModelCallResult:
        return handler(self._augment_request(request))

    @override
    async def awrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
    ) -> ModelCallResult:
        return await handler(self._augment_request(request))

    # ------------------------------------------------------------------
    # before_agent: clean up stale pending hints from previous runs

    @override
    def before_agent(self, state: AgentState, runtime: Runtime) -> dict | None:
        self._clear_stale_pending(runtime)
        self._reset_run_states(runtime)
        return None

    @override
    async def abefore_agent(self, state: AgentState, runtime: Runtime) -> dict | None:
        self._clear_stale_pending(runtime)
        self._reset_run_states(runtime)
        return None