deer-flow/backend/packages/harness/deerflow/agents/thread_state.py

import copy
import uuid
from collections.abc import Mapping, Sequence
from functools import cache
from typing import Annotated, Any, NotRequired, TypedDict, cast, get_type_hints

from langchain.agents import AgentState
from langchain_core.messages import (
    AnyMessage,
    BaseMessageChunk,
    RemoveMessage,
    convert_to_messages,
    message_chunk_to_message,
)
from langgraph.channels import DeltaChannel
from langgraph.graph.message import REMOVE_ALL_MESSAGES

import deerflow.checkpoint_patches as _checkpoint_patches  # noqa: F401 - import-time saver fixes
from deerflow.agents.goal_state import GoalOutcomeState, GoalState
from deerflow.agents.task_continuity.state import TaskNotesChannel, merge_task_notes
from deerflow.config.database_config import DEFAULT_CHECKPOINT_SNAPSHOT_FREQUENCY, CheckpointChannelMode
from deerflow.subagents.status_contract import SUBAGENT_STATUS_VALUES


def _resolve_snapshot_frequency(snapshot_frequency: int | None) -> int:
    """Resolve the effective cadence: explicit value, else process-frozen,
    else default. Imported lazily — ``deerflow.runtime.__init__`` reaches this
    module via ``checkpoint_state``, so a top-level import would cycle."""
    if snapshot_frequency is not None:
        return snapshot_frequency
    from deerflow.runtime.checkpoint_mode import resolve_checkpoint_snapshot_frequency

    return resolve_checkpoint_snapshot_frequency()


class SandboxState(TypedDict):
    sandbox_id: NotRequired[str | None]


class ThreadDataState(TypedDict):
    workspace_path: NotRequired[str | None]
    uploads_path: NotRequired[str | None]
    outputs_path: NotRequired[str | None]


class BackgroundTaskState(TypedDict):
    task_id: str
    task_name: str
    status: str
    updated_at: str


class BlobRefData(TypedDict):
    """JSON-safe checkpoint representation of ``deerflow.storage.BlobRef``."""

    sha256: str
    size: int
    kind: str
    content_type: NotRequired[str | None]


class ViewedImageData(TypedDict):
    """Metadata for a viewed image file.

    Only lightweight metadata is persisted in checkpoint state. Image bytes are
    read on-demand from the active sandbox or from a synchronized host copy whose
    size and SHA-256 match the previously viewed bytes. This avoids duplicating
    large base64 payloads across every checkpoint (see #4138).
    """

    mime_type: str
    size: int
    actual_path: str
    sha256: str
    source_sandbox_id: NotRequired[str]
    blob_ref: NotRequired[BlobRefData]


def merge_sandbox(existing: SandboxState | None, new: SandboxState | None) -> SandboxState | None:
    """Reducer for sandbox state - accepts idempotent writes only.

    Multiple sandbox tools can initialize lazily in the same graph step and
    emit the same sandbox_id via Command(update=...). LangGraph needs an
    explicit reducer for that shared state key. Different sandbox ids in the
    same thread indicate a lifecycle/isolation bug, so fail closed instead of
    choosing one silently.
    """
    if new is None:
        return existing
    if existing is None:
        return new

    existing_id = existing.get("sandbox_id")
    new_id = new.get("sandbox_id")
    if existing_id == new_id:
        return existing
    raise ValueError(f"Conflicting sandbox state updates: {existing_id!r} != {new_id!r}")


SandboxStateField = Annotated[NotRequired[SandboxState | None], merge_sandbox]


def merge_artifacts(existing: list[str] | None, new: list[str] | None) -> list[str]:
    """Reducer for artifacts list - merges and deduplicates artifacts."""
    if existing is None:
        return new or []
    if new is None:
        return existing
    # Use dict.fromkeys to deduplicate while preserving order
    return list(dict.fromkeys(existing + new))


def merge_viewed_images(existing: dict[str, ViewedImageData] | None, new: dict[str, ViewedImageData] | None) -> dict[str, ViewedImageData]:
    """Reducer for viewed_images dict - merges image dictionaries.

    Special case: If new is an empty dict {}, it clears the existing images.
    This allows middlewares to clear the viewed_images state after processing.

    Dict order tracks recency: a key updated by ``new`` is re-inserted at the
    end, so a re-viewed image becomes the newest entry. The image injection
    middleware admits entries newest-first and tells the model to re-view an
    omitted image, so the recovery only works if the re-viewed entry actually
    moves to the newest position (#6825, review on #6829).
    """
    if existing is None:
        return new or {}
    if new is None:
        return existing
    # Special case: empty dict means clear all viewed images
    if len(new) == 0:
        return {}
    # Merge dictionaries; new values override existing ones for same keys, and
    # updated keys move to the end so insertion order reflects recency.
    merged = {key: value for key, value in existing.items() if key not in new}
    merged.update(new)
    return merged


