perf(checkpoint): linearize message write merging (#4421)

* perf(checkpoint): linearize message write merging

* test(checkpoint): address message reducer review

backend/AGENTS.md

@@ -291,6 +291,15 @@ tool graph or subagent executor during state/schema imports.
 **ThreadState** (`packages/harness/deerflow/agents/thread_state.py`):
 - Extends `AgentState` with: `sandbox`, `thread_data`, `title`, `artifacts`, `todos`, `uploaded_files`, `viewed_images`, `goal`, `promoted`, `delegations`, `skill_context`, `summary_text`
 - Uses custom reducers: `merge_artifacts` (deduplicate), `merge_viewed_images` (merge/clear), `merge_goal` (preserve the active goal across ordinary state updates unless the goal writer replaces it), `merge_promoted` (catalog-hash-scoped deferred tool promotions), `merge_delegations` (append task delegation entries, same id latest wins, terminal status never downgraded, capped to the most recent entries), and `merge_skill_context` (dedupe active-skill references by path, keep the most recently read entries; entries store a name/path/description reference, not the SKILL.md body). `summary_text` is a LastValue channel updated by summarization and projected into model requests as durable context data instead of being stored as a `messages` item.
+- Delta-mode `merge_message_writes` normalizes the current message state once,
+  then folds normalized writes in order with message-ID position indexes and
+  deferred tombstone compaction. It preserves public `add_messages` behavior,
+  including duplicate IDs, replacement position, removal errors,
+  `REMOVE_ALL_MESSAGES`, null-write errors, and missing-ID allocation order,
+  without rescanning the accumulated state for every write. Keep this
+  full-parity contract covered by differential tests: LangGraph's private
+  `_messages_delta_reducer` is also linear, but intentionally omits some of
+  those public `add_messages` semantics and cannot be substituted directly.
 
 **Runtime Configuration** (via `config.configurable`):
 - `thinking_enabled` - Enable model's extended thinking

backend/packages/harness/deerflow/agents/thread_state.py

@@ -1,12 +1,19 @@
 import copy
+import uuid
 from collections.abc import Mapping, Sequence
 from functools import cache
-from typing import Annotated, Any, NotRequired, TypedDict, get_type_hints
+from typing import Annotated, Any, NotRequired, TypedDict, cast, get_type_hints
 
 from langchain.agents import AgentState
-from langchain_core.messages import AnyMessage
+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 add_messages
+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 GoalState
@@ -258,11 +265,91 @@ class ThreadState(AgentState):
     summary_text: NotRequired[str | None]
 
 
+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]:
-    result = list(state)
+    """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:
-        result = list(add_messages(result, write))
+        if write is None:
-    return result
+            _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]
 
 
 DELTA_MESSAGES_FIELD = Annotated[

backend/tests/test_delta_channel_state.py

@@ -1,3 +1,4 @@
+import copy
 from typing import get_type_hints
 
 import pytest
@@ -6,7 +7,7 @@ from hypothesis import strategies as st
 from langchain.agents import AgentState, create_agent
 from langchain.agents.middleware import AgentMiddleware
 from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
-from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, RemoveMessage
+from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, RemoveMessage, ToolMessageChunk
 from langgraph.channels import DeltaChannel
 from langgraph.checkpoint.memory import InMemorySaver
 from langgraph.graph import StateGraph
@@ -29,6 +30,57 @@ def _fold(state: list, writes: list) -> list:
     return result
 
 
+def _outcome(call):
+    try:
+        return ("result", call())
+    except Exception as exc:
+        return ("error", type(exc), str(exc))
+
+
+@st.composite
+def _message_merge_cases(draw):
+    message_ids = ["a", "b", "c", "missing"]
+    state_ids = draw(st.lists(st.sampled_from(message_ids[:-1]), max_size=6))
+    state = [
+        {
+            "role": draw(st.sampled_from(["user", "assistant"])),
+            "content": f"state-{index}",
+            "id": message_id,
+        }
+        for index, message_id in enumerate(state_ids)
+    ]
+
+    operation = st.one_of(
+        st.tuples(
+            st.just("message"),
+            st.sampled_from(message_ids),
+            st.sampled_from(["user", "assistant", "ai_chunk", "tool_chunk"]),
+            st.text(max_size=12),
+        ),
+        st.tuples(
+            st.just("remove"),
+            st.sampled_from([*message_ids, REMOVE_ALL_MESSAGES]),
+            st.none(),
+            st.none(),
+        ),
+    )
+    raw_writes = draw(st.lists(st.lists(operation, max_size=6), max_size=8))
+    writes = []
+    for raw_write in raw_writes:
+        write = []
+        for kind, message_id, role, content in raw_write:
+            if kind == "remove":
+                write.append(RemoveMessage(id=message_id))
+            elif role == "ai_chunk":
+                write.append(AIMessageChunk(id=message_id, content=content))
+            elif role == "tool_chunk":
+                write.append(ToolMessageChunk(id=message_id, content=content, tool_call_id=f"call-{message_id}"))
+            else:
+                write.append({"role": role, "content": content, "id": message_id})
+        writes.append(write)
+    return state, writes
+
+
 @pytest.mark.parametrize(
     "writes",
     [
@@ -45,6 +97,16 @@ def test_merge_message_writes_matches_sequential_add_messages(writes: list) -> N
     assert merge_message_writes([], writes) == _fold([], writes)
 
 
+@given(case=_message_merge_cases())
+def test_merge_message_writes_randomized_differential(case: tuple[list, list]) -> None:
+    state, writes = case
+
+    expected = _outcome(lambda: _fold(copy.deepcopy(state), copy.deepcopy(writes)))
+    actual = _outcome(lambda: merge_message_writes(copy.deepcopy(state), copy.deepcopy(writes)))
+
+    assert actual == expected
+
+
 @given(split=st.integers(min_value=0, max_value=3))
 def test_merge_message_writes_is_batching_invariant(split: int) -> None:
     state = [HumanMessage(id="h0", content="seed")]
@@ -58,6 +120,20 @@ def test_merge_message_writes_is_batching_invariant(split: int) -> None:
     assert merge_message_writes(merge_message_writes(state, xs), ys) == merge_message_writes(state, writes)
 
 
+@given(case=_message_merge_cases(), data=st.data())
+def test_merge_message_writes_randomized_batching_invariance(case: tuple[list, list], data: st.DataObject) -> None:
+    state, writes = case
+    split = data.draw(st.integers(min_value=0, max_value=len(writes)))
+
+    expected = _outcome(lambda: merge_message_writes(copy.deepcopy(state), copy.deepcopy(writes)))
+
+    def batched():
+        intermediate = merge_message_writes(copy.deepcopy(state), copy.deepcopy(writes[:split]))
+        return merge_message_writes(intermediate, copy.deepcopy(writes[split:]))
+
+    assert _outcome(batched) == expected
+
+
 def test_merge_message_writes_matches_unknown_remove_error() -> None:
     writes = [[RemoveMessage(id="missing")]]
 
@@ -69,6 +145,117 @@ def test_merge_message_writes_matches_unknown_remove_error() -> None:
     assert str(actual.value) == str(expected.value)
 
 
+@pytest.mark.parametrize(
+    ("state", "writes"),
+    [
+        (
+            [HumanMessage(id="duplicate", content="first"), HumanMessage(id="duplicate", content="second")],
+            [[AIMessage(id="duplicate", content="replacement")]],
+        ),
+        (
+            [HumanMessage(id="duplicate", content="first"), HumanMessage(id="duplicate", content="second")],
+            [[RemoveMessage(id="duplicate")]],
+        ),
+        (
+            [HumanMessage(id="same", content="old")],
+            [[RemoveMessage(id="same"), AIMessage(id="same", content="same-write replacement")]],
+        ),
+        (
+            [HumanMessage(id="same", content="old")],
+            [[RemoveMessage(id="same")], [AIMessage(id="same", content="later-write append")]],
+        ),
+        (
+            [HumanMessage(id="seed", content="old")],
+            [
+                [
+                    RemoveMessage(id="unknown-but-ignored"),
+                    RemoveMessage(id=REMOVE_ALL_MESSAGES),
+                    RemoveMessage(id="suffix-is-returned-verbatim"),
+                ]
+            ],
+        ),
+    ],
+    ids=[
+        "duplicate-id-replacement",
+        "duplicate-id-removal",
+        "same-write-remove-then-replace",
+        "cross-write-remove-then-append",
+        "remove-all-short-circuit",
+    ],
+)
+def test_merge_message_writes_preserves_add_messages_edge_semantics(state: list, writes: list) -> None:
+    assert merge_message_writes(copy.deepcopy(state), copy.deepcopy(writes)) == _fold(copy.deepcopy(state), copy.deepcopy(writes))
+
+
+@pytest.mark.parametrize(
+    ("state", "writes"),
+    [
+        ([], [None]),
+        ([HumanMessage(id="seed", content="seed")], [None]),
+        ([], [[HumanMessage(id="added", content="added")], None]),
+        (
+            [HumanMessage(id="removed", content="removed")],
+            [[RemoveMessage(id="removed")], None],
+        ),
+    ],
+)
+def test_merge_message_writes_preserves_null_write_errors(state: list, writes: list) -> None:
+    expected = _outcome(lambda: _fold(copy.deepcopy(state), copy.deepcopy(writes)))
+    actual = _outcome(lambda: merge_message_writes(copy.deepcopy(state), copy.deepcopy(writes)))
+
+    assert actual == expected
+
+
+def test_merge_message_writes_preserves_missing_id_allocation_order(monkeypatch: pytest.MonkeyPatch) -> None:
+    state = [HumanMessage(content="state")]
+    writes = [
+        [AIMessage(content="first"), HumanMessage(content="second")],
+        [AIMessage(content="third")],
+    ]
+
+    expected_ids = iter(["state-id", "first-id", "second-id", "third-id"])
+    monkeypatch.setattr("langgraph.graph.message.uuid.uuid4", lambda: next(expected_ids))
+    expected = _fold(copy.deepcopy(state), copy.deepcopy(writes))
+
+    actual_ids = iter(["state-id", "first-id", "second-id", "third-id"])
+    monkeypatch.setattr("langgraph.graph.message.uuid.uuid4", lambda: next(actual_ids))
+    actual = merge_message_writes(copy.deepcopy(state), copy.deepcopy(writes))
+
+    assert actual == expected
+
+
+def test_merge_message_writes_normalizes_state_and_each_write_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    import deerflow.agents.thread_state as thread_state
+
+    state = [HumanMessage(id="state", content="state")]
+    writes = [
+        [AIMessage(id="first", content="first")],
+        [AIMessage(id="second", content="second")],
+        [AIMessage(id="third", content="third")],
+    ]
+    original = thread_state.convert_to_messages
+    normalized_inputs = []
+
+    def record_conversion(messages):
+        normalized_inputs.append(messages)
+        return original(messages)
+
+    monkeypatch.setattr(thread_state, "convert_to_messages", record_conversion)
+
+    merge_message_writes(state, writes)
+
+    assert normalized_inputs == [state, *writes]
+
+
+def test_merge_message_writes_empty_batch_does_not_assign_state_ids() -> None:
+    state = [HumanMessage(content="unchanged")]
+
+    result = merge_message_writes(state, [])
+
+    assert result == state
+    assert state[0].id is None
+
+
 @pytest.mark.parametrize(
     "write",
     [