deer-flow/backend/tests/test_thread_regenerate_prepare.py

from __future__ import annotations

import asyncio
import copy
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from _router_auth_helpers import call_unwrapped
from fastapi import HTTPException
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
from langgraph.checkpoint.base import empty_checkpoint, uuid6
from langgraph.checkpoint.memory import InMemorySaver

from deerflow.runtime import DisconnectMode, RunRecord, RunStatus
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY


def _checkpoint(
    checkpoint_id: str,
    messages: list[object],
    *,
    metadata: dict | None = None,
    goal: dict | None = None,
    next_tasks: tuple[str, ...] = (),
):
    channel_values = {"messages": messages}
    if goal is not None:
        channel_values["goal"] = goal
    return SimpleNamespace(
        config={
            "configurable": {
                "thread_id": "thread-1",
                "checkpoint_ns": "",
                "checkpoint_id": checkpoint_id,
                "checkpoint_map": None,
            }
        },
        checkpoint={"channel_values": channel_values},
        metadata=metadata or {},
        next=next_tasks,
    )


async def _put_memory_checkpoint(
    checkpointer: InMemorySaver,
    thread_id: str,
    messages: list[object],
    *,
    step: int,
    parent_config: dict | None = None,
    metadata: dict | None = None,
) -> dict:
    checkpoint = empty_checkpoint()
    checkpoint["id"] = str(uuid6())
    checkpoint["channel_values"] = {"messages": messages}
    checkpoint["channel_versions"] = {"messages": step}
    checkpoint_metadata = {
        "step": step,
        "source": "loop",
        "writes": {"test": {"messages": messages}},
        "parents": {},
    }
    checkpoint_metadata.update(metadata or {})
    return await checkpointer.aput(
        parent_config or {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}},
        checkpoint,
        checkpoint_metadata,
        {"messages": step},
    )


async def _collect_checkpoints(checkpointer: InMemorySaver, config: dict) -> list:
    return [checkpoint async for checkpoint in checkpointer.alist(config)]


class FakeCheckpointer:
    def __init__(self, history, *, latest=None, materialized_history=None, materialized_latest=None):
        self.history = history
        self.latest = latest
        self.materialized_history = materialized_history
        self.materialized_latest = materialized_latest
        self.alist_limits = []

    async def aget_tuple(self, config):
        checkpoint_id = config.get("configurable", {}).get("checkpoint_id")
        if checkpoint_id:
            return next((item for item in self.history if item.config["configurable"]["checkpoint_id"] == checkpoint_id), None)
        return self.latest or (self.history[0] if self.history else None)

    async def alist(self, config, limit=200):
        self.alist_limits.append(limit)
        for item in self.history[:limit]:
            yield item


def _snapshot(checkpoint_id: str, messages: list[object], *, metadata: dict | None = None):
    return SimpleNamespace(
        values={"messages": messages},
        config={
            "configurable": {
                "thread_id": "thread-1",
                "checkpoint_ns": "",
                "checkpoint_id": checkpoint_id,
                "checkpoint_map": None,
            }
        },
        metadata=metadata or {},
    )


def _human_input_response(
    value: str,
    *,
    request_id: str,
    source: str = "ask_clarification",
) -> dict:
    return {
        "version": 1,
        "kind": "human_input_response",
        "source": source,
        "request_id": request_id,
        "response_kind": "text",
        "value": value,
    }


class FakeAccessor:
    def __init__(self, checkpointer: FakeCheckpointer):
        self.checkpointer = checkpointer

    @staticmethod
    def _from_raw(checkpoint):
        return SimpleNamespace(
            values=dict(checkpoint.checkpoint.get("channel_values", {})),
            config=checkpoint.config,
            metadata=checkpoint.metadata,
            parent_config=getattr(checkpoint, "parent_config", None),
            next=getattr(checkpoint, "next", ()),
        )

    async def aget(self, config):
        materialized_latest = getattr(self.checkpointer, "materialized_latest", None)
        if materialized_latest is not None and not config.get("configurable", {}).get("checkpoint_id"):
            return materialized_latest
        raw = await self.checkpointer.aget_tuple(config)
        return self._from_raw(raw) if raw is not None else SimpleNamespace(values={}, config={}, metadata={})

    async def ahistory(self, config, *, limit=None):
        alist_limits = getattr(self.checkpointer, "alist_limits", None)
        if alist_limits is not None:
            alist_limits.append(limit)
        history = getattr(self.checkpointer, "materialized_history", None)
        if history is None:
            if hasattr(self.checkpointer, "history"):
                history = [self._from_raw(item) for item in self.checkpointer.history]
            else:
                history = [self._from_raw(item) async for item in self.checkpointer.alist(config, limit=limit)]
        return history[:limit]


@pytest.fixture(autouse=True)
def _patch_checkpoint_accessor(monkeypatch):
    from app.gateway import services
    from app.gateway.routers import thread_runs

    def build_accessor(request, *, thread_id, assistant_id=None, checkpoint_id=None):
        config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
        if checkpoint_id is not None:
            config["configurable"]["checkpoint_id"] = checkpoint_id
        return FakeAccessor(request.app.state.checkpointer), config

    async def build_thread_accessor(request, *, thread_id, checkpoint_id=None):
        return build_accessor(request, thread_id=thread_id, checkpoint_id=checkpoint_id)

    monkeypatch.setattr(services, "build_checkpoint_state_accessor", build_accessor)
    monkeypatch.setattr(thread_runs, "build_thread_checkpoint_state_accessor", build_thread_accessor)


class FakeEventStore:
    def __init__(self, rows):
        self.rows = rows

    async def list_messages(self, thread_id, *, limit=50, before_seq=None, after_seq=None, user_id=None):
        return self.rows[-limit:]


def _event_store_for_ai(
    content: str,
    *,
    message_id: str = "ai-1",
    run_id: str = "run-old",
) -> FakeEventStore:
    return FakeEventStore(
        [
            {
                "run_id": run_id,
                "event_type": "llm.ai.response",
                "category": "message",
                "content": {
                    "id": message_id,
                    "type": "ai",
                    "content": content,
                },
                "metadata": {"caller": "lead_agent"},
            }
        ]
    )


class FakeRunManager:
    def __init__(self, records):
        self.records = records

    async def list_by_thread(self, thread_id, *, user_id=None, limit=100):
        return self.records[:limit]

    async def get(self, run_id, *, user_id=None):
        return next((record for record in self.records if record.run_id == run_id), None)


def _request(checkpointer, event_store, *, run_manager=None, user_id="user-1"):
    from app.gateway.auth_disabled import AUTH_SOURCE_SESSION

    return SimpleNamespace(
        app=SimpleNamespace(
            state=SimpleNamespace(
                checkpointer=checkpointer,
                run_event_store=event_store,
                run_manager=run_manager or FakeRunManager([]),
            )
        ),
        state=SimpleNamespace(user=SimpleNamespace(id=user_id), auth_source=AUTH_SOURCE_SESSION),
    )


