From 9a43d8276d22a3f2308924a6b3a5d5aa3fe7939a Mon Sep 17 00:00:00 2001 From: Aari Date: Tue, 28 Jul 2026 22:12:27 +0800 Subject: [PATCH] fix(gateway): replay edit and rerun from a settled checkpoint (#4534) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Editing the only turn of a thread reran the original prompt: the model answered the question the edit was replacing while the UI showed the edited text, and the edit vanished on reload. The replay-base lookup decided whether a checkpoint predates the target user message by message id alone. DynamicContextMiddleware re-keys the first user turn to `{id}__user` mid-run, so every checkpoint written before it holds the same prompt under an id the lookup cannot match. The scan walked past those and anchored inside the run that produced the turn — a checkpoint that still contains the original prompt and owns the injection node's pending writes, which the replay then re-added after the edited message. Require the replay base to be a settled checkpoint (no pending tasks) in both the lineage walk and the chronological fallback. That rule is middleware agnostic: the first turn now anchors on the thread's empty initial checkpoint and later turns on the previous run's tail, which also drops the existing reliance on LangGraph discarding a stale `__start__` write. Edit replay additionally passes `head_checkpoint` so it resolves its base lineage-first like regenerate does, and a replayed user message is restored to its pre-swap id: replaying `{id}__user` into a state that has no reminder yet makes the middleware treat the turn as already injected and silently drops its date and memory block. Frontend: a prepared replay masks the turn it supersedes, so the optimistic-message baseline is taken from the post-mask human count. The pre-mask count can never be exceeded when the replay puts exactly one human message back, and on the first turn the runtime re-keys the replacement message so identity comparison cannot stand in for the count. Fixes #4531 --- backend/AGENTS.md | 14 ++ backend/app/gateway/checkpoint_lineage.py | 29 +++- backend/app/gateway/routers/thread_runs.py | 15 +- .../middlewares/dynamic_context_middleware.py | 21 ++- backend/tests/test_checkpoint_lineage.py | 149 ++++++++++++++++++ .../tests/test_thread_regenerate_prepare.py | 111 ++++++++++++- frontend/src/core/threads/hooks.ts | 25 +++ .../unit/core/threads/message-merge.test.ts | 53 +++++++ 8 files changed, 404 insertions(+), 13 deletions(-) create mode 100644 backend/tests/test_checkpoint_lineage.py diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 03faf4389..bcbc6995e 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -533,6 +533,20 @@ full and delta checkpoint modes. Only an explicitly absent legacy parent link ma use chronological compatibility lookup; cycles, dangling links, and depth-limit exhaustion fail closed. Existing single-checkpoint branches are never repaired by copying a raw checkpoint because delta state is not self-contained in one tuple. +Both lookups additionally require the replay base to be a **settled** checkpoint +(`has_pending_tasks` — no scheduled `next` tasks). A checkpoint with pending tasks +is a mid-run snapshot: resuming from it replays the writes of the node that was +about to run. Message ids alone cannot exclude those, because middleware may +rewrite a message's id inside the run that produced it — `DynamicContextMiddleware` +moves the first user turn to `{id}__user` and gives `{id}` to the injected +reminder, so every checkpoint written before it holds the same prompt under an +unmatched id. Selecting one of those re-added the original prompt *after* the +edited one, and the model answered the question the edit was replacing (#4531). +`next` is not derivable on the degraded raw-checkpoint read path, which reports no +tasks; absence of evidence stays permissive there rather than failing closed. +Edit replay resolves its base through the same lineage-first path as regenerate; +it must pass `head_checkpoint` or it silently degrades to the chronological scan +that cannot tell sibling branches apart. ### Sandbox System (`packages/harness/deerflow/sandbox/`) diff --git a/backend/app/gateway/checkpoint_lineage.py b/backend/app/gateway/checkpoint_lineage.py index 6b121c318..275784c7d 100644 --- a/backend/app/gateway/checkpoint_lineage.py +++ b/backend/app/gateway/checkpoint_lineage.py @@ -45,6 +45,23 @@ def is_duration_only_checkpoint(checkpoint_tuple: Any) -> bool: return isinstance(writes, dict) and "runtime_run_duration" in writes +def has_pending_tasks(checkpoint_tuple: Any) -> bool: + """Return whether *checkpoint_tuple* still has graph work scheduled. + + A replay base must be a state the thread was at rest in. A mid-run + checkpoint owns the writes of the node that was about to run, and resuming + from it replays them — re-adding the very turn the replay is meant to + replace. Message ids alone cannot detect this, because middleware may + rewrite a message's id in the same run that produced it. + + ``next`` is not derivable on the degraded raw-checkpoint read path, which + reports no tasks at all. Absence of evidence therefore stays permissive: + those reads keep selecting the same base they always did. + """ + + return bool(getattr(checkpoint_tuple, "next", None)) + + def _message_id(message: Any) -> str | None: value = getattr(message, "id", None) if value is None and isinstance(message, dict): @@ -96,7 +113,8 @@ async def find_checkpoint_before_message( Following ``parent_config`` is important after a regenerate: a thread can contain sibling checkpoint branches, and a global time-ordered scan can otherwise select a checkpoint from the wrong branch. Duration-only metadata checkpoints do not - represent an addressable conversation state and are skipped. + represent an addressable conversation state and are skipped, and so are + checkpoints that still have pending tasks (see :func:`has_pending_tasks`). """ if message_id not in {_message_id(message) for message in checkpoint_messages(head_checkpoint)}: @@ -132,7 +150,7 @@ async def find_checkpoint_before_message( continue parent_message_ids = {_message_id(message) for message in checkpoint_messages(parent)} - if message_id not in parent_message_ids: + if message_id not in parent_message_ids and not has_pending_tasks(parent): return parent current = parent @@ -148,8 +166,9 @@ def find_checkpoint_before_message_chronologically( This is a compatibility fallback for imported or legacy checkpoints that do not carry ``parent_config`` links. Callers must prefer the lineage walk when links are available because a chronological scan cannot distinguish sibling - checkpoint branches. Duration-only checkpoints are ignored, and only - checkpoints with an addressable id can become the replay base. + checkpoint branches. Duration-only checkpoints are ignored, and only settled + checkpoints (see :func:`has_pending_tasks`) with an addressable id can become + the replay base. """ previous_checkpoint = None @@ -159,6 +178,6 @@ def find_checkpoint_before_message_chronologically( message_ids = {_message_id(message) for message in checkpoint_messages(checkpoint_tuple)} if message_id in message_ids: return previous_checkpoint, True - if checkpoint_configurable(checkpoint_tuple).get("checkpoint_id"): + if checkpoint_configurable(checkpoint_tuple).get("checkpoint_id") and not has_pending_tasks(checkpoint_tuple): previous_checkpoint = checkpoint_tuple return None, False diff --git a/backend/app/gateway/routers/thread_runs.py b/backend/app/gateway/routers/thread_runs.py index 5a0c226ee..c12fa4f1f 100644 --- a/backend/app/gateway/routers/thread_runs.py +++ b/backend/app/gateway/routers/thread_runs.py @@ -38,6 +38,7 @@ from app.gateway.pagination import trim_run_message_page from app.gateway.run_models import RunCreateRequest from app.gateway.services import build_checkpoint_state_accessor, build_thread_checkpoint_state_accessor, sse_consumer, start_run, wait_for_run_completion from app.gateway.utils import sanitize_log_param +from deerflow.agents.middlewares.dynamic_context_middleware import strip_injected_user_message_id_suffix from deerflow.runtime import CancelOutcome, RunRecord, RunStatus, serialize_channel_values_for_api from deerflow.runtime.secret_context import redact_config_secrets, redact_metadata_secrets from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY, get_original_user_content_text, message_to_text @@ -344,7 +345,12 @@ def _clean_human_message_for_regenerate(message: Any) -> dict[str, Any]: "content": [{"type": "text", "text": content}], "additional_kwargs": additional_kwargs, } - message_id = _message_id(message) + # Replay the id the client originally sent. The dynamic-context reminder + # re-keys the first user message of a thread to `{id}__user`, and replaying + # that persisted id into a state that has no reminder yet makes the + # middleware treat the turn as already injected, silently dropping the date + # and memory block the original turn had. + message_id = strip_injected_user_message_id_suffix(_message_id(message)) if message_id: clean_message["id"] = message_id name = _message_name(message) @@ -653,7 +659,12 @@ async def _prepare_edit_regenerate_payload( if not source_ai_id: raise HTTPException(status_code=409, detail="The source assistant message is missing an id") - base_checkpoint_tuple = await _find_base_checkpoint_before_human(thread_id, source_human_id, request) + base_checkpoint_tuple = await _find_base_checkpoint_before_human( + thread_id, + source_human_id, + request, + head_checkpoint=latest_checkpoint, + ) target_run_id = await _find_target_run_id(thread_id, source_ai_id, source_ai, source_human, request) source_record = await _require_successful_source_run(thread_id, target_run_id, request) checkpoint = _checkpoint_response(base_checkpoint_tuple) diff --git a/backend/packages/harness/deerflow/agents/middlewares/dynamic_context_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/dynamic_context_middleware.py index e4e35d23a..599a35ec5 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/dynamic_context_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/dynamic_context_middleware.py @@ -60,6 +60,23 @@ _DYNAMIC_CONTEXT_REMINDER_KEY = "dynamic_context_reminder" # so it is never exposed to user-influenceable memory content. _REMINDER_DATE_KEY = "reminder_date" _SUMMARY_MESSAGE_NAME = "summary" +# Suffix the ID-swap gives the real user message; the reminder SystemMessage +# takes the original id so ``add_messages`` can replace it in place. +INJECTED_USER_MESSAGE_ID_SUFFIX = "__user" + + +def strip_injected_user_message_id_suffix(message_id: str | None) -> str | None: + """Return the id *message_id* had before the reminder ID-swap. + + Replaying a persisted user turn must feed the graph the id the client + originally sent: a ``{id}__user`` message is skipped as an injection target, + so replaying one into a state that has no reminder yet silently drops the + date and memory block for that turn. + """ + + if isinstance(message_id, str) and message_id.endswith(INJECTED_USER_MESSAGE_ID_SUFFIX): + return message_id[: -len(INJECTED_USER_MESSAGE_ID_SUFFIX)] or message_id + return message_id def _extract_date(content: str) -> str | None: @@ -120,7 +137,7 @@ def _is_user_injection_target(message: object) -> bool: # (id__user__user__user...) and ghost-message re-execution. # Using endswith (not substring "in") avoids false positives on IDs that # happen to contain "__user" in the middle. - if message.id and str(message.id).endswith("__user"): + if message.id and str(message.id).endswith(INJECTED_USER_MESSAGE_ID_SUFFIX): return False return True @@ -232,7 +249,7 @@ class DynamicContextMiddleware(AgentMiddleware): messages.append( HumanMessage( content=original.content, - id=f"{stable_id}__user", + id=f"{stable_id}{INJECTED_USER_MESSAGE_ID_SUFFIX}", name=original.name, additional_kwargs=original.additional_kwargs, ) diff --git a/backend/tests/test_checkpoint_lineage.py b/backend/tests/test_checkpoint_lineage.py new file mode 100644 index 000000000..10d560924 --- /dev/null +++ b/backend/tests/test_checkpoint_lineage.py @@ -0,0 +1,149 @@ +"""Replay-base resolution rules shared by regenerate, edit replay and branching. + +The load-bearing rule here is that a replay base must be a checkpoint the +thread was actually at rest in. A mid-run checkpoint still owns the interrupted +node's pending writes, so resuming from it replays them. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage + +from app.gateway.checkpoint_lineage import ( + CheckpointParentMissingError, + find_checkpoint_before_message, + find_checkpoint_before_message_chronologically, +) + +THREAD_ID = "thread-1" + + +def _snapshot(checkpoint_id: str, messages: list[object], *, next_tasks: tuple[str, ...] = (), parent_id: str | None = None): + parent_config = None + if parent_id is not None: + parent_config = {"configurable": {"thread_id": THREAD_ID, "checkpoint_ns": "", "checkpoint_id": parent_id}} + return SimpleNamespace( + values={"messages": messages}, + config={ + "configurable": { + "thread_id": THREAD_ID, + "checkpoint_ns": "", + "checkpoint_id": checkpoint_id, + "checkpoint_map": None, + } + }, + metadata={}, + parent_config=parent_config, + next=next_tasks, + ) + + +class _Accessor: + """Minimal accessor over a fixed snapshot list, keyed by checkpoint id.""" + + def __init__(self, snapshots: list[object]) -> None: + self.snapshots = snapshots + + async def aget(self, config): + checkpoint_id = config.get("configurable", {}).get("checkpoint_id") + return next( + (item for item in self.snapshots if item.config["configurable"]["checkpoint_id"] == checkpoint_id), + SimpleNamespace(values={}, config={}, metadata=None, parent_config=None, next=()), + ) + + +def _first_turn_history() -> list[object]: + """Newest-first history of a thread whose only turn is still its first. + + ``DynamicContextMiddleware`` swaps the first user message's id mid-run: + the injected reminder takes ``{id}`` and the real user message becomes + ``{id}__user``. So the pre-injection checkpoints hold the very same prompt + under an id the replay-base lookup does not recognise (#4531). + """ + system = SystemMessage(id="h1", content="date") + swapped_human = HumanMessage(id="h1__user", content="question") + raw_human = HumanMessage(id="h1", content="question") + ai = AIMessage(id="ai-1", content="answer") + return [ + _snapshot("ckpt-head", [system, swapped_human, ai], parent_id="ckpt-after-inject"), + _snapshot("ckpt-after-inject", [system, swapped_human], next_tasks=("LoopDetectionMiddleware.before_agent",), parent_id="ckpt-mid"), + _snapshot("ckpt-mid", [raw_human], next_tasks=("DynamicContextMiddleware.before_agent",), parent_id="ckpt-input"), + _snapshot("ckpt-input", [], next_tasks=("__start__",), parent_id="ckpt-empty"), + _snapshot("ckpt-empty", []), + ] + + +def test_chronological_scan_skips_checkpoints_with_pending_tasks(): + history = _first_turn_history() + + base, found = find_checkpoint_before_message_chronologically(history, "h1__user") + + assert found is True + assert base is not None + assert base.config["configurable"]["checkpoint_id"] == "ckpt-empty" + + +def test_lineage_walk_skips_checkpoints_with_pending_tasks(): + history = _first_turn_history() + accessor = _Accessor(history) + + base = asyncio.run(find_checkpoint_before_message(accessor, history[0], "h1__user", max_depth=10)) + + assert base.config["configurable"]["checkpoint_id"] == "ckpt-empty" + + +def test_chronological_scan_prefers_the_previous_turn_boundary(): + """A later turn resolves to the previous run's settled tail, not its input checkpoint.""" + human1 = HumanMessage(id="h1__user", content="first") + ai1 = AIMessage(id="ai-1", content="first answer") + human2 = HumanMessage(id="h2", content="second") + history = [ + _snapshot("ckpt-turn2-head", [human1, ai1, human2, AIMessage(id="ai-2", content="second answer")]), + _snapshot("ckpt-turn2-mid", [human1, ai1, human2], next_tasks=("model",)), + _snapshot("ckpt-turn2-input", [human1, ai1], next_tasks=("__start__",)), + _snapshot("ckpt-turn1-tail", [human1, ai1]), + ] + + base, found = find_checkpoint_before_message_chronologically(history, "h2") + + assert found is True + assert base.config["configurable"]["checkpoint_id"] == "ckpt-turn1-tail" + + +def test_lineage_walk_reports_missing_parent_when_no_settled_ancestor_exists(): + """Fail closed rather than fork a mid-run checkpoint.""" + human = HumanMessage(id="h1__user", content="question") + history = [ + _snapshot("ckpt-head", [human, AIMessage(id="ai-1", content="answer")], parent_id="ckpt-mid"), + _snapshot("ckpt-mid", [HumanMessage(id="h1", content="question")], next_tasks=("DynamicContextMiddleware.before_agent",)), + ] + accessor = _Accessor(history) + + with pytest.raises(CheckpointParentMissingError): + asyncio.run(find_checkpoint_before_message(accessor, history[0], "h1__user", max_depth=10)) + + +def test_unknown_pending_tasks_do_not_block_selection(): + """Raw full-mode reads cannot derive tasks; absence of evidence stays permissive.""" + human = HumanMessage(id="h1", content="question") + history = [ + SimpleNamespace( + values={"messages": [human]}, + config={"configurable": {"thread_id": THREAD_ID, "checkpoint_ns": "", "checkpoint_id": "ckpt-head"}}, + metadata={}, + ), + SimpleNamespace( + values={"messages": []}, + config={"configurable": {"thread_id": THREAD_ID, "checkpoint_ns": "", "checkpoint_id": "ckpt-base"}}, + metadata={}, + ), + ] + + base, found = find_checkpoint_before_message_chronologically(history, "h1") + + assert found is True + assert base.config["configurable"]["checkpoint_id"] == "ckpt-base" diff --git a/backend/tests/test_thread_regenerate_prepare.py b/backend/tests/test_thread_regenerate_prepare.py index 24ff667a9..56ed3e16e 100644 --- a/backend/tests/test_thread_regenerate_prepare.py +++ b/backend/tests/test_thread_regenerate_prepare.py @@ -7,7 +7,7 @@ from unittest.mock import AsyncMock, patch import pytest from fastapi import HTTPException -from langchain_core.messages import AIMessage, HumanMessage, ToolMessage +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage from langgraph.checkpoint.base import empty_checkpoint, uuid6 from langgraph.checkpoint.memory import InMemorySaver @@ -21,6 +21,7 @@ def _checkpoint( *, metadata: dict | None = None, goal: dict | None = None, + next_tasks: tuple[str, ...] = (), ): channel_values = {"messages": messages} if goal is not None: @@ -36,6 +37,7 @@ def _checkpoint( }, checkpoint={"channel_values": channel_values}, metadata=metadata or {}, + next=next_tasks, ) @@ -117,6 +119,7 @@ class FakeAccessor: config=checkpoint.config, metadata=checkpoint.metadata, parent_config=getattr(checkpoint, "parent_config", None), + next=getattr(checkpoint, "next", ()), ) async def aget(self, config): @@ -629,7 +632,30 @@ def test_prepare_edit_regenerate_payload_returns_new_human_and_edit_metadata(): } -def _edit_source_run_fixtures() -> tuple[FakeEventStore, FakeRunManager]: +def _first_turn_checkpointer() -> FakeCheckpointer: + """First-turn history where the user message's id is swapped mid-run. + + ``DynamicContextMiddleware`` moves the first user message to ``{id}__user`` + and gives ``{id}`` to the injected reminder, so every checkpoint written + before that node ran holds the same prompt under an id the replay-base + lookup cannot match (#4531). + """ + system = SystemMessage(id="human-1", content="date") + swapped_human = HumanMessage(id="human-1__user", content="original question") + raw_human = HumanMessage(id="human-1", content="original question") + ai = AIMessage(id="ai-1", content="answer v1") + return FakeCheckpointer( + [ + _checkpoint("ckpt-head", [system, swapped_human, ai]), + _checkpoint("ckpt-after-inject", [system, swapped_human], next_tasks=("LoopDetectionMiddleware.before_agent",)), + _checkpoint("ckpt-mid", [raw_human], next_tasks=("DynamicContextMiddleware.before_agent",)), + _checkpoint("ckpt-input", [], next_tasks=("__start__",)), + _checkpoint("ckpt-empty", []), + ] + ) + + +def _answer_run_fixtures() -> tuple[FakeEventStore, FakeRunManager]: event_store = FakeEventStore( [ { @@ -655,7 +681,7 @@ def test_prepare_edit_regenerate_payload_preserves_a_rename_the_replay_base_pred base.checkpoint["channel_values"]["title"] = "auto generated title" latest = _checkpoint("ckpt-ai", [human, ai]) latest.checkpoint["channel_values"]["title"] = "User renamed title" - event_store, run_manager = _edit_source_run_fixtures() + event_store, run_manager = _answer_run_fixtures() response = asyncio.run( thread_runs._prepare_edit_regenerate_payload( @@ -683,7 +709,7 @@ def test_prepare_edit_regenerate_payload_lets_an_untitled_base_name_the_edited_t ai = AIMessage(id="ai-1", content="answer v1") latest = _checkpoint("ckpt-ai", [human, ai]) latest.checkpoint["channel_values"]["title"] = "title of the replaced question" - event_store, run_manager = _edit_source_run_fixtures() + event_store, run_manager = _answer_run_fixtures() response = asyncio.run( thread_runs._prepare_edit_regenerate_payload( @@ -698,6 +724,83 @@ def test_prepare_edit_regenerate_payload_lets_an_untitled_base_name_the_edited_t assert "title" not in response.input +def test_prepare_regenerate_payload_replays_the_pre_swap_user_message_id(): + """Replaying `{id}__user` would make the reminder middleware skip the turn. + + The replay base predates the injection, so the turn must re-enter the graph + under the id the client originally sent or it loses its date/memory block. + """ + from app.gateway.routers import thread_runs + + event_store, run_manager = _answer_run_fixtures() + + response = asyncio.run( + thread_runs._prepare_regenerate_payload( + "thread-1", + "ai-1", + _request(_first_turn_checkpointer(), event_store, run_manager=run_manager), + ) + ) + + assert response.checkpoint["checkpoint_id"] == "ckpt-empty" + assert response.input["messages"][0]["id"] == "human-1" + + +def test_prepare_edit_regenerate_payload_skips_mid_run_replay_base_on_first_turn(): + """The replay base must predate the turn, not sit inside the run that produced it. + + ``ckpt-mid`` still holds the original prompt (under its pre-swap id) and owns + the injection node's pending writes, so replaying from it re-adds the prompt + the edit is meant to replace. + """ + from app.gateway.routers import thread_runs + + event_store, run_manager = _answer_run_fixtures() + + response = asyncio.run( + thread_runs._prepare_edit_regenerate_payload( + "thread-1", + "human-1__user", + "edited question", + _request(_first_turn_checkpointer(), event_store, run_manager=run_manager), + ) + ) + + assert response.checkpoint["checkpoint_id"] == "ckpt-empty" + assert response.metadata["regenerate_checkpoint_id"] == "ckpt-empty" + + +def test_prepare_edit_regenerate_payload_prefers_checkpoint_lineage(): + """Edit replay must resolve its base the same lineage-first way regenerate does. + + A chronological scan cannot tell sibling branches apart (#4358), so the head + checkpoint has to reach the lineage walk. + """ + from app.gateway.routers import thread_runs + + event_store, run_manager = _answer_run_fixtures() + checkpointer = _first_turn_checkpointer() + from app.gateway.checkpoint_lineage import CheckpointParentMissingError + + walk = AsyncMock(side_effect=CheckpointParentMissingError("no parent link")) + + with patch.object(thread_runs, "find_checkpoint_before_message", walk): + response = asyncio.run( + thread_runs._prepare_edit_regenerate_payload( + "thread-1", + "human-1__user", + "edited question", + _request(checkpointer, event_store, run_manager=run_manager), + ) + ) + + assert walk.await_count == 1 + head_checkpoint = walk.await_args.args[1] + assert head_checkpoint.config["configurable"]["checkpoint_id"] == "ckpt-head" + # Legacy checkpoints without parent links still degrade to the bounded scan. + assert response.checkpoint["checkpoint_id"] == "ckpt-empty" + + @pytest.mark.parametrize( ("replacement_text", "detail"), [ diff --git a/frontend/src/core/threads/hooks.ts b/frontend/src/core/threads/hooks.ts index 9399b357c..f10455c5e 100644 --- a/frontend/src/core/threads/hooks.ts +++ b/frontend/src/core/threads/hooks.ts @@ -918,6 +918,27 @@ function getMessagesAfterBaseline( }); } +/** + * Human-message baseline for a prepared replay (regenerate / edit). + * + * A replay masks the turn it supersedes, so those messages leave the live + * message list the moment the mask is applied. Baselining on the pre-mask count + * means the replacement only ever restores the count instead of exceeding it, + * and the optimistic copy is never recognised as confirmed. That matters for + * the first turn of a thread, where the runtime re-keys the replacement message + * so identity comparison cannot stand in for the count either. + */ +export function countHumanMessagesExcludingSuperseded( + messages: Message[], + supersededMessageIds: readonly string[], +): number { + const superseded = new Set(supersededMessageIds); + return messages.filter( + (message) => + message.type === "human" && (!message.id || !superseded.has(message.id)), + ).length; +} + export function getVisibleOptimisticMessages( optimisticMessages: Message[], previousHumanMessageCount: number, @@ -2068,6 +2089,10 @@ export function useThreadStream({ const prepared = await prepare(); preparedSupersededRunId = prepared.target_run_id; preparedSupersededMessageIds = getSupersededMessageIds(prepared); + prevHumanMsgCountRef.current = countHumanMessagesExcludingSuperseded( + persistedMessages, + preparedSupersededMessageIds, + ); const replacementHumanMessageId = "replacement_human_message_id" in prepared && typeof prepared.replacement_human_message_id === "string" diff --git a/frontend/tests/unit/core/threads/message-merge.test.ts b/frontend/tests/unit/core/threads/message-merge.test.ts index 643267efa..444a99874 100644 --- a/frontend/tests/unit/core/threads/message-merge.test.ts +++ b/frontend/tests/unit/core/threads/message-merge.test.ts @@ -7,6 +7,7 @@ import { buildVisibleHistoryMessages, areOptimisticMessagesConfirmed, computeSummarizationTransientMessages, + countHumanMessagesExcludingSuperseded, flattenThreadHistoryPages, getSummarizationMiddlewareMessages, getThreadHistoryNextPageParam, @@ -387,6 +388,58 @@ test("mergeMessages shows server human instead of optimistic duplicate after fir ]); }); +test("edit replay of the only turn hides the optimistic copy once the server human arrives", () => { + // The runtime re-keys the first user message of a thread, so the persisted + // replacement never matches the optimistic id and only the count can confirm + // it. Masking the superseded turn drops the live count to zero first. + const supersededHuman = { + id: "human-1__user", + type: "human", + content: "introduce Li Bai", + } as Message; + const optimisticHuman = { + id: "replacement-1", + type: "human", + content: "introduce Du Fu", + } as Message; + const serverHuman = { + id: "replacement-1__user", + type: "human", + content: "introduce Du Fu", + } as Message; + + const baseline = countHumanMessagesExcludingSuperseded( + [supersededHuman], + ["human-1__user", "ai-1"], + ); + expect(baseline).toBe(0); + + expect(getVisibleOptimisticMessages([optimisticHuman], baseline, 0)).toEqual([ + optimisticHuman, + ]); + expect(getVisibleOptimisticMessages([optimisticHuman], baseline, 1)).toEqual( + [], + ); + expect(mergeMessages([], [serverHuman], [])).toEqual([serverHuman]); +}); + +test("countHumanMessagesExcludingSuperseded keeps turns the replay does not supersede", () => { + const keptHuman = { id: "human-1", type: "human", content: "one" } as Message; + const supersededHuman = { + id: "human-2", + type: "human", + content: "two", + } as Message; + const ai = { id: "ai-1", type: "ai", content: "answer" } as Message; + + expect( + countHumanMessagesExcludingSuperseded( + [keptHuman, ai, supersededHuman], + ["human-2", "ai-2"], + ), + ).toBe(1); +}); + test("getVisibleOptimisticMessages keeps optimistic user input until server human arrives", () => { const optimisticHuman = { id: "opt-human-1",