From da56c8dafb00aef99de34499c37b6dac1a94d7c0 Mon Sep 17 00:00:00 2001 From: rayhpeng Date: Wed, 12 Aug 2026 11:46:22 +0800 Subject: [PATCH] refactor(events): share the seq-stamping expression between the two stampers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The walrus-plus-merge expression was duplicated verbatim between stamp_messages_with_seq and _MessageSeqStamper.stamp — two counterparts of one rule where silent divergence is the likely failure mode if only one side is edited. Both now call attach_message_seq next to MESSAGE_SEQ_KEY in message_identity.py. The trailing isinstance(message, Mapping) guard was unreachable (a non-Mapping entry already got identity = None) and is gone with the extraction. Raised by review on #4696. Co-Authored-By: Claude Fable 5 --- .../runtime/events/message_identity.py | 12 +++++++- .../deerflow/runtime/events/message_seq.py | 7 ++--- .../harness/deerflow/runtime/runs/worker.py | 7 ++--- backend/tests/test_run_event_store.py | 29 +++++++++++++++++++ 4 files changed, 44 insertions(+), 11 deletions(-) diff --git a/backend/packages/harness/deerflow/runtime/events/message_identity.py b/backend/packages/harness/deerflow/runtime/events/message_identity.py index 8f2fcab66..1a05d8beb 100644 --- a/backend/packages/harness/deerflow/runtime/events/message_identity.py +++ b/backend/packages/harness/deerflow/runtime/events/message_identity.py @@ -20,7 +20,7 @@ from typing import Any from deerflow.utils.messages import strip_injected_user_message_id_suffix -__all__ = ["MESSAGE_SEQ_KEY", "message_identity"] +__all__ = ["MESSAGE_SEQ_KEY", "attach_message_seq", "message_identity"] #: ``additional_kwargs`` key carrying a message's thread-feed seq to clients. #: Server-owned display metadata: it is attached when a frame is serialized and @@ -48,3 +48,13 @@ def message_identity(message: Mapping[str, Any]) -> str | None: if message.get("type") == "human": message_id = strip_injected_user_message_id_suffix(message_id) or message_id return f"message:{message_id}" + + +def attach_message_seq(message: Mapping[str, Any], seq: int) -> dict[str, Any]: + """Return a shallow copy of *message* with *seq* under ``MESSAGE_SEQ_KEY``. + + The one stamping expression shared by the worker's run-scoped stamper and + the request-scoped ``stamp_messages_with_seq``, so the two counterparts of + the same rule cannot silently diverge. The input is never mutated. + """ + return {**message, "additional_kwargs": {**(message.get("additional_kwargs") or {}), MESSAGE_SEQ_KEY: seq}} diff --git a/backend/packages/harness/deerflow/runtime/events/message_seq.py b/backend/packages/harness/deerflow/runtime/events/message_seq.py index 0c63ff5ae..4dcd79d6f 100644 --- a/backend/packages/harness/deerflow/runtime/events/message_seq.py +++ b/backend/packages/harness/deerflow/runtime/events/message_seq.py @@ -19,7 +19,7 @@ import logging from collections.abc import Mapping, Sequence from typing import Any -from deerflow.runtime.events.message_identity import MESSAGE_SEQ_KEY, message_identity +from deerflow.runtime.events.message_identity import attach_message_seq, message_identity logger = logging.getLogger(__name__) @@ -49,7 +49,4 @@ async def stamp_messages_with_seq(store: Any, thread_id: str, messages: Sequence logger.warning("Failed to resolve message seqs for thread %s", thread_id, exc_info=True) return list(messages) - return [ - {**message, "additional_kwargs": {**(message.get("additional_kwargs") or {}), MESSAGE_SEQ_KEY: seq}} if identity is not None and (seq := found.get(identity)) is not None and isinstance(message, Mapping) else message - for message, identity in zip(messages, identities, strict=True) - ] + return [attach_message_seq(message, seq) if identity is not None and (seq := found.get(identity)) is not None else message for message, identity in zip(messages, identities, strict=True)] diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 31ff86391..8444d147d 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -48,7 +48,7 @@ from deerflow.runtime.checkpoint_state import ( graph_writable_channels, ) from deerflow.runtime.context_keys import CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY -from deerflow.runtime.events.message_identity import MESSAGE_SEQ_KEY, message_identity +from deerflow.runtime.events.message_identity import attach_message_seq, message_identity from deerflow.runtime.goal import ( DEFAULT_MAX_GOAL_CONTINUATIONS, DEFAULT_MAX_NO_PROGRESS_CONTINUATIONS, @@ -2313,10 +2313,7 @@ class _MessageSeqStamper: self._seqs.update(found) self._missing.update(unresolved - found.keys()) - stamped = [ - {**message, "additional_kwargs": {**(message.get("additional_kwargs") or {}), MESSAGE_SEQ_KEY: seq}} if identity is not None and (seq := self._seqs.get(identity)) is not None and isinstance(message, Mapping) else message - for message, identity in zip(messages, identities, strict=True) - ] + stamped = [attach_message_seq(message, seq) if identity is not None and (seq := self._seqs.get(identity)) is not None else message for message, identity in zip(messages, identities, strict=True)] return {**payload, "messages": stamped} diff --git a/backend/tests/test_run_event_store.py b/backend/tests/test_run_event_store.py index eb94b2b90..a90796096 100644 --- a/backend/tests/test_run_event_store.py +++ b/backend/tests/test_run_event_store.py @@ -993,6 +993,35 @@ class TestGetMessageSeqs: await close_engine() +class TestAttachMessageSeq: + """The one stamping expression shared by the worker's `_MessageSeqStamper` + and the request-scoped `stamp_messages_with_seq` — a single helper so the + two counterparts cannot silently diverge.""" + + def test_attaches_the_seq_under_the_server_owned_key(self): + from deerflow.runtime.events.message_identity import attach_message_seq + + stamped = attach_message_seq({"type": "human", "id": "u1"}, 7) + + assert stamped["additional_kwargs"] == {"deerflow_seq": 7} + + def test_existing_additional_kwargs_are_preserved(self): + from deerflow.runtime.events.message_identity import attach_message_seq + + stamped = attach_message_seq({"type": "ai", "id": "a1", "additional_kwargs": {"run_id": "r1"}}, 3) + + assert stamped["additional_kwargs"] == {"run_id": "r1", "deerflow_seq": 3} + + def test_the_input_message_is_not_mutated(self): + from deerflow.runtime.events.message_identity import attach_message_seq + + message = {"type": "human", "id": "u1", "additional_kwargs": {"run_id": "r1"}} + + attach_message_seq(message, 5) + + assert message["additional_kwargs"] == {"run_id": "r1"} + + class TestStampMessagesWithSeq: """Attach the feed seq to an arbitrary list of checkpoint messages.