def test_run_wait_readers_return_materialized_final_values() -> None:
    from app.gateway import services
    from app.gateway.routers import runs, thread_runs

    snapshot = SimpleNamespace(
        values={
            "messages": [
                HumanMessage(id="h1", content="question"),
                AIMessage(id="a1", content="answer"),
            ]
        },
        config={
            "configurable": {
                "thread_id": "thread-1",
                "checkpoint_ns": "",
                "checkpoint_id": "ckpt-2",
            }
        },
        parent_config={"configurable": {"checkpoint_id": "ckpt-1"}},
        metadata={"step": 2},
        next=(),
        tasks=(),
        created_at=None,
    )
    accessor = SimpleNamespace(aget=AsyncMock(return_value=snapshot))
    record = SimpleNamespace(
        run_id="run-1",
        thread_id="thread-1",
        task=None,
        status=RunStatus.success,
        error=None,
    )
    request = SimpleNamespace()
    body = thread_runs.RunCreateRequest(
        assistant_id="lead-agent",
        config={"configurable": {"thread_id": "thread-1"}},
    )

    async def _scenario() -> tuple[dict, dict]:
        with (
            patch.object(thread_runs, "get_stream_bridge", return_value=object()),
            patch.object(thread_runs, "get_run_manager", return_value=object()),
            patch.object(thread_runs, "start_run", AsyncMock(return_value=record)),
            patch.object(
                services,
                "build_checkpoint_state_accessor",
                create=True,
                new=MagicMock(return_value=(accessor, snapshot.config)),
            ),
            patch.object(runs, "get_stream_bridge", return_value=object()),
            patch.object(runs, "get_run_manager", return_value=object()),
            patch.object(runs, "start_run", AsyncMock(return_value=record)),
            patch.object(
                services,
                "build_checkpoint_state_accessor",
                create=True,
                new=MagicMock(return_value=(accessor, snapshot.config)),
            ),
        ):
            thread_result = await call_unwrapped(thread_runs.wait_run, "thread-1", body, request)
            stateless_result = await call_unwrapped(runs.stateless_wait, body, request)
        return thread_result, stateless_result

    thread_result, stateless_result = asyncio.run(_scenario())

    assert [message["id"] for message in thread_result["messages"]] == ["h1", "a1"]
    assert [message["id"] for message in stateless_result["messages"]] == ["h1", "a1"]


def test_run_wait_readers_preserve_terminal_error_without_checkpoint() -> None:
    from app.gateway import services
    from app.gateway.routers import runs, thread_runs

    snapshot = SimpleNamespace(
        values={},
        config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}},
        parent_config=None,
        metadata={},
        next=(),
        tasks=(),
        created_at=None,
    )
    accessor = SimpleNamespace(aget=AsyncMock(return_value=snapshot))
    record = SimpleNamespace(
        run_id="run-1",
        thread_id="thread-1",
        task=None,
        status=RunStatus.error,
        error="run failed before checkpoint",
    )
    request = SimpleNamespace()
    body = thread_runs.RunCreateRequest(config={"configurable": {"thread_id": "thread-1"}})

    async def _scenario() -> tuple[dict, dict]:
        with (
            patch.object(thread_runs, "get_stream_bridge", return_value=object()),
            patch.object(thread_runs, "get_run_manager", return_value=object()),
            patch.object(thread_runs, "start_run", AsyncMock(return_value=record)),
            patch.object(
                services,
                "build_checkpoint_state_accessor",
                new=MagicMock(return_value=(accessor, snapshot.config)),
            ),
            patch.object(runs, "get_stream_bridge", return_value=object()),
            patch.object(runs, "get_run_manager", return_value=object()),
            patch.object(runs, "start_run", AsyncMock(return_value=record)),
            patch.object(
                services,
                "build_checkpoint_state_accessor",
                new=MagicMock(return_value=(accessor, snapshot.config)),
            ),
        ):
            thread_result = await call_unwrapped(thread_runs.wait_run, "thread-1", body, request)
            stateless_result = await call_unwrapped(runs.stateless_wait, body, request)
        return thread_result, stateless_result

    thread_result, stateless_result = asyncio.run(_scenario())

    expected = {
        "status": "error",
        "error": "run failed before checkpoint",
        "__error__": {"error": "RunError", "message": "run failed before checkpoint"},
    }
    assert thread_result == expected
    assert stateless_result == expected


@pytest.mark.parametrize("route_name", ["thread", "stateless"])
def test_run_wait_readers_preserve_terminal_error_when_accessor_builder_fails(route_name: str) -> None:
    from app.gateway import services
    from app.gateway.routers import runs, thread_runs

    record = SimpleNamespace(
        run_id="run-1",
        thread_id="thread-1",
        task=None,
        status=RunStatus.error,
        error="run failed before checkpoint",
    )
    request = SimpleNamespace()
    body = thread_runs.RunCreateRequest(config={"configurable": {"thread_id": "thread-1"}})

    async def _scenario() -> dict:
        if route_name == "thread":
            with (
                patch.object(thread_runs, "get_stream_bridge", return_value=object()),
                patch.object(thread_runs, "get_run_manager", return_value=object()),
                patch.object(thread_runs, "start_run", AsyncMock(return_value=record)),
                patch.object(
                    services,
                    "build_checkpoint_state_accessor",
                    new=MagicMock(side_effect=RuntimeError("graph construction failed")),
                ),
            ):
                return await call_unwrapped(thread_runs.wait_run, "thread-1", body, request)

        with (
            patch.object(runs, "get_stream_bridge", return_value=object()),
            patch.object(runs, "get_run_manager", return_value=object()),
            patch.object(runs, "start_run", AsyncMock(return_value=record)),
            patch.object(
                services,
                "build_checkpoint_state_accessor",
                new=MagicMock(side_effect=RuntimeError("graph construction failed")),
            ),
        ):
            return await call_unwrapped(runs.stateless_wait, body, request)

    result = asyncio.run(_scenario())

    expected = {
        "status": "error",
        "error": "run failed before checkpoint",
        "__error__": {"error": "RunError", "message": "run failed before checkpoint"},
    }
    assert result == expected