def merge_todos(existing: list | None, new: list | None) -> list | None:
    """Reducer for todos list - keeps the last non-None value.

    Semantics:
    - If `new` is None (node didn't touch todos), preserve `existing`.
    - If `new` is provided (even empty list), it represents an explicit
      update and wins over `existing`.
    """
    if new is None:
        return existing
    return new


def merge_goal(existing: GoalState | None, new: GoalState | None) -> GoalState | None:
    """Reducer for goal state - preserves existing when a node does not touch it."""
    if new is None:
        return existing
    return new


class PromotedTools(TypedDict):
    catalog_hash: str
    names: list[str]
    turn_id: NotRequired[str]


def merge_promoted(existing: PromotedTools | None, new: PromotedTools | None) -> PromotedTools | None:
    """Reducer for deferred-tool promotions, scoped by catalog and user turn.

    - new None/empty -> preserve existing (node didn't touch promotions).
    - catalog_hash or explicit turn_id changed -> replace wholesale, dropping
      stale names.
    - same scope -> union names, dedupe, preserve order.
    - a missing new turn_id inherits the existing scope. ``tool_search`` writes
      this shape after the filter establishes the current turn before the model.
    """
    if not new:
        return existing
    new_turn_id = new.get("turn_id")
    replace = existing is None or existing.get("catalog_hash") != new["catalog_hash"] or (new_turn_id is not None and existing.get("turn_id") != new_turn_id)
    existing_names = existing.get("names") if existing is not None else None
    if not isinstance(existing_names, list) or not all(isinstance(name, str) for name in existing_names):
        existing_names = []
    names = new["names"] if replace else existing_names + new["names"]
    merged: PromotedTools = {
        "catalog_hash": new["catalog_hash"],
        "names": list(dict.fromkeys(names)),
    }
    turn_id = new_turn_id if new_turn_id is not None else (existing or {}).get("turn_id")
    if turn_id is not None:
        merged["turn_id"] = turn_id
    return merged


TERMINAL_STATUSES: frozenset[str] = frozenset(SUBAGENT_STATUS_VALUES)
_DELEGATION_LEDGER_MAX_ENTRIES = 50


class DelegationEntry(TypedDict):
    id: str
    run_id: NotRequired[str]
    description: str
    subagent_type: str
    status: str
    result_brief: NotRequired[str]
    result_sha256: NotRequired[str]
    result_ref: NotRequired[str]
    # Why a guardrail cap ended the run early (#3875 Phase 2): token_capped /
    # turn_capped / loop_capped. The status stays completed/failed; this field
    # is the additive signal that distinguishes a capped run from a clean one.
    stop_reason: NotRequired[str]
    # RFC #4651 PR2: parent-side citation-check verdict (advisory execution
    # evidence), stamped at task write-back; absent on legacy history.
    receipt_verdict: NotRequired[dict]
    # RFC #4651 PR4: deterministic acceptance-checklist verdict, same
    # provenance as receipt_verdict.
    acceptance_verdict: NotRequired[dict]
    created_at: str
    # The AIMessage that issued the call. Tool-call ids are unique per response,
    # so a provider can reuse one in a later response of the same run; this keeps
    # each dispatch its own entry. Absent on entries written before it existed.
    dispatch_message_id: NotRequired[str]


type DelegationKey = tuple[str | None, str, str | None]


def delegation_key(entry: Mapping[str, Any]) -> DelegationKey:
    """Identify one delegation: ``(run_id, tool-call id, dispatch message id)``.

    DurableContext uses this key too. The cap's ``_delegation_id`` reads the
    same fields from entries; keep the two in step.
    """
    return entry.get("run_id") or None, entry["id"], entry.get("dispatch_message_id") or None


def find_delegation_key(keys: Sequence[DelegationKey], key: DelegationKey) -> DelegationKey | None:
    """Return the key in ``keys`` that an entry identified by ``key`` updates, or None.

    An exact match wins. Otherwise an entry without a run id updates the most
    recent entry with its call id, as updates did before run-scoped
    identities; an entry without a dispatch id updates the most recent entry
    of its run with that call id; and a dispatch-scoped entry adopts the most
    recent entry of its run and call id that predates dispatch ids. Two
    different dispatches never match each other.
    """
    if key in keys:
        return key
    run_id, entry_id, dispatch_id = key
    for prior in reversed(keys):
        if prior[1] != entry_id:
            continue
        if run_id is None:
            return prior
        if prior[0] == run_id and (dispatch_id is None or prior[2] is None):
            return prior
    return None


