From 8c19a2eb363a945dc0169950c9831af0efb924f3 Mon Sep 17 00:00:00 2001 From: Vanzeren <53075619+Vanzeren@users.noreply.github.com> Date: Sat, 25 Jul 2026 21:19:24 +0800 Subject: [PATCH] perf(checkpoint): linearize message write merging (#4421) * perf(checkpoint): linearize message write merging * test(checkpoint): address message reducer review --- backend/AGENTS.md | 9 + .../harness/deerflow/agents/thread_state.py | 99 ++++++++- backend/tests/test_delta_channel_state.py | 189 +++++++++++++++++- 3 files changed, 290 insertions(+), 7 deletions(-) diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 417ebc455..5d3044df0 100644 --- a/backend/AGENTS.md +++ b/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 diff --git a/backend/packages/harness/deerflow/agents/thread_state.py b/backend/packages/harness/deerflow/agents/thread_state.py index 385794a67..5784c267d 100644 --- a/backend/packages/harness/deerflow/agents/thread_state.py +++ b/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)) - return result + 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] DELTA_MESSAGES_FIELD = Annotated[ diff --git a/backend/tests/test_delta_channel_state.py b/backend/tests/test_delta_channel_state.py index d9dcc7bc3..fdcbfa6e1 100644 --- a/backend/tests/test_delta_channel_state.py +++ b/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", [