def test_prepare_regenerate_payload_returns_clean_input_and_base_checkpoint():
    from app.gateway.routers import thread_runs

    human = HumanMessage(
        id="human-1",
        content="<current_uploads>injected</current_uploads>\n\n/data-analysis analyze data.csv",
        additional_kwargs={
            ORIGINAL_USER_CONTENT_KEY: "/data-analysis analyze data.csv",
            "files": [{"filename": "data.csv", "path": "/mnt/user-data/uploads/data.csv"}],
        },
    )
    ai = AIMessage(id="ai-1", content="answer v1")
    base = _checkpoint("ckpt-base", [])
    after_human = _checkpoint("ckpt-human", [human])
    latest = _checkpoint("ckpt-ai", [human, ai])
    checkpointer = FakeCheckpointer([latest, after_human, base])
    event_store = FakeEventStore(
        [
            {
                "run_id": "run-old",
                "event_type": "llm.ai.response",
                "category": "message",
                "content": {"id": "ai-1", "type": "ai", "content": "answer v1"},
                "metadata": {"caller": "lead_agent"},
            }
        ]
    )

    original_builder = thread_runs.build_thread_checkpoint_state_accessor
    thread_builder = AsyncMock(side_effect=original_builder)
    with patch.object(thread_runs, "build_thread_checkpoint_state_accessor", thread_builder):
        response = asyncio.run(thread_runs._prepare_regenerate_payload("thread-1", "ai-1", _request(checkpointer, event_store)))

    assert [call.kwargs["thread_id"] for call in thread_builder.await_args_list] == ["thread-1", "thread-1"]

    assert response.checkpoint == {
        "checkpoint_ns": "",
        "checkpoint_id": "ckpt-base",
        "checkpoint_map": None,
    }
    assert response.target_run_id == "run-old"
    assert response.metadata == {
        "regenerate_from_message_id": "ai-1",
        "regenerate_from_run_id": "run-old",
        "regenerate_checkpoint_id": "ckpt-base",
    }
    assert "title" not in response.input
    regenerated_human = response.input["messages"][0]
    assert regenerated_human["id"] == "human-1"
    assert regenerated_human["content"] == [{"type": "text", "text": "/data-analysis analyze data.csv"}]
    assert regenerated_human["additional_kwargs"] == {"files": [{"filename": "data.csv", "path": "/mnt/user-data/uploads/data.csv"}]}


@pytest.mark.parametrize("materialized_delta", [False, True], ids=["full", "delta"])
def test_prepare_regenerate_payload_replays_latest_confirmed_human_input_response(
    materialized_delta: bool,
):
    from app.gateway.routers import thread_runs

    original_human = HumanMessage(id="human-1", content="Help me plan a trip")
    first_clarification = ToolMessage(
        id="tool-city",
        tool_call_id="call-city",
        content="Which city?",
    )
    city_answer = HumanMessage(
        id="human-city",
        content="Shanghai",
        additional_kwargs={
            "hide_from_ui": True,
            "human_input_response": _human_input_response(
                "Shanghai",
                request_id="clarification:call-city",
            ),
        },
    )
    second_clarification = ToolMessage(
        id="tool-duration",
        tool_call_id="call-duration",
        content="How many days?",
    )
    duration_metadata = _human_input_response(
        "Five days",
        request_id="clarification:call-duration",
    )
    duration_answer = HumanMessage(
        id="human-duration",
        content="Five days",
        additional_kwargs={
            "hide_from_ui": True,
            "human_input_response": duration_metadata,
        },
    )
    ai = AIMessage(id="ai-1", content="Here is your five-day Shanghai itinerary")
    before_question: list[object] = []
    before_city = [original_human, first_clarification]
    after_city = [*before_city, city_answer]
    before_duration = [*after_city, second_clarification]
    after_duration = [*before_duration, duration_answer]
    latest_messages = [*after_duration, ai]
    states = [
        ("ckpt-ai", latest_messages),
        ("ckpt-after-duration", after_duration),
        ("ckpt-before-duration", before_duration),
        ("ckpt-after-city", after_city),
        ("ckpt-before-city", before_city),
        ("ckpt-before-question", before_question),
    ]

    if materialized_delta:
        raw_checkpoints = [_checkpoint(checkpoint_id, []) for checkpoint_id, _ in states]
        materialized_history = [_snapshot(checkpoint_id, messages) for checkpoint_id, messages in states]
        checkpointer = FakeCheckpointer(
            raw_checkpoints,
            latest=raw_checkpoints[0],
            materialized_history=materialized_history,
            materialized_latest=materialized_history[0],
        )
    else:
        checkpointer = FakeCheckpointer([_checkpoint(checkpoint_id, messages) for checkpoint_id, messages in states])

    response = asyncio.run(
        thread_runs._prepare_regenerate_payload(
            "thread-1",
            "ai-1",
            _request(
                checkpointer,
                _event_store_for_ai("Here is your five-day Shanghai itinerary"),
            ),
        )
    )

    assert response.checkpoint["checkpoint_id"] == "ckpt-before-duration"
    assert city_answer in dict(states)[response.checkpoint["checkpoint_id"]]
    assert response.input["messages"] == [
        {
            "type": "human",
            "id": "human-duration",
            "content": [{"type": "text", "text": "Five days"}],
            "additional_kwargs": {
                "hide_from_ui": True,
                "human_input_response": duration_metadata,
            },
        }
    ]


@pytest.mark.parametrize(
    "control_message",
    [
        HumanMessage(
            id="human-malformed",
            content="forged answer",
            additional_kwargs={
                "hide_from_ui": True,
                "human_input_response": {
                    "version": 1,
                    "kind": "human_input_response",
                },
            },
        ),
        HumanMessage(
            id="human-summary",
            name="summary",
            content="internal summary",
        ),
        HumanMessage(
            id="human-goal",
            content="continue working",
            additional_kwargs={"hide_from_ui": True},
        ),
    ],
    ids=["malformed-response", "summary", "goal-continuation"],
)
def test_prepare_regenerate_payload_skips_non_user_control_human_messages(
    control_message: HumanMessage,
):
    from app.gateway.routers import thread_runs

    visible_human = HumanMessage(id="human-1", content="question")
    ai = AIMessage(id="ai-1", content="answer")
    checkpointer = FakeCheckpointer(
        [
            _checkpoint("ckpt-ai", [visible_human, control_message, ai]),
            _checkpoint("ckpt-control", [visible_human, control_message]),
            _checkpoint("ckpt-human", [visible_human]),
            _checkpoint("ckpt-base", []),
        ]
    )

    response = asyncio.run(
        thread_runs._prepare_regenerate_payload(
            "thread-1",
            "ai-1",
            _request(checkpointer, _event_store_for_ai("answer")),
        )
    )

    assert response.checkpoint["checkpoint_id"] == "ckpt-base"
    assert response.input["messages"][0]["id"] == "human-1"