def merge_delegations(existing: list[DelegationEntry] | None, new: list[DelegationEntry] | None) -> list[DelegationEntry]:
    """Reducer for the delegation ledger.

    - new None/empty -> preserve existing.
    - append entries, replacing the same delegation (see ``find_delegation_key``)
      with the latest version while preserving first-seen order. Provider ids
      can repeat in later runs of the same thread and in later responses of
      the same run.
    - terminal status is never overwritten by a non-terminal status.
    """
    if not new:
        return existing or []

    by_key: dict[DelegationKey, DelegationEntry] = {}
    order: list[DelegationKey] = []
    for entry in [*(existing or []), *new]:
        target = find_delegation_key(order, delegation_key(entry))
        previous = by_key.get(target) if target is not None else None
        if previous is None:
            key = delegation_key(entry)
            order.append(key)
            by_key[key] = entry
            continue
        if previous["status"] in TERMINAL_STATUSES and entry["status"] not in TERMINAL_STATUSES:
            continue
        inherited = {field: previous[field] for field in ("created_at", "run_id", "dispatch_message_id") if previous.get(field) and (field == "created_at" or not entry.get(field))}
        if inherited:
            entry = {**entry, **inherited}
        key = delegation_key(entry)
        if key != target:
            # A legacy entry gained its dispatch id: re-key it in place so a
            # later dispatch reusing the call id starts its own entry.
            order[order.index(target)] = key
            del by_key[target]
        by_key[key] = entry
    merged = [by_key[key] for key in order]
    if len(merged) > _DELEGATION_LEDGER_MAX_ENTRIES:
        merged = merged[-_DELEGATION_LEDGER_MAX_ENTRIES:]
    return merged


_SKILL_CONTEXT_MAX_ENTRIES = 8
_SKILL_DESCRIPTION_MAX_CHARS = 500


class SkillEntry(TypedDict):
    name: str
    path: str
    description: str
    loaded_at: int


def _normalize_skill_entry(entry: Mapping[str, object]) -> SkillEntry:
    """Drop legacy payload keys before storing skill_context back to state."""
    description = entry.get("description")
    loaded_at = entry.get("loaded_at")
    return {
        "name": str(entry.get("name") or ""),
        "path": str(entry["path"]),
        "description": " ".join(description.split())[:_SKILL_DESCRIPTION_MAX_CHARS] if isinstance(description, str) else "",
        "loaded_at": loaded_at if isinstance(loaded_at, int) else 0,
    }


def merge_skill_context(existing: list[SkillEntry] | None, new: list[SkillEntry] | None) -> list[SkillEntry]:
    """Reducer for the skill-context channel.

    - new None/empty -> preserve existing.
    - legacy entries are normalized to references; verbatim body keys are dropped.
    - dedup by ``path``; later reads refresh recency and replace the reference.
    - cap by keeping the most recently read entries. ``loaded_at`` is
      observational only because message indices reset after compaction.
    """
    normalized_existing = [_normalize_skill_entry(entry) for entry in existing or []]
    if not new:
        return normalized_existing

    by_path: dict[str, SkillEntry] = {}
    order: list[str] = []
    for entry in normalized_existing:
        path = entry["path"]
        if path not in by_path:
            order.append(path)
        by_path[path] = entry

    for entry in (_normalize_skill_entry(entry) for entry in new):
        path = entry["path"]
        if path in by_path:
            order.remove(path)
        order.append(path)
        by_path[path] = entry

    merged = [by_path[path] for path in order]
    if len(merged) > _SKILL_CONTEXT_MAX_ENTRIES:
        merged = merged[-_SKILL_CONTEXT_MAX_ENTRIES:]
    return merged


# Absolute ceiling applied by the reducer itself. The operator-configured cap
# (`tool_artifacts.max_entries`) is enforced per-agent by
# ArtifactCaptureMiddleware when it emits updates, so two agents in one process
# can carry different caps without sharing state.
_ARTIFACT_MAX_ENTRIES_CEILING = 1000


class ArtifactEntry(TypedDict):
    """A captured artifact reference from a tool result (issue #4676).

    Stored in ``ThreadState.tool_artifacts`` so it survives context compaction.
    The model sees only the short ``handle``; the real reference (path, URL,
    task id) lives here and is resolved at tool-call time.
    """

    handle: str
    tool_name: str
    tool_call_id: str
    call_index: int
    artifact_type: str
    display_name: str
    real_ref: str
    mime_type: NotRequired[str]
    created_at: str
    consumed_by: NotRequired[list[str]]