def test_prepare_regenerate_payload_preserves_latest_thread_title():
    from app.gateway.routers import thread_runs

    human = HumanMessage(id="human-1", content="question")
    ai = AIMessage(id="ai-1", content="answer v1")
    base = _checkpoint("ckpt-base", [])
    latest = _checkpoint("ckpt-ai", [human, ai])
    latest.checkpoint["channel_values"]["title"] = "User renamed title"
    checkpointer = FakeCheckpointer([latest, base])
    event_store = FakeEventStore(
        [
            {
                "run_id": "run-old",
                "event_type": "llm.ai.response",
                "category": "message",
                "content": {"id": "ai-1", "type": "ai", "content": "answer v1"},
                "metadata": {"caller": "lead_agent"},
            }
        ]
    )

    response = asyncio.run(
        thread_runs._prepare_regenerate_payload(
            "thread-1",
            "ai-1",
            _request(checkpointer, event_store),
        )
    )

    assert response.checkpoint["checkpoint_id"] == "ckpt-base"
    assert response.input["title"] == "User renamed title"


def test_prepare_regenerate_payload_supports_latest_interrupted_response_missing_from_checkpoint():
    from app.gateway.routers.thread_runs import _prepare_regenerate_payload

    human = HumanMessage(
        id="human-1",
        content="question",
        additional_kwargs={"run_id": "run-interrupted"},
    )
    base = _checkpoint("ckpt-base", [])
    latest = _checkpoint("ckpt-human", [human])
    checkpointer = FakeCheckpointer([latest, base])
    run_manager = FakeRunManager(
        [
            SimpleNamespace(
                run_id="run-interrupted",
                thread_id="thread-1",
                status=RunStatus.interrupted,
            )
        ]
    )

    response = asyncio.run(
        _prepare_regenerate_payload(
            "thread-1",
            "lc_run--partial-response",
            _request(
                checkpointer,
                FakeEventStore([]),
                run_manager=run_manager,
            ),
        )
    )

    assert response.checkpoint["checkpoint_id"] == "ckpt-base"
    assert response.target_run_id == "run-interrupted"
    assert response.metadata == {
        "regenerate_from_message_id": "lc_run--partial-response",
        "regenerate_from_run_id": "run-interrupted",
        "regenerate_checkpoint_id": "ckpt-base",
    }
    assert response.input["messages"][0]["id"] == "human-1"


@pytest.mark.parametrize(
    ("status", "record_thread_id"),
    [
        (RunStatus.success, "thread-1"),
        (RunStatus.interrupted, "another-thread"),
    ],
)
def test_prepare_regenerate_payload_does_not_accept_unverified_missing_response(
    status: RunStatus,
    record_thread_id: str,
):
    from app.gateway.routers.thread_runs import _prepare_regenerate_payload

    human = HumanMessage(
        id="human-1",
        content="question",
        additional_kwargs={"run_id": "source-run"},
    )
    checkpointer = FakeCheckpointer(
        [
            _checkpoint("ckpt-human", [human]),
            _checkpoint("ckpt-base", []),
        ]
    )
    run_manager = FakeRunManager(
        [
            SimpleNamespace(
                run_id="source-run",
                thread_id=record_thread_id,
                status=status,
            )
        ]
    )

    with pytest.raises(HTTPException) as exc:
        asyncio.run(
            _prepare_regenerate_payload(
                "thread-1",
                "missing-response",
                _request(
                    checkpointer,
                    FakeEventStore([]),
                    run_manager=run_manager,
                ),
            )
        )

    assert exc.value.status_code == 404
    assert exc.value.detail == "Message missing-response not found"


def test_prepare_regenerate_payload_does_not_mutate_legacy_single_checkpoint_branch():
    from app.gateway.routers.thread_runs import _prepare_regenerate_payload

    checkpointer = InMemorySaver()
    source_thread_id = "source-thread"
    branch_thread_id = "legacy-branch"
    source_run_id = "source-run"
    human = HumanMessage(id="human-1", content="question", additional_kwargs={"run_id": source_run_id})
    ai = AIMessage(id="ai-1", content="answer")

    async def _seed() -> str:
        source_base_config = await _put_memory_checkpoint(checkpointer, source_thread_id, [], step=0)
        after_human = await _put_memory_checkpoint(
            checkpointer,
            source_thread_id,
            [human],
            step=1,
            parent_config=source_base_config,
        )
        source_head_config = await _put_memory_checkpoint(
            checkpointer,
            source_thread_id,
            [human, ai],
            step=2,
            parent_config=after_human,
        )
        source_head = await checkpointer.aget_tuple(source_head_config)
        assert source_head is not None

        legacy_head = copy.deepcopy(source_head.checkpoint)
        legacy_head_id = str(uuid6())
        legacy_head["id"] = legacy_head_id
        legacy_metadata = copy.deepcopy(source_head.metadata)
        legacy_metadata.update(
            {
                "source": "branch",
                "deerflow_branch": True,
                "branch_parent_thread_id": source_thread_id,
                "branch_parent_checkpoint_id": source_head_config["configurable"]["checkpoint_id"],
                "branch_parent_message_id": "ai-1",
            }
        )
        await checkpointer.aput(
            {"configurable": {"thread_id": branch_thread_id, "checkpoint_ns": ""}},
            legacy_head,
            legacy_metadata,
            dict(legacy_head["channel_versions"]),
        )
        return legacy_head_id

    legacy_head_id = asyncio.run(_seed())
    request = _request(checkpointer, FakeEventStore([]))

    with pytest.raises(HTTPException) as exc:
        asyncio.run(_prepare_regenerate_payload(branch_thread_id, "ai-1", request))

    assert exc.value.status_code == 409
    assert exc.value.detail == "Could not find an addressable checkpoint before the target user message"
    latest = asyncio.run(checkpointer.aget_tuple({"configurable": {"thread_id": branch_thread_id, "checkpoint_ns": ""}}))
    assert latest is not None
    assert latest.config["configurable"]["checkpoint_id"] == legacy_head_id
    branch_history = asyncio.run(_collect_checkpoints(checkpointer, {"configurable": {"thread_id": branch_thread_id, "checkpoint_ns": ""}}))
    assert [item.config["configurable"]["checkpoint_id"] for item in branch_history] == [legacy_head_id]


def test_prepare_regenerate_payload_rejects_legacy_branch_when_source_checkpoint_is_missing():
    from app.gateway.routers.thread_runs import _prepare_regenerate_payload

    checkpointer = InMemorySaver()
    branch_thread_id = "legacy-orphan"
    human = HumanMessage(id="human-1", content="question", additional_kwargs={"run_id": "source-run"})
    ai = AIMessage(id="ai-1", content="answer")

    async def _seed() -> None:
        await _put_memory_checkpoint(
            checkpointer,
            branch_thread_id,
            [human, ai],
            step=1,
            metadata={
                "source": "branch",
                "deerflow_branch": True,
                "branch_parent_thread_id": "deleted-source",
                "branch_parent_checkpoint_id": "missing-checkpoint",
                "branch_parent_message_id": "ai-1",
            },
        )

    asyncio.run(_seed())
    with pytest.raises(HTTPException) as exc:
        asyncio.run(
            _prepare_regenerate_payload(
                branch_thread_id,
                "ai-1",
                _request(checkpointer, FakeEventStore([])),
            )
        )

    assert exc.value.status_code == 409
    assert exc.value.detail == "Could not find an addressable checkpoint before the target user message"
    branch_history = asyncio.run(_collect_checkpoints(checkpointer, {"configurable": {"thread_id": branch_thread_id, "checkpoint_ns": ""}}))
    assert len(branch_history) == 1


def test_prepare_edit_regenerate_payload_returns_new_human_and_edit_metadata():
    from app.gateway.routers import thread_runs

    human = HumanMessage(
        id="human-1",
        content="<current_uploads>injected</current_uploads>\n\noriginal question",
        name="researcher",
        additional_kwargs={
            ORIGINAL_USER_CONTENT_KEY: "original question",
            "files": [{"filename": "data.csv", "path": "/mnt/user-data/uploads/data.csv"}],
            "referenced_message_contexts": [{"message_id": "ai-prev", "quote": "quoted"}],
            "hide_from_ui": False,
            "run_id": "old-run",
            "timestamp": "2026-07-22T00:00:00Z",
            "middleware_private": "do-not-copy",
        },
    )
    ai = AIMessage(id="ai-1", content="answer v1")
    base = _checkpoint("ckpt-base", [])
    after_human = _checkpoint("ckpt-human", [human])
    latest = _checkpoint("ckpt-ai", [human, ai])
    checkpointer = FakeCheckpointer([latest, after_human, base])
    event_store = FakeEventStore(
        [
            {
                "run_id": "run-old",
                "event_type": "llm.ai.response",
                "category": "message",
                "content": {"id": "ai-1", "type": "ai", "content": "answer v1"},
                "metadata": {"caller": "lead_agent"},
            }
        ]
    )
    run_manager = FakeRunManager(
        [
            SimpleNamespace(
                run_id="run-old",
                status=RunStatus.success,
                metadata={},
                last_ai_message="answer v1",
            )
        ]
    )

    response = asyncio.run(
        thread_runs._prepare_edit_regenerate_payload(
            "thread-1",
            "human-1",
            "  updated question\nwith details  ",
            _request(checkpointer, event_store, run_manager=run_manager),
        )
    )

    assert response.checkpoint == {
        "checkpoint_ns": "",
        "checkpoint_id": "ckpt-base",
        "checkpoint_map": None,
    }
    assert response.target_run_id == "run-old"
    assert response.replacement_human_message_id != "human-1"
    assert response.source_message_ids == ["human-1", "ai-1"]
    assert response.metadata == {
        "replay_kind": "edit",
        "regenerate_from_message_id": "ai-1",
        "regenerate_from_run_id": "run-old",
        "regenerate_checkpoint_id": "ckpt-base",
        "edit_from_message_id": "human-1",
        "edit_message_id": response.replacement_human_message_id,
        "edit_version_group_id": "human-1",
    }
    replacement = response.input["messages"][0]
    assert replacement == {
        "type": "human",
        "id": response.replacement_human_message_id,
        "name": "researcher",
        "content": [{"type": "text", "text": "updated question\nwith details"}],
        "additional_kwargs": {
            "files": [{"filename": "data.csv", "path": "/mnt/user-data/uploads/data.csv"}],
            "referenced_message_contexts": [{"message_id": "ai-prev", "quote": "quoted"}],
        },
    }


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="<system-reminder>date</system-reminder>")
    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(
        [
            {
                "run_id": "run-old",
                "event_type": "llm.ai.response",
                "category": "message",
                "content": {"id": "ai-1", "type": "ai", "content": "answer v1"},
                "metadata": {"caller": "lead_agent"},
            }
        ]
    )
    run_manager = FakeRunManager([SimpleNamespace(run_id="run-old", status=RunStatus.success, metadata={}, last_ai_message="answer v1")])
    return event_store, run_manager


def test_prepare_edit_regenerate_payload_preserves_a_rename_the_replay_base_predates():
    """Replaying a turn must not roll the thread back to an older title (#4457)."""
    from app.gateway.routers import thread_runs

    human = HumanMessage(id="human-1", content="original question")
    ai = AIMessage(id="ai-1", content="answer v1")
    base = _checkpoint("ckpt-base", [])
    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 = _answer_run_fixtures()

    response = asyncio.run(
        thread_runs._prepare_edit_regenerate_payload(
            "thread-1",
            "human-1",
            "edited question",
            _request(FakeCheckpointer([latest, base]), event_store, run_manager=run_manager),
        )
    )

    assert response.checkpoint["checkpoint_id"] == "ckpt-base"
    assert response.input["title"] == "User renamed title"


def test_prepare_edit_regenerate_payload_lets_an_untitled_base_name_the_edited_turn():
    """A base with no title belongs to a thread that has not been named yet.

    Pinning the current title there would keep a name generated from the prompt
    the edit just replaced, so leave the channel empty and let the title
    middleware name the rewritten turn.
    """
    from app.gateway.routers import thread_runs

    human = HumanMessage(id="human-1", content="original question")
    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 = _answer_run_fixtures()

    response = asyncio.run(
        thread_runs._prepare_edit_regenerate_payload(
            "thread-1",
            "human-1",
            "edited question",
            _request(FakeCheckpointer([latest, _checkpoint("ckpt-base", [])]), event_store, run_manager=run_manager),
        )
    )

    assert response.checkpoint["checkpoint_id"] == "ckpt-base"
    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"),
    [
        ("  \n\t", "Edited message cannot be empty"),
        ("  original question  ", "Edited message is unchanged"),
    ],
)
def test_prepare_edit_regenerate_payload_rejects_empty_or_unchanged_text(replacement_text: str, detail: str):
    from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload

    human = HumanMessage(id="human-1", content="original question")
    ai = AIMessage(id="ai-1", content="answer")
    checkpointer = FakeCheckpointer(
        [
            _checkpoint("ckpt-ai", [human, ai]),
            _checkpoint("ckpt-human", [human]),
            _checkpoint("ckpt-base", []),
        ]
    )
    event_store = FakeEventStore(
        [
            {
                "run_id": "run-old",
                "event_type": "llm.ai.response",
                "category": "message",
                "content": {"id": "ai-1", "type": "ai", "content": "answer"},
                "metadata": {"caller": "lead_agent"},
            }
        ]
    )
    run_manager = FakeRunManager([SimpleNamespace(run_id="run-old", status=RunStatus.success, metadata={}, last_ai_message="answer")])

    with pytest.raises(HTTPException) as exc:
        asyncio.run(
            _prepare_edit_regenerate_payload(
                "thread-1",
                "human-1",
                replacement_text,
                _request(checkpointer, event_store, run_manager=run_manager),
            )
        )

    assert exc.value.status_code == 409
    assert exc.value.detail == detail