def merge_tool_artifacts(existing: list[ArtifactEntry] | None, new: list[ArtifactEntry] | None) -> list[ArtifactEntry]:
    """Reducer for the tool-artifact registry channel.

    - new None/empty -> preserve existing.
    - append entries, replacing same handle with the latest version while
      preserving first-seen order (latest wins, e.g. for consumption updates).
    - a trailing ``{"op": "trim_to", "keep": N}`` directive (emitted by the
      capture middleware) makes the configured cap a sliding window: the
      oldest entries beyond N are evicted. Without a directive nothing is
      evicted, so updates stay purely additive.
    - absolute ceiling of 1000 applies regardless.
    """
    if not new:
        return existing or []

    by_handle: dict[str, ArtifactEntry] = {}
    order: list[str] = []
    for entry in [*(existing or []), *new]:
        if not isinstance(entry, dict) or "handle" not in entry:
            continue
        handle = entry["handle"]
        if handle not in by_handle:
            order.append(handle)
        by_handle[handle] = entry

    merged = [by_handle[handle] for handle in order]

    for item in reversed(new):
        if isinstance(item, dict) and item.get("op") == "trim_to":
            keep = min(int(item.get("keep", len(merged))), _ARTIFACT_MAX_ENTRIES_CEILING)
            if keep < len(merged):
                merged = merged[-keep:]
            break

    if len(merged) > _ARTIFACT_MAX_ENTRIES_CEILING:
        merged = merged[-_ARTIFACT_MAX_ENTRIES_CEILING:]
    return merged


class ThreadState(AgentState):
    sandbox: SandboxStateField
    thread_data: NotRequired[ThreadDataState | None]
    title: NotRequired[str | None]
    artifacts: Annotated[list[str], merge_artifacts]
    todos: Annotated[list | None, merge_todos]
    goal: Annotated[GoalState | None, merge_goal]
    goal_outcome: NotRequired[GoalOutcomeState | None]
    applied_run_inputs: NotRequired[dict | None]
    uploaded_files: NotRequired[list[dict] | None]
    viewed_images: Annotated[dict[str, ViewedImageData], merge_viewed_images]  # image_path -> metadata (no base64)
    promoted: Annotated[PromotedTools | None, merge_promoted]
    delegations: Annotated[list[DelegationEntry], merge_delegations]
    skill_context: Annotated[list[SkillEntry], merge_skill_context]
    tool_artifacts: Annotated[list[ArtifactEntry], merge_tool_artifacts]
    tool_artifact_processed: Annotated[list[str], merge_artifacts]
    task_notes: Annotated[dict | None, TaskNotesChannel(dict | None, merge_task_notes)]
    task_history: NotRequired[dict | None]
    summary_text: NotRequired[str | None]
    background_tasks: NotRequired[list[BackgroundTaskState]]


def _normalize_messages(value: Any) -> list[AnyMessage]:
    values = value if isinstance(value, list) else [value]
    messages = [message_chunk_to_message(cast(BaseMessageChunk, message)) for message in convert_to_messages(values)]
    for message in messages:
        if message.id is None:
            message.id = str(uuid.uuid4())
    return messages


def _index_messages(
    messages: list[AnyMessage | None],
) -> tuple[dict[str, int], dict[str, list[int]]]:
    latest_position: dict[str, int] = {}
    positions_by_id: dict[str, list[int]] = {}
    for position, message in enumerate(messages):
        if message is None:
            continue
        message_id = cast(str, message.id)
        latest_position[message_id] = position
        positions_by_id.setdefault(message_id, []).append(position)
    return latest_position, positions_by_id


def _raise_null_write(has_messages: bool) -> None:
    # ``add_messages(left, None)`` reports only ``left`` when the accumulated
    # message list is non-empty; with an empty list, it reports only ``right``.
    received = "left" if has_messages else "right"
    raise ValueError(f"Must specify non-null arguments for both 'left' and 'right'. Only received: '{received}'.")