def test_prepare_edit_regenerate_payload_requires_latest_human_turn():
    from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload

    old_human = HumanMessage(id="human-old", content="old question")
    old_ai = AIMessage(id="ai-old", content="old answer")
    latest_human = HumanMessage(id="human-latest", content="latest question")
    latest_ai = AIMessage(id="ai-latest", content="latest answer")
    checkpointer = FakeCheckpointer(
        [
            _checkpoint("ckpt-latest", [old_human, old_ai, latest_human, latest_ai]),
            _checkpoint("ckpt-base", []),
        ]
    )

    with pytest.raises(HTTPException) as exc:
        asyncio.run(
            _prepare_edit_regenerate_payload(
                "thread-1",
                "human-old",
                "edited old question",
                _request(checkpointer, FakeEventStore([])),
            )
        )

    assert exc.value.status_code == 409
    assert exc.value.detail == "Only the latest completed user turn can be edited"


def test_prepare_edit_regenerate_payload_allows_answered_historical_clarification():
    from app.gateway.routers import thread_runs

    old_human = HumanMessage(id="human-old", content="ambiguous request")
    clarification = ToolMessage(
        id="tool-clarify",
        tool_call_id="call-clarify",
        content="need input",
        artifact={
            "human_input": {
                "version": 1,
                "kind": "human_input_request",
                "request_id": "clarify-1",
                "prompt": "Which format?",
            }
        },
    )
    clarification_answer = HumanMessage(
        id="human-clarify-answer",
        content="Use Markdown",
        additional_kwargs={
            "hide_from_ui": True,
            "human_input_response": {
                "version": 1,
                "kind": "human_input_response",
                "request_id": "clarify-1",
                "source": "user",
                "value": "Use Markdown",
            },
        },
    )
    latest_human = HumanMessage(id="human-latest", content="latest question")
    latest_ai = AIMessage(id="ai-latest", content="latest answer")
    historical_messages = [old_human, clarification, clarification_answer]
    checkpointer = FakeCheckpointer(
        [
            _checkpoint("ckpt-ai", [*historical_messages, latest_human, latest_ai]),
            _checkpoint("ckpt-human", [*historical_messages, latest_human]),
            _checkpoint("ckpt-base", historical_messages),
        ]
    )
    event_store = FakeEventStore(
        [
            {
                "run_id": "run-latest",
                "event_type": "llm.ai.response",
                "category": "message",
                "content": {"id": "ai-latest", "type": "ai", "content": "latest answer"},
                "metadata": {"caller": "lead_agent"},
            }
        ]
    )
    run_manager = FakeRunManager(
        [
            SimpleNamespace(
                run_id="run-latest",
                status=RunStatus.success,
                metadata={},
                last_ai_message="latest answer",
            )
        ]
    )

    response = asyncio.run(
        thread_runs._prepare_edit_regenerate_payload(
            "thread-1",
            "human-latest",
            "updated latest question",
            _request(checkpointer, event_store, run_manager=run_manager),
        )
    )

    assert response.checkpoint["checkpoint_id"] == "ckpt-base"
    assert response.target_run_id == "run-latest"
    assert response.source_message_ids == ["human-latest", "ai-latest"]


def test_prepare_edit_regenerate_payload_rejects_open_clarification_turn():
    from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload

    human = HumanMessage(id="human-1", content="question")
    tool = ToolMessage(
        id="tool-1",
        tool_call_id="call-1",
        content="need input",
        artifact={"human_input": {"request_id": "clarify-1", "status": "pending"}},
    )
    checkpointer = FakeCheckpointer(
        [
            _checkpoint("ckpt-tool", [human, tool]),
            _checkpoint("ckpt-human", [human]),
            _checkpoint("ckpt-base", []),
        ]
    )

    with pytest.raises(HTTPException) as exc:
        asyncio.run(
            _prepare_edit_regenerate_payload(
                "thread-1",
                "human-1",
                "updated question",
                _request(checkpointer, FakeEventStore([])),
            )
        )

    assert exc.value.status_code == 409
    assert exc.value.detail == "Only completed assistant text turns can be edited"


def test_prepare_edit_regenerate_payload_rejects_active_goal():
    from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload

    human = HumanMessage(id="human-1", content="question")
    ai = AIMessage(id="ai-1", content="answer")
    checkpointer = FakeCheckpointer(
        [
            _checkpoint("ckpt-ai", [human, ai], goal={"status": "active", "objective": "finish"}),
            _checkpoint("ckpt-human", [human]),
            _checkpoint("ckpt-base", []),
        ]
    )

    with pytest.raises(HTTPException) as exc:
        asyncio.run(
            _prepare_edit_regenerate_payload(
                "thread-1",
                "human-1",
                "updated question",
                _request(checkpointer, FakeEventStore([])),
            )
        )

    assert exc.value.status_code == 409
    assert exc.value.detail == "Cannot edit while a goal is active"


@pytest.mark.parametrize(
    "latest_values",
    [
        # Clearing a met goal leaves only its record.
        {"goal_outcome": {"status": "achieved", "objective": "finish", "reply_message_id": "ai-1"}},
        # POST /state stores an unbuilt goal the history head does not show as active.
        {"goal": {"objective": "finish"}},
    ],
    ids=["met-goal-record", "goal-without-status"],
)
def test_prepare_edit_regenerate_payload_allows_editing_without_an_active_goal(latest_values):
    from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload

    human = HumanMessage(id="human-1", content="question")
    ai = AIMessage(id="ai-1", content="answer v1")
    latest = _checkpoint("ckpt-ai", [human, ai])
    latest.checkpoint["channel_values"].update(latest_values)
    event_store, run_manager = _answer_run_fixtures()

    response = asyncio.run(
        _prepare_edit_regenerate_payload(
            "thread-1",
            "human-1",
            "updated question",
            _request(FakeCheckpointer([latest, _checkpoint("ckpt-human", [human]), _checkpoint("ckpt-base", [])]), event_store, run_manager=run_manager),
        )
    )

    assert response.checkpoint["checkpoint_id"] == "ckpt-base"


def test_prepare_edit_regenerate_payload_requires_successful_source_run():
    from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload

    human = HumanMessage(id="human-1", content="question")
    ai = AIMessage(id="ai-1", content="answer")
    checkpointer = FakeCheckpointer(
        [
            _checkpoint("ckpt-ai", [human, ai]),
            _checkpoint("ckpt-human", [human]),
            _checkpoint("ckpt-base", []),
        ]
    )
    event_store = FakeEventStore(
        [
            {
                "run_id": "run-old",
                "event_type": "llm.ai.response",
                "category": "message",
                "content": {"id": "ai-1", "type": "ai", "content": "answer"},
                "metadata": {"caller": "lead_agent"},
            }
        ]
    )
    run_manager = FakeRunManager([SimpleNamespace(run_id="run-old", status=RunStatus.error, metadata={}, last_ai_message="answer")])

    with pytest.raises(HTTPException) as exc:
        asyncio.run(
            _prepare_edit_regenerate_payload(
                "thread-1",
                "human-1",
                "updated question",
                _request(checkpointer, event_store, run_manager=run_manager),
            )
        )

    assert exc.value.status_code == 409
    assert exc.value.detail == "Only successful assistant runs can be edited and rerun"


def test_prepare_regenerate_uses_materialized_history_when_raw_messages_are_omitted():
    from app.gateway.routers.thread_runs import _prepare_regenerate_payload

    earlier_human = HumanMessage(id="human-0", content="earlier question")
    earlier_ai = AIMessage(id="ai-0", content="earlier answer")
    target_human = HumanMessage(id="human-1", content="question")
    target_ai = AIMessage(id="ai-1", content="answer")

    raw_latest = _checkpoint("ckpt-ai", [])
    raw_after_human = _checkpoint("ckpt-human", [])
    raw_base = _checkpoint("ckpt-base", [])
    materialized_history = [
        _snapshot("ckpt-ai", [earlier_human, earlier_ai, target_human, target_ai]),
        _snapshot("ckpt-human", [earlier_human, earlier_ai, target_human]),
        _snapshot("ckpt-base", [earlier_human, earlier_ai]),
    ]
    checkpointer = FakeCheckpointer(
        [raw_latest, raw_after_human, raw_base],
        latest=raw_latest,
        materialized_history=materialized_history,
        materialized_latest=materialized_history[0],
    )
    event_store = FakeEventStore(
        [
            {
                "run_id": "run-target",
                "event_type": "ai_message",
                "category": "message",
                "content": {"id": "ai-1", "type": "ai", "content": "answer"},
                "metadata": {"caller": "lead_agent"},
            }
        ]
    )

    response = asyncio.run(
        _prepare_regenerate_payload(
            "thread-1",
            "ai-1",
            _request(checkpointer, event_store),
        )
    )

    assert response.checkpoint["checkpoint_id"] == "ckpt-base"
    assert response.metadata["regenerate_checkpoint_id"] == "ckpt-base"
    assert response.input["messages"][0]["id"] == "human-1"
    assert checkpointer.alist_limits == [400]


def test_prepare_regenerate_rejects_cyclic_lineage_without_chronological_fallback():
    from app.gateway.routers import thread_runs

    human = HumanMessage(id="human-1", content="question")
    ai = AIMessage(id="ai-1", content="answer")

    def linked_snapshot(checkpoint_id: str, messages: list[object], parent_id: str) -> SimpleNamespace:
        return SimpleNamespace(
            values={"messages": messages},
            config={
                "configurable": {
                    "thread_id": "thread-1",
                    "checkpoint_ns": "",
                    "checkpoint_id": checkpoint_id,
                }
            },
            metadata={},
            parent_config={
                "configurable": {
                    "thread_id": "thread-1",
                    "checkpoint_ns": "",
                    "checkpoint_id": parent_id,
                }
            },
        )

    head = linked_snapshot("head", [human, ai], "cycle")
    cycle = linked_snapshot("cycle", [human], "head")
    wrong_sibling_base = linked_snapshot("wrong-sibling", [], "root")
    by_id = {"head": head, "cycle": cycle}

    async def aget(config):
        return by_id[config["configurable"]["checkpoint_id"]]

    accessor = SimpleNamespace(
        aget=aget,
        ahistory=AsyncMock(return_value=[head, wrong_sibling_base]),
    )
    builder = AsyncMock(
        return_value=(
            accessor,
            {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}},
        )
    )

    with patch.object(thread_runs, "build_thread_checkpoint_state_accessor", builder):
        with pytest.raises(HTTPException) as exc:
            asyncio.run(
                thread_runs._find_base_checkpoint_before_human(
                    "thread-1",
                    "human-1",
                    _request(FakeCheckpointer([]), FakeEventStore([])),
                    head_checkpoint=head,
                )
            )

    assert exc.value.status_code == 409
    assert exc.value.detail == "Could not safely resolve the checkpoint before the target user message"
    accessor.ahistory.assert_not_awaited()


def test_prepare_regenerate_rejects_dangling_parent_without_chronological_fallback():
    from app.gateway.routers import thread_runs

    human = HumanMessage(id="human-1", content="question")
    ai = AIMessage(id="ai-1", content="answer")
    head = SimpleNamespace(
        values={"messages": [human, ai]},
        config={
            "configurable": {
                "thread_id": "thread-1",
                "checkpoint_ns": "",
                "checkpoint_id": "head",
            }
        },
        metadata={},
        parent_config={
            "configurable": {
                "thread_id": "thread-1",
                "checkpoint_ns": "",
                "checkpoint_id": "missing",
            }
        },
    )
    missing = SimpleNamespace(
        values={},
        config=head.parent_config,
        metadata=None,
        created_at=None,
        parent_config=None,
    )
    accessor = SimpleNamespace(
        aget=AsyncMock(return_value=missing),
        ahistory=AsyncMock(return_value=[head, _snapshot("wrong-sibling", [])]),
    )
    builder = AsyncMock(
        return_value=(
            accessor,
            {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}},
        )
    )

    with patch.object(thread_runs, "build_thread_checkpoint_state_accessor", builder):
        with pytest.raises(HTTPException) as exc:
            asyncio.run(
                thread_runs._find_base_checkpoint_before_human(
                    "thread-1",
                    "human-1",
                    _request(FakeCheckpointer([]), FakeEventStore([])),
                    head_checkpoint=head,
                )
            )

    assert exc.value.status_code == 409
    assert exc.value.detail == "Could not safely resolve the checkpoint before the target user message"
    accessor.ahistory.assert_not_awaited()


def test_prepare_regenerate_payload_rejects_non_latest_assistant():
    from app.gateway.routers.thread_runs import _prepare_regenerate_payload

    human = HumanMessage(id="human-1", content="question")
    old_ai = AIMessage(id="ai-old", content="old")
    latest_ai = AIMessage(id="ai-latest", content="latest")
    base = _checkpoint("ckpt-base", [])
    after_human = _checkpoint("ckpt-human", [human])
    latest = _checkpoint("ckpt-latest", [human, old_ai, latest_ai])
    checkpointer = FakeCheckpointer([latest, after_human, base])
    event_store = FakeEventStore(
        [
            {
                "run_id": "run-old",
                "event_type": "ai_message",
                "category": "message",
                "content": {"id": "ai-old", "type": "ai", "content": "old"},
                "metadata": {"caller": "lead_agent"},
            }
        ]
    )

    with pytest.raises(HTTPException) as exc:
        asyncio.run(_prepare_regenerate_payload("thread-1", "ai-old", _request(checkpointer, event_store)))

    assert exc.value.status_code == 409
    assert exc.value.detail == "Only the latest assistant message can be regenerated"