def merge_message_writes(state: list[AnyMessage], writes: Sequence[Any]) -> list[AnyMessage]:
    """Fold DeltaChannel writes with ``add_messages`` semantics in linear time.

    LangGraph's private ``_messages_delta_reducer`` is also linear, but does
    not preserve the public reducer's full coercion, ID, removal, and
    ``REMOVE_ALL_MESSAGES`` behavior.
    """
    if not writes:
        return list(state)
    if writes[0] is None:
        _raise_null_write(bool(state))

    messages: list[AnyMessage | None] = _normalize_messages(state)
    latest_position, positions_by_id = _index_messages(messages)

    for write in writes:
        if write is None:
            _raise_null_write(bool(latest_position))
        normalized_write = _normalize_messages(write)
        remove_all_idx = None
        for position, message in enumerate(normalized_write):
            if isinstance(message, RemoveMessage) and message.id == REMOVE_ALL_MESSAGES:
                remove_all_idx = position

        if remove_all_idx is not None:
            messages = list(normalized_write[remove_all_idx + 1 :])
            latest_position, positions_by_id = _index_messages(messages)
            continue

        ids_to_remove: set[str] = set()
        for message in normalized_write:
            message_id = cast(str, message.id)
            existing_position = latest_position.get(message_id)
            if existing_position is not None:
                if isinstance(message, RemoveMessage):
                    ids_to_remove.add(message_id)
                else:
                    ids_to_remove.discard(message_id)
                    messages[existing_position] = message
                continue

            if isinstance(message, RemoveMessage):
                raise ValueError(f"Attempting to delete a message with an ID that doesn't exist ('{message_id}')")

            position = len(messages)
            messages.append(message)
            latest_position[message_id] = position
            positions_by_id[message_id] = [position]

        for message_id in ids_to_remove:
            for position in positions_by_id.pop(message_id):
                messages[position] = None
            del latest_position[message_id]

    return [message for message in messages if message is not None]


def delta_messages_field(snapshot_frequency: int = DEFAULT_CHECKPOINT_SNAPSHOT_FREQUENCY) -> Any:
    """Messages field annotation with a ``DeltaChannel`` at the given cadence."""
    return Annotated[
        list[AnyMessage],
        DeltaChannel(merge_message_writes, snapshot_frequency=snapshot_frequency),
    ]


DELTA_MESSAGES_FIELD = delta_messages_field()


class DeltaThreadState(ThreadState):
    messages: DELTA_MESSAGES_FIELD


THREAD_STATE_REDUCER_FIELDS = frozenset(
    {
        "messages",
        "sandbox",
        "artifacts",
        "todos",
        "goal",
        "viewed_images",
        "promoted",
        "delegations",
        "skill_context",
        "tool_artifacts",
        "tool_artifact_processed",
        "task_notes",
    }
)


def get_thread_state_schema(mode: CheckpointChannelMode, snapshot_frequency: int | None = None) -> type:
    if mode != "delta":
        return ThreadState
    return _delta_thread_state_schema(_resolve_snapshot_frequency(snapshot_frequency))


@cache
def _delta_thread_state_schema(snapshot_frequency: int) -> type:
    """Delta thread schema keyed by cadence; the default keeps the static
    ``DeltaThreadState`` identity so existing type checks keep holding."""
    if snapshot_frequency == DEFAULT_CHECKPOINT_SNAPSHOT_FREQUENCY:
        return DeltaThreadState
    annotations = get_type_hints(ThreadState, include_extras=True)
    annotations["messages"] = delta_messages_field(snapshot_frequency)
    return TypedDict(
        f"DeltaThreadState_f{snapshot_frequency}",
        annotations,
        total=getattr(ThreadState, "__total__", True),
    )


def adapt_state_schema_for_mode(schema: type, mode: CheckpointChannelMode, snapshot_frequency: int | None = None) -> type:
    if mode == "full":
        return schema
    return _adapt_state_schema_for_delta(schema, _resolve_snapshot_frequency(snapshot_frequency))


@cache
def _adapt_state_schema_for_delta(schema: type, snapshot_frequency: int) -> type:
    annotations = get_type_hints(schema, include_extras=True)
    annotations["messages"] = delta_messages_field(snapshot_frequency)
    return TypedDict(
        f"Delta{schema.__module__.replace('.', '_')}_{schema.__name__}_f{snapshot_frequency}",
        annotations,
        total=getattr(schema, "__total__", True),
    )


def normalize_middleware_state_schemas(middleware: Sequence[Any], mode: CheckpointChannelMode, snapshot_frequency: int | None = None) -> list[Any]:
    if mode == "full":
        return list(middleware)
    resolved_frequency = _resolve_snapshot_frequency(snapshot_frequency)
    normalized = []
    for item in middleware:
        schema = getattr(item, "state_schema", None)
        if schema is None:
            normalized.append(item)
            continue
        adapted = copy.copy(item)
        adapted.state_schema = adapt_state_schema_for_mode(schema, mode, resolved_frequency)
        normalized.append(adapted)
    return normalized