@pytest.mark.parametrize(
    ("content", "summary"),
    [
        ("answer", "answer"),
        ("<think>Reasoning.</think>answer", "<think>Reasoning.</think>answer"),
        ("<think>Reasoning.</think>answer", "answer"),
        ("<think>" + "r" * 2400 + "</think>" + "a" * 2500, "a" * 2000),
    ],
    ids=["plain", "legacy-summary", "visible-summary", "long-reasoning"],
)
def test_prepare_regenerate_payload_falls_back_to_matching_run_when_events_are_missing(content, summary):
    from app.gateway.routers.thread_runs import _prepare_regenerate_payload

    human = HumanMessage(id="human-1", content="question")
    ai = AIMessage(id="ai-1", content=content)
    base = _checkpoint("ckpt-base", [])
    after_human = _checkpoint("ckpt-human", [human])
    latest = _checkpoint("ckpt-ai", [human, ai])
    checkpointer = FakeCheckpointer([latest, after_human, base])
    run_manager = FakeRunManager(
        [
            RunRecord(run_id="run-latest", thread_id="thread-1", assistant_id="lead_agent", status=RunStatus.success, on_disconnect=DisconnectMode.continue_, last_ai_message=summary),
            SimpleNamespace(run_id="run-older", status=RunStatus.error, last_ai_message="answer"),
        ]
    )

    response = asyncio.run(
        _prepare_regenerate_payload(
            "thread-1",
            "ai-1",
            _request(checkpointer, FakeEventStore([]), run_manager=run_manager),
        )
    )

    assert response.target_run_id == "run-latest"
    assert response.metadata["regenerate_from_run_id"] == "run-latest"


def test_prepare_regenerate_payload_uses_server_stamped_human_run_id_without_parent_events():
    from app.gateway.routers.thread_runs import _prepare_regenerate_payload

    human = HumanMessage(id="human-1", content="question", additional_kwargs={"run_id": "parent-run"})
    ai = AIMessage(id="ai-1", content="answer")
    base = _checkpoint("ckpt-base", [])
    after_human = _checkpoint("ckpt-human", [human])
    latest = _checkpoint(
        "ckpt-ai",
        [human, ai],
        metadata={
            "deerflow_branch": True,
            "branch_parent_thread_id": "parent-thread",
            "branch_parent_checkpoint_id": "parent-checkpoint",
        },
    )
    checkpointer = FakeCheckpointer([latest, after_human, base])
    event_store = FakeEventStore([])

    response = asyncio.run(_prepare_regenerate_payload("thread-1", "ai-1", _request(checkpointer, event_store)))

    assert response.target_run_id == "parent-run"


def test_prepare_regenerate_payload_rejects_unverified_run_fallback_when_events_are_missing():
    from app.gateway.routers.thread_runs import _prepare_regenerate_payload

    human = HumanMessage(id="human-1", content="question")
    ai = AIMessage(id="ai-1", content="answer")
    base = _checkpoint("ckpt-base", [])
    after_human = _checkpoint("ckpt-human", [human])
    latest = _checkpoint("ckpt-ai", [human, ai])
    checkpointer = FakeCheckpointer([latest, after_human, base])
    run_manager = FakeRunManager(
        [
            SimpleNamespace(run_id="run-latest", status=RunStatus.success, last_ai_message="different"),
        ]
    )

    with pytest.raises(HTTPException) as exc:
        asyncio.run(
            _prepare_regenerate_payload(
                "thread-1",
                "ai-1",
                _request(checkpointer, FakeEventStore([]), run_manager=run_manager),
            )
        )

    assert exc.value.status_code == 409
    assert exc.value.detail == "Could not find source run for assistant message"


def test_prepare_regenerate_payload_requires_addressable_checkpoint_before_human():
    from app.gateway.routers.thread_runs import _prepare_regenerate_payload

    human = HumanMessage(id="human-1", content="question")
    ai = AIMessage(id="ai-1", content="answer")
    latest = _checkpoint("ckpt-ai", [human, ai])
    checkpointer = FakeCheckpointer([latest])
    event_store = FakeEventStore(
        [
            {
                "run_id": "run-old",
                "event_type": "llm.ai.response",
                "category": "message",
                "content": {"id": "ai-1", "type": "ai", "content": "answer"},
                "metadata": {"caller": "lead_agent"},
            }
        ]
    )

    with pytest.raises(HTTPException) as exc:
        asyncio.run(_prepare_regenerate_payload("thread-1", "ai-1", _request(checkpointer, event_store)))

    assert exc.value.status_code == 409
    assert exc.value.detail == "Could not find an addressable checkpoint before the target user message"
    assert checkpointer.alist_limits == [400]


def test_prepare_regenerate_payload_reports_recent_checkpoint_scan_limit():
    from app.gateway.routers.thread_runs import _prepare_regenerate_payload

    human = HumanMessage(id="human-1", content="question")
    ai = AIMessage(id="ai-1", content="answer")
    latest = _checkpoint("ckpt-latest", [human, ai])
    history_without_human = [_checkpoint(f"ckpt-{index}", []) for index in range(201)]
    checkpointer = FakeCheckpointer(history_without_human, latest=latest)
    event_store = FakeEventStore(
        [
            {
                "run_id": "run-old",
                "event_type": "llm.ai.response",
                "category": "message",
                "content": {"id": "ai-1", "type": "ai", "content": "answer"},
                "metadata": {"caller": "lead_agent"},
            }
        ]
    )

    with pytest.raises(HTTPException) as exc:
        asyncio.run(_prepare_regenerate_payload("thread-1", "ai-1", _request(checkpointer, event_store)))

    assert exc.value.status_code == 409
    assert exc.value.detail == "Could not locate target user message in recent checkpoint history (limit=200)"
    assert checkpointer.alist_limits == [400]


def test_find_base_checkpoint_ignores_duration_only_checkpoints() -> None:
    from app.gateway.routers.thread_runs import _find_base_checkpoint_before_human

    human = HumanMessage(id="human-1", content="question")
    duration_checkpoints = [
        _checkpoint(
            f"duration-{index}",
            [],
            metadata={"writes": {"runtime_run_duration": {"run_ids": [f"run-{index}"]}}},
        )
        for index in range(200)
    ]
    base = _checkpoint("ckpt-base", [])
    after_human = _checkpoint("ckpt-human", [human])
    checkpointer = FakeCheckpointer([*duration_checkpoints, after_human, base])

    result = asyncio.run(_find_base_checkpoint_before_human("thread-1", "human-1", _request(checkpointer, FakeEventStore([]))))

    assert result.config == base.config
    assert checkpointer.alist_limits == [400]