deer-flow/backend/tests/test_run_worker_rollback.py

import asyncio
import copy
import logging
import threading
import weakref
from contextlib import suppress
from contextvars import ContextVar
from types import SimpleNamespace
from typing import Annotated, Any, NotRequired, TypedDict
from unittest.mock import AsyncMock, MagicMock, call, patch

import pytest
from langchain_core.messages import AIMessage, AIMessageChunk, AnyMessage, HumanMessage
from langgraph.channels.delta import DeltaChannel
from langgraph.checkpoint.base import empty_checkpoint
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages
from langgraph.types import Overwrite

from deerflow.agents.thread_state import merge_artifacts, merge_message_writes
from deerflow.config.run_ownership_config import RunOwnershipConfig
from deerflow.runtime import context_compaction
from deerflow.runtime.checkpoint_state import CheckpointStateAccessor, build_state_mutation_graph
from deerflow.runtime.context_compaction import compact_thread_context
from deerflow.runtime.context_keys import CHECKPOINT_AGENT_NAME_METADATA_KEY, CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY
from deerflow.runtime.events.store.memory import MemoryRunEventStore
from deerflow.runtime.journal import RunJournal
from deerflow.runtime.runs.manager import CancelOutcome, ConflictError, RunManager
from deerflow.runtime.runs.schemas import RunStatus
from deerflow.runtime.runs.store.memory import MemoryRunStore
from deerflow.runtime.runs.worker import (
    RollbackPoint,
    RunContext,
    _agent_factory_supports_app_config,
    _build_runtime_context,
    _bump_channel_version,
    _capture_rollback_point,
    _collect_pre_existing_message_ids,
    _ensure_interrupted_title,
    _extract_llm_error_fallback_message,
    _install_runtime_context,
    _LargeFileToolChunkBatcher,
    _rollback_to_pre_run_checkpoint,
    _try_extract_from_message,
    run_agent,
)
from deerflow.sandbox.lease import (
    SANDBOX_COMMAND_SCOPE_CONTEXT_KEY,
    SANDBOX_LEASE_OWNER_CONTEXT_KEY,
    ensure_sandbox_lease_owner,
    get_sandbox_lease_manager,
)
from deerflow.sandbox.sandbox_provider import reset_sandbox_provider, set_sandbox_provider


class FakeCheckpointer:
    def __init__(self):
        self.adelete_thread = AsyncMock()
        self.aget_tuple = AsyncMock(return_value=None)
        self.aput_writes = AsyncMock()


def _lease_test_bridge():
    return SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )


@pytest.mark.anyio
async def test_run_agent_releases_execution_lease_when_graph_raises():
    provider = MagicMock()
    provider.get.return_value = MagicMock()
    manager = get_sandbox_lease_manager(provider)
    run_manager = RunManager()
    record = await run_manager.create("thread-lead-error")
    owner_ids: list[str] = []

    class FailingAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, stream_mode, subgraphs
            context = config["configurable"]["__pregel_runtime"].context
            owner_id = ensure_sandbox_lease_owner(context)
            assert owner_id is not None
            owner_ids.append(owner_id)
            context["sandbox_id"] = "shared"
            manager.retain(
                owner_id,
                "shared",
                thread_id=record.thread_id,
                user_id="anonymous",
            )
            raise RuntimeError("lead model failed")
            yield  # pragma: no cover

    set_sandbox_provider(provider)
    try:
        await run_agent(
            _lease_test_bridge(),
            run_manager,
            record,
            ctx=RunContext(checkpointer=None),
            agent_factory=lambda **_kwargs: FailingAgent(),
            graph_input={},
            config={},
        )
        await asyncio.sleep(0)

        assert len(owner_ids) == 1
        assert manager.binding_for(owner_ids[0]) is None
        provider.get.return_value.release_command_scope.assert_called_once_with(owner_ids[0])
        provider.release.assert_called_once_with("shared")
    finally:
        reset_sandbox_provider()


@pytest.mark.anyio
async def test_run_agent_releases_execution_lease_when_cancelled():
    provider = MagicMock()
    provider.get.return_value = MagicMock()
    manager = get_sandbox_lease_manager(provider)
    run_manager = RunManager()
    record = await run_manager.create("thread-lead-cancel")
    lease_bound = asyncio.Event()
    owner_ids: list[str] = []

    class BlockingAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, stream_mode, subgraphs
            context = config["configurable"]["__pregel_runtime"].context
            owner_id = ensure_sandbox_lease_owner(context)
            assert owner_id is not None
            owner_ids.append(owner_id)
            context["sandbox_id"] = "shared"
            manager.retain(
                owner_id,
                "shared",
                thread_id=record.thread_id,
                user_id="anonymous",
            )
            lease_bound.set()
            await asyncio.Event().wait()
            yield  # pragma: no cover

    set_sandbox_provider(provider)
    try:
        task = asyncio.create_task(
            run_agent(
                _lease_test_bridge(),
                run_manager,
                record,
                ctx=RunContext(checkpointer=None),
                agent_factory=lambda **_kwargs: BlockingAgent(),
                graph_input={},
                config={},
            )
        )
        await asyncio.wait_for(lease_bound.wait(), timeout=1)
        task.cancel()
        await task
        await asyncio.sleep(0)

        assert len(owner_ids) == 1
        assert manager.binding_for(owner_ids[0]) is None
        provider.get.return_value.release_command_scope.assert_called_once_with(owner_ids[0])
        provider.release.assert_called_once_with("shared")
    finally:
        reset_sandbox_provider()


@pytest.mark.anyio
async def test_run_agent_cleans_up_when_mcp_task_projection_is_cancelled():
    class CleanupTrackingRunManager(RunManager):
        def __init__(self) -> None:
            super().__init__()
            self.cleanup_calls: list[tuple[str, float]] = []

        async def cleanup(self, run_id: str, *, delay: float = 300) -> None:
            self.cleanup_calls.append((run_id, delay))

    projection_started = asyncio.Event()

    class BlockingTaskRepository:
        async def list_by_thread(self, thread_id, *, user_id, limit):
            del thread_id, user_id, limit
            projection_started.set()
            await asyncio.Event().wait()

    run_manager = CleanupTrackingRunManager()
    record = await run_manager.create("thread-mcp-projection-cancelled", user_id="alice")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    agent_factory = MagicMock(side_effect=AssertionError("cancelled preflight built the agent"))
    run_task = asyncio.create_task(
        run_agent(
            bridge,
            run_manager,
            record,
            ctx=RunContext(
                checkpointer=None,
                event_store=MemoryRunEventStore(),
                mcp_task_repo=BlockingTaskRepository(),
            ),
            agent_factory=agent_factory,
            graph_input={},
            config={},
        )
    )
    await asyncio.wait_for(projection_started.wait(), timeout=1)

    run_task.cancel("MCP projection interrupted")
    await run_task
    await asyncio.sleep(0)

    agent_factory.assert_not_called()
    assert record.status == RunStatus.interrupted
    assert record.finalizing is False
    bridge.publish_end.assert_awaited_once_with(record.run_id)
    bridge.cleanup.assert_awaited_once_with(record.run_id, delay=60)
    assert run_manager.cleanup_calls == [(record.run_id, 300)]


@pytest.mark.anyio
async def test_pending_cancel_stops_waiting_for_prior_finalization():
    run_manager = RunManager()
    prior = await run_manager.create("thread-cancel-while-waiting")
    prior.status = RunStatus.interrupted
    prior.finalizing = True
    record = await run_manager.create("thread-cancel-while-waiting")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    factory_called = False

    def agent_factory(**_kwargs):
        nonlocal factory_called
        factory_called = True
        raise AssertionError("cancelled pending run must not build an agent")

    record.task = asyncio.create_task(
        run_agent(
            bridge,
            run_manager,
            record,
            ctx=RunContext(checkpointer=None),
            agent_factory=agent_factory,
            graph_input={"messages": []},
            config={},
        )
    )
    await asyncio.sleep(0)

    outcome = await run_manager.cancel(record.run_id)
    await asyncio.wait_for(record.task, timeout=0.2)

    assert outcome == CancelOutcome.cancelled
    assert prior.finalizing is True
    assert factory_called is False
    assert record.status == RunStatus.interrupted
    bridge.publish_end.assert_awaited_once_with(record.run_id)


@pytest.mark.anyio
async def test_remote_cancel_wins_when_graph_finishes_before_owner_heartbeat():
    store = MemoryRunStore()
    ownership = RunOwnershipConfig(
        heartbeat_enabled=True,
        lease_seconds=30,
        grace_seconds=10,
    )
    owner = RunManager(
        store=store,
        worker_id="worker-a",
        run_ownership_config=ownership,
    )
    peer = RunManager(
        store=store,
        worker_id="worker-b",
        run_ownership_config=ownership,
    )
    record = await owner.create_or_reject("thread-cancel-race")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    started = asyncio.Event()
    release = asyncio.Event()

    class _FinishingAgent:
        async def astream(
            self,
            graph_input,
            config=None,
            stream_mode=None,
            subgraphs=False,
        ):
            del graph_input, config, stream_mode, subgraphs
            started.set()
            await release.wait()
            yield {"messages": []}

    task = asyncio.create_task(
        run_agent(
            bridge,
            owner,
            record,
            ctx=RunContext(
                checkpointer=None,
                event_store=MemoryRunEventStore(),
            ),
            agent_factory=lambda **_kwargs: _FinishingAgent(),
            graph_input={},
            config={},
        )
    )
    record.task = task

    await asyncio.wait_for(started.wait(), timeout=1)
    assert await peer.cancel(record.run_id, action="rollback") == CancelOutcome.requested
    release.set()
    await asyncio.wait_for(task, timeout=1)

    stored = await store.get(record.run_id)
    assert stored is not None
    assert stored["status"] == "error"
    assert stored["error"] == "Rolled back by user"
    assert record.status == RunStatus.error


def _make_rollback_point(*, checkpoint_id="ckpt-1", messages=("before",), pending_writes=()):
    materialized_messages = tuple(messages)
    return RollbackPoint(
        config={
            "configurable": {
                "thread_id": "thread-1",
                "checkpoint_ns": "",
                "checkpoint_id": checkpoint_id,
            }
        },
        state_values={},
        messages=materialized_messages,
        metadata={"source": "input"},
        pending_writes=tuple(pending_writes),
    )


def _stub_mutation_graph(monkeypatch, *, restored_config):
    """Replace the rollback mutation graph with a stub returning ``restored_config``."""
    mock_graph = SimpleNamespace()
    mock_graph.aupdate_state = AsyncMock(return_value=restored_config)
    monkeypatch.setattr(
        "deerflow.runtime.runs.worker.build_state_mutation_graph",
        lambda *args, **kwargs: mock_graph,
    )
    return mock_graph


def _rollback_accessor(checkpointer):
    return CheckpointStateAccessor(graph=SimpleNamespace(), checkpointer=checkpointer, mode="full")


class _DeltaChannelState(TypedDict):
    messages: Annotated[list[AnyMessage], DeltaChannel(merge_message_writes, snapshot_frequency=1000)]


class _FullChannelState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]


class _CompactionFullState(_FullChannelState):
    summary_text: NotRequired[str | None]


class _CompactionDeltaState(_DeltaChannelState):
    summary_text: NotRequired[str | None]


def _build_message_append_graph(state_schema: type, checkpointer: Any):
    async def _append_message(state: dict[str, Any]) -> dict[str, Any]:
        return {"messages": [HumanMessage(content=f"turn-{len(state.get('messages') or [])}")]}

    builder = StateGraph(state_schema)
    builder.add_node("append_message", _append_message)
    builder.set_entry_point("append_message")
    builder.set_finish_point("append_message")
    return builder.compile(checkpointer=checkpointer)


@pytest.mark.parametrize("tool_name", ["write_file", "str_replace"])
def test_large_file_tool_chunk_batcher_streams_bounded_batches(tool_name: str):
    batcher = _LargeFileToolChunkBatcher(batch_size=2)
    first = AIMessageChunk(
        content="",
        id="ai-1",
        tool_call_chunks=[
            {
                "id": "call-1",
                "index": 0,
                "name": tool_name,
                "args": '{"path":"/mnt/user-data/outputs/report.md","content":"Hel',
            }
        ],
    )
    continuation = AIMessageChunk(
        content="",
        id="ai-1",
        tool_call_chunks=[{"index": 0, "name": None, "args": 'lo"}'}],
    )

    assert batcher.push((first, {})) == []
    published = batcher.push((continuation, {}))
    assert len(published) == 1
    message, metadata = published[0]
    assert metadata == {}
    assert message.tool_calls[0]["args"]["content"] == "Hello"
    assert batcher.flush() == []


def test_large_file_tool_chunk_batcher_preserves_visible_and_non_file_chunks():
    batcher = _LargeFileToolChunkBatcher()
    visible_text = AIMessageChunk(content="Writing the report now.", id="ai-1")
    search_tool = AIMessageChunk(
        content="",
        id="ai-2",
        tool_call_chunks=[
            {
                "id": "call-2",
                "index": 0,
                "name": "web_search",
                "args": '{"query":"vector databases"}',
            }
        ],
    )
    write_with_reasoning = AIMessageChunk(
        content="",
        id="ai-3",
        additional_kwargs={"reasoning_content": "Choosing a filename."},
        tool_call_chunks=[
            {
                "id": "call-3",
                "index": 0,
                "name": "write_file",
                "args": '{"path":"/mnt/user-data/outputs/report.md"}',
            }
        ],
    )

    assert batcher.push((visible_text, {})) == [(visible_text, {})]
    assert batcher.push((search_tool, {})) == [(search_tool, {})]
    visible_reasoning = batcher.push((write_with_reasoning, {}))
    assert len(visible_reasoning) == 1
    filtered_message, filtered_metadata = visible_reasoning[0]
    assert filtered_metadata == {}
    assert filtered_message.additional_kwargs == {"reasoning_content": "Choosing a filename."}
    assert filtered_message.tool_call_chunks == []
    pending_file_chunks = batcher.flush()
    assert len(pending_file_chunks) == 1
    assert pending_file_chunks[0][0].tool_call_chunks[0]["name"] == "write_file"


def test_large_file_tool_chunk_batcher_separates_subgraph_namespaces():
    batcher = _LargeFileToolChunkBatcher()
    first = AIMessageChunk(
        content="",
        id="shared-ai-id",
        tool_call_chunks=[{"id": "call-a", "index": 0, "name": "write_file", "args": '{"path":"a.md","content":"A'}],
    )
    second = AIMessageChunk(
        content="",
        id="shared-ai-id",
        tool_call_chunks=[{"id": "call-b", "index": 0, "name": "write_file", "args": '{"path":"b.md","content":"B'}],
    )

    assert batcher.push((first, {"langgraph_checkpoint_ns": "task-a"})) == []
    published = batcher.push((second, {"langgraph_checkpoint_ns": "task-b"}))

    assert len(published) == 1
    assert published[0][1]["langgraph_checkpoint_ns"] == "task-a"
    assert batcher.flush()[0][1]["langgraph_checkpoint_ns"] == "task-b"


@pytest.mark.parametrize("metadata", [None, "not-a-dict"])
def test_large_file_tool_chunk_batcher_accepts_non_dict_metadata(metadata: Any):
    batcher = _LargeFileToolChunkBatcher(batch_size=1)
    message = AIMessageChunk(
        content="",
        id="ai-file",
        tool_call_chunks=[
            {
                "id": "call-file",
                "index": 0,
                "name": "write_file",
                "args": '{"path":"report.md","content":"draft"}',
            }
        ],
    )

    published = batcher.push((message, metadata))

    assert len(published) == 1
    assert published[0][0].tool_call_chunks[0]["args"] == '{"path":"report.md","content":"draft"}'
    assert published[0][1] == {}


def test_large_file_tool_chunk_batcher_does_not_retain_non_file_names():
    batcher = _LargeFileToolChunkBatcher()

    for index in range(100):
        message = AIMessageChunk(
            content="",
            id=f"ai-{index}",
            tool_call_chunks=[
                {
                    "id": f"call-{index}",
                    "index": 0,
                    "name": "web_search",
                    "args": '{"query":"deerflow"}',
                }
            ],
        )

        assert batcher.push((message, {})) == [(message, {})]

    assert batcher.tool_names == {}


def test_large_file_tool_chunk_batcher_starts_batching_after_split_name_matches():
    batcher = _LargeFileToolChunkBatcher(batch_size=1)
    name_prefix = AIMessageChunk(
        content="",
        id="ai-file",
        tool_call_chunks=[{"id": "call-file", "index": 0, "name": "write_", "args": ""}],
    )
    name_suffix = AIMessageChunk(
        content="",
        id="ai-file",
        tool_call_chunks=[
            {
                "index": 0,
                "name": "file",
                "args": '{"path":"report.md","content":"draft"}',
            }
        ],
    )

    assert batcher.push((name_prefix, {})) == [(name_prefix, {})]
    assert set(batcher.tool_names.values()) == {"write_"}
    assert len(batcher.push((name_suffix, {}))) == 1
    assert set(batcher.tool_names.values()) == {"write_file"}


def test_large_file_tool_chunk_batcher_keeps_identity_across_batches_then_releases_it():
    batcher = _LargeFileToolChunkBatcher(batch_size=1)
    first = AIMessageChunk(
        content="",
        id="ai-file",
        tool_call_chunks=[
            {
                "id": "call-file",
                "index": 0,
                "name": "write_file",
                "args": '{"path":"report.md","content":"Hel',
            }
        ],
    )
    continuation = AIMessageChunk(
        content="",
        id="ai-file",
        tool_call_chunks=[{"index": 0, "name": None, "args": 'lo"}'}],
    )

    assert len(batcher.push((first, {}))) == 1
    assert len(batcher.push((continuation, {}))) == 1
    assert set(batcher.tool_names.values()) == {"write_file"}

    assert batcher.finish() == []
    assert batcher.tool_names == {}


def test_build_runtime_context_includes_app_config_when_present():
    app_config = object()

    context = _build_runtime_context("thread-1", "run-1", None, app_config)

    assert context["thread_id"] == "thread-1"
    assert context["run_id"] == "run-1"
    assert context["app_config"] is app_config


def test_install_runtime_context_preserves_existing_thread_id_and_threads_app_config():
    app_config = object()
    config = {"context": {"thread_id": "caller-thread"}}

    _install_runtime_context(
        config,
        {
            "thread_id": "record-thread",
            "run_id": "run-1",
            "app_config": app_config,
        },
    )

    assert config["context"]["thread_id"] == "caller-thread"
    assert config["context"]["run_id"] == "run-1"
    assert config["context"]["app_config"] is app_config


def test_install_runtime_context_overrides_internal_pre_existing_message_ids():
    config = {"context": {CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY: {"spoofed"}}}

    _install_runtime_context(
        config,
        {
            "thread_id": "record-thread",
            "run_id": "run-1",
            CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY: frozenset({"old-ai"}),
        },
    )

    assert config["context"][CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] == frozenset({"old-ai"})


def test_install_runtime_context_removes_caller_sandbox_execution_identities():
    config = {
        "context": {
            SANDBOX_LEASE_OWNER_CONTEXT_KEY: "forged-owner",
            SANDBOX_COMMAND_SCOPE_CONTEXT_KEY: "forged-scope",
        }
    }

    _install_runtime_context(
        config,
        {
            "thread_id": "record-thread",
            "run_id": "run-1",
        },
    )

    assert SANDBOX_LEASE_OWNER_CONTEXT_KEY not in config["context"]
    assert SANDBOX_COMMAND_SCOPE_CONTEXT_KEY not in config["context"]


def test_build_runtime_context_ignores_caller_project_context_key():
    """The worker merge refuses a caller-supplied pinned project snapshot (§12)."""
    from deerflow.runtime.context_keys import PROJECT_CONTEXT_KEY

    ctx = _build_runtime_context("thread-1", "run-1", {PROJECT_CONTEXT_KEY: {"project_id": "forged"}, "agent_name": "kept"})

    assert PROJECT_CONTEXT_KEY not in ctx
    assert ctx["agent_name"] == "kept"


def test_install_runtime_context_removes_caller_project_context_key_when_unpinned():
    from deerflow.runtime.context_keys import PROJECT_CONTEXT_KEY

    config = {"context": {PROJECT_CONTEXT_KEY: {"project_id": "forged"}}}

    _install_runtime_context(config, {"thread_id": "record-thread", "run_id": "run-1"})

    assert PROJECT_CONTEXT_KEY not in config["context"]


def test_install_runtime_context_preserves_the_pinned_project_context():
    from deerflow.runtime.context_keys import PROJECT_CONTEXT_KEY

    stamped = {"project_id": "p-1", "name": "Roadmap", "instructions": "x"}
    config = {"context": {PROJECT_CONTEXT_KEY: stamped}}

    _install_runtime_context(config, {"thread_id": "record-thread", "run_id": "run-1", PROJECT_CONTEXT_KEY: stamped})

    assert config["context"][PROJECT_CONTEXT_KEY] is stamped


def test_pin_admission_project_context_hoists_only_the_stamped_value():
    from deerflow.runtime.context_keys import PROJECT_CONTEXT_KEY
    from deerflow.runtime.runs.worker import _pin_admission_project_context

    stamped = {"project_id": "p-1", "name": "Roadmap", "instructions": "x"}
    runtime_ctx = {"thread_id": "record-thread", "run_id": "run-1"}
    _pin_admission_project_context({"context": {PROJECT_CONTEXT_KEY: stamped}}, runtime_ctx)
    assert runtime_ctx[PROJECT_CONTEXT_KEY] is stamped

    empty = {"thread_id": "record-thread", "run_id": "run-1"}
    _pin_admission_project_context({"context": {}}, empty)
    _pin_admission_project_context({}, empty)
    _pin_admission_project_context({"context": "not-a-mapping"}, empty)
    assert PROJECT_CONTEXT_KEY not in empty


@pytest.mark.parametrize(
    ("stream_modes", "emit_values"),
    [
        (["messages-tuple", "values"], True),
        (["messages-tuple", "updates", "custom"], False),
    ],
)
@pytest.mark.anyio
async def test_run_agent_batches_incremental_file_args_and_keeps_complete_values(stream_modes: list[str], emit_values: bool):
    run_manager = RunManager()
    record = await run_manager.create("thread-file-stream")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    complete_message = AIMessage(
        content="",
        id="ai-file",
        tool_calls=[
            {
                "id": "call-file",
                "name": "write_file",
                "args": {
                    "path": "/mnt/user-data/outputs/report.md",
                    "content": "Hello world",
                },
                "type": "tool_call",
            }
        ],
    )

    class DummyAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, config, stream_mode, subgraphs
            yield (
                "messages",
                (
                    AIMessageChunk(
                        content="",
                        id="ai-file",
                        tool_call_chunks=[
                            {
                                "id": "call-file",
                                "index": 0,
                                "name": "write_file",
                                "args": '{"path":"/mnt/user-data/outputs/report.md","content":"Hello',
                            }
                        ],
                    ),
                    {},
                ),
            )
            yield (
                "messages",
                (
                    AIMessageChunk(
                        content="",
                        id="ai-file",
                        tool_call_chunks=[{"index": 0, "name": None, "args": ' world"}'}],
                    ),
                    {},
                ),
            )
            if emit_values:
                yield ("values", {"messages": [complete_message]})
            else:
                yield ("updates", {"agent": {}})

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(checkpointer=None),
        agent_factory=lambda **_kwargs: DummyAgent(),
        graph_input={},
        config={},
        stream_modes=stream_modes,
    )

    message_events = [call.args for call in bridge.publish.await_args_list if call.args[1] == "messages"]
    assert len(message_events) == 1
    assert message_events[0][2][0]["tool_calls"][0]["args"]["content"] == "Hello world"
    if emit_values:
        values_events = [call.args[2] for call in bridge.publish.await_args_list if call.args[1] == "values"]
        assert any(event["messages"][0]["tool_calls"][0]["args"]["content"] == "Hello world" for event in values_events)


@pytest.mark.parametrize(
    ("stream_error", "flush_publish_error", "expected_error"),
    [
        (True, False, "stream failed"),
        (True, True, "stream failed"),
        (False, True, "flush publish failed"),
    ],
)
@pytest.mark.anyio
async def test_run_agent_handles_pending_file_args_when_stream_or_flush_raises(stream_error: bool, flush_publish_error: bool, expected_error: str):
    run_manager = RunManager()
    record = await run_manager.create("thread-file-stream-error")

    async def publish(_run_id: str, event: str, _data: Any):
        if flush_publish_error and event == "messages":
            raise RuntimeError("flush publish failed")

    bridge = SimpleNamespace(
        publish=AsyncMock(side_effect=publish),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )

    class DummyAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, config, stream_mode, subgraphs
            yield (
                "messages",
                (
                    AIMessageChunk(
                        content="",
                        id="ai-file",
                        tool_call_chunks=[
                            {
                                "id": "call-file",
                                "index": 0,
                                "name": "write_file",
                                "args": '{"path":"report.md","content":"partial',
                            }
                        ],
                    ),
                    {},
                ),
            )
            if stream_error:
                raise RuntimeError("stream failed")

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(checkpointer=None),
        agent_factory=lambda **_kwargs: DummyAgent(),
        graph_input={},
        config={},
        stream_modes=["messages-tuple", "values"],
    )

    message_events = [call.args for call in bridge.publish.await_args_list if call.args[1] == "messages"]
    assert len(message_events) == 1
    assert message_events[0][2][0]["tool_call_chunks"][0]["args"].endswith('"content":"partial')
    error_events = [call.args for call in bridge.publish.await_args_list if call.args[1] == "error"]
    assert error_events[0][2]["message"] == expected_error


@pytest.mark.anyio
async def test_run_agent_keeps_file_chunks_unbatched_without_values_mode():
    run_manager = RunManager()
    record = await run_manager.create("thread-file-messages-only")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    chunks = [
        AIMessageChunk(
            content="",
            id="ai-file",
            tool_call_chunks=[
                {
                    "id": "call-file",
                    "index": 0,
                    "name": "write_file",
                    "args": '{"path":"report.md","content":"Hel',
                }
            ],
        ),
        AIMessageChunk(
            content="",
            id="ai-file",
            tool_call_chunks=[{"index": 0, "name": None, "args": 'lo"}'}],
        ),
    ]

    class DummyAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, config, stream_mode, subgraphs
            for chunk in chunks:
                yield (chunk, {})

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(checkpointer=None),
        agent_factory=lambda **_kwargs: DummyAgent(),
        graph_input={},
        config={},
        stream_modes=["messages-tuple"],
    )

    message_events = [call.args[2] for call in bridge.publish.await_args_list if call.args[1] == "messages"]
    assert len(message_events) == 2
    assert message_events[0][0]["tool_call_chunks"][0]["args"].endswith('"content":"Hel')
    assert message_events[1][0]["tool_call_chunks"][0]["args"] == 'lo"}'


@pytest.mark.anyio
async def test_run_agent_threads_explicit_app_config_into_config_only_factory():
    run_manager = RunManager()
    record = await run_manager.create("thread-1")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    app_config = object()
    captured: dict[str, object] = {}

    class DummyAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            captured["astream_context"] = config["context"]
            yield {"messages": []}

    def factory(*, config):
        captured["factory_context"] = config["context"]
        return DummyAgent()

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(checkpointer=None, app_config=app_config),
        agent_factory=factory,
        graph_input={},
        config={},
    )
    await asyncio.sleep(0)

    assert captured["factory_context"]["app_config"] is app_config
    assert captured["astream_context"]["app_config"] is app_config
    fetched = await run_manager.get(record.run_id)
    assert fetched is not None
    assert fetched.status == RunStatus.success
    bridge.publish_end.assert_awaited_once_with(record.run_id)
    bridge.cleanup.assert_awaited_once_with(record.run_id, delay=60)


@pytest.mark.anyio
async def test_run_agent_schedules_terminal_run_record_cleanup():
    class CleanupTrackingRunManager(RunManager):
        def __init__(self) -> None:
            super().__init__()
            self.cleanup_calls: list[tuple[str, float]] = []

        async def cleanup(self, run_id: str, *, delay: float = 300) -> None:
            self.cleanup_calls.append((run_id, delay))

    run_manager = CleanupTrackingRunManager()
    record = await run_manager.create("thread-terminal-cleanup")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )

    class DummyAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, config, stream_mode, subgraphs
            yield {"messages": []}

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(checkpointer=None),
        agent_factory=lambda **_kwargs: DummyAgent(),
        graph_input={},
        config={},
    )
    await asyncio.sleep(0)

    assert run_manager.cleanup_calls == [(record.run_id, 300)]


@pytest.mark.anyio
async def test_run_agent_schedules_terminal_cleanup_when_publish_end_fails(monkeypatch):
    import deerflow.runtime.runs.worker as worker_module

    class CleanupTrackingRunManager(RunManager):
        def __init__(self) -> None:
            super().__init__()
            self.cleanup_calls: list[tuple[str, float]] = []

        async def cleanup(self, run_id: str, *, delay: float = 300) -> None:
            self.cleanup_calls.append((run_id, delay))

    run_manager = CleanupTrackingRunManager()
    record = await run_manager.create("thread-terminal-publish-failure")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(side_effect=RuntimeError("end publication unavailable")),
        cleanup=AsyncMock(),
    )
    schedule_collection = MagicMock()
    monkeypatch.setattr(worker_module, "_schedule_terminal_cycle_collection", schedule_collection)

    class DummyAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, config, stream_mode, subgraphs
            yield {"messages": []}

    with pytest.raises(RuntimeError, match="end publication unavailable"):
        await run_agent(
            bridge,
            run_manager,
            record,
            ctx=RunContext(checkpointer=None),
            agent_factory=lambda **_kwargs: DummyAgent(),
            graph_input={},
            config={},
        )
    await asyncio.sleep(0)

    bridge.publish_end.assert_awaited_once_with(record.run_id)
    bridge.cleanup.assert_awaited_once_with(record.run_id, delay=60)
    assert run_manager.cleanup_calls == [(record.run_id, 300)]
    schedule_collection.assert_called_once_with()


@pytest.mark.anyio
async def test_run_agent_schedules_terminal_cleanup_when_completion_hook_is_cancelled(monkeypatch):
    import deerflow.runtime.runs.worker as worker_module
    from deerflow.runtime.journal import RunJournal

    class CleanupTrackingRunManager(RunManager):
        def __init__(self) -> None:
            super().__init__()
            self.cleanup_calls: list[tuple[str, float]] = []

        async def cleanup(self, run_id: str, *, delay: float = 300) -> None:
            self.cleanup_calls.append((run_id, delay))

    completion_hook_entered = asyncio.Event()

    async def block_completion(_record) -> None:
        completion_hook_entered.set()
        await asyncio.Event().wait()

    run_manager = CleanupTrackingRunManager()
    record = await run_manager.create("thread-terminal-completion-cancelled")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    schedule_collection = MagicMock()
    monkeypatch.setattr(worker_module, "_schedule_terminal_cycle_collection", schedule_collection)
    captured: dict[str, Any] = {}

    class DummyAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, stream_mode, subgraphs
            callbacks = config.get("callbacks") or []
            captured["journal"] = next(callback for callback in callbacks if isinstance(callback, RunJournal))
            yield {"messages": []}

    config: dict[str, Any] = {}
    run_task = asyncio.create_task(
        run_agent(
            bridge,
            run_manager,
            record,
            ctx=RunContext(
                checkpointer=None,
                event_store=MemoryRunEventStore(),
                on_run_completed=block_completion,
            ),
            agent_factory=lambda **_kwargs: DummyAgent(),
            graph_input={},
            config=config,
        )
    )
    await asyncio.wait_for(completion_hook_entered.wait(), timeout=1)
    run_task.cancel("completion hook interrupted")
    with pytest.raises(asyncio.CancelledError, match="completion hook interrupted"):
        await run_task
    await asyncio.sleep(0)

    journal = captured["journal"]
    assert "__pregel_runtime" not in config["configurable"]
    assert journal not in config["callbacks"]
    assert journal._closed is True
    assert journal._store is None
    assert record.finalizing is False
    bridge.publish_end.assert_awaited_once_with(record.run_id)
    bridge.cleanup.assert_awaited_once_with(record.run_id, delay=60)
    assert run_manager.cleanup_calls == [(record.run_id, 300)]
    schedule_collection.assert_called_once_with()


@pytest.mark.anyio
async def test_run_agent_closes_stream_when_abort_breaks_iteration():
    run_manager = RunManager()
    record = await run_manager.create("thread-stream-close")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )

    class CloseTrackingStream:
        def __init__(self) -> None:
            self.yielded = False
            self.closed = False

        def __aiter__(self):
            return self

        async def __anext__(self):
            if self.yielded:
                raise StopAsyncIteration
            self.yielded = True
            record.abort_event.set()
            return {"messages": []}

        async def aclose(self) -> None:
            self.closed = True

    stream = CloseTrackingStream()

    class DummyAgent:
        def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, config, stream_mode, subgraphs
            return stream

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(checkpointer=None),
        agent_factory=lambda **_kwargs: DummyAgent(),
        graph_input={},
        config={},
        stream_modes=["values"],
    )

    assert stream.closed is True
    assert record.status == RunStatus.interrupted


@pytest.mark.parametrize("stream_modes", [["values"], ["messages-tuple", "values"]])
@pytest.mark.parametrize("abort_before_break", [True, False], ids=["early-break", "exhaustion-race"])
@pytest.mark.anyio
async def test_run_agent_ignores_stream_close_failure_after_abort(stream_modes, abort_before_break, caplog):
    run_manager = RunManager()
    record = await run_manager.create("thread-stream-close-failure")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )

    class CloseFailingStream:
        def __init__(self) -> None:
            self.yielded = False
            self.close_attempted = False

        def __aiter__(self):
            return self

        async def __anext__(self):
            if self.yielded:
                if not abort_before_break:
                    record.abort_event.set()
                raise StopAsyncIteration
            self.yielded = True
            if abort_before_break:
                record.abort_event.set()
            chunk = {"messages": []}
            return chunk if len(stream_modes) == 1 else ("values", chunk)

        async def aclose(self) -> None:
            self.close_attempted = True
            raise RuntimeError("stream close failed")

    stream = CloseFailingStream()

    class DummyAgent:
        def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, config, stream_mode, subgraphs
            return stream

    with caplog.at_level(logging.WARNING, logger="deerflow.runtime.runs.worker"):
        await run_agent(
            bridge,
            run_manager,
            record,
            ctx=RunContext(checkpointer=None),
            agent_factory=lambda **_kwargs: DummyAgent(),
            graph_input={},
            config={},
            stream_modes=stream_modes,
        )

    assert stream.close_attempted is True
    assert record.status == RunStatus.interrupted
    assert record.error is None
    assert "Could not close aborted agent stream" in caplog.text
    bridge.publish_end.assert_awaited_once_with(record.run_id)


@pytest.mark.anyio
async def test_terminal_cleanup_tasks_do_not_inherit_run_context():
    marker: ContextVar[str | None] = ContextVar("run_cleanup_marker", default=None)
    seen: dict[str, str | None] = {}
    bridge_cleaned = asyncio.Event()
    manager_cleaned = asyncio.Event()

    class CleanupTrackingRunManager(RunManager):
        async def cleanup(self, run_id: str, *, delay: float = 300) -> None:
            seen["manager"] = marker.get()
            await super().cleanup(run_id, delay=0)
            manager_cleaned.set()

    class CleanupTrackingBridge:
        async def publish(self, *args, **kwargs) -> None:
            pass

        async def publish_end(self, run_id: str) -> None:
            pass

        async def cleanup(self, run_id: str, *, delay: float = 0) -> None:
            seen["bridge"] = marker.get()
            bridge_cleaned.set()

    class DummyAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, config, stream_mode, subgraphs
            yield {"messages": []}

    run_manager = CleanupTrackingRunManager()
    record = await run_manager.create("thread-contextless-cleanup")
    token = marker.set("run-context")
    try:
        await run_agent(
            CleanupTrackingBridge(),
            run_manager,
            record,
            ctx=RunContext(checkpointer=None),
            agent_factory=lambda **_kwargs: DummyAgent(),
            graph_input={},
            config={},
        )
        await asyncio.wait_for(
            asyncio.gather(bridge_cleaned.wait(), manager_cleaned.wait()),
            timeout=1,
        )
    finally:
        marker.reset(token)

    assert seen == {"bridge": None, "manager": None}


@pytest.mark.anyio
async def test_terminal_cycle_collection_is_coalesced_contextless_and_off_loop(monkeypatch, caplog):
    import deerflow.runtime.runs.worker as worker_module

    marker: ContextVar[str | None] = ContextVar("terminal_gc_marker", default=None)
    loop_thread_id = threading.get_ident()
    seen: list[tuple[str | None, int]] = []
    loop = asyncio.get_running_loop()
    collected = asyncio.Event()

    def collect() -> int:
        seen.append((marker.get(), threading.get_ident()))
        loop.call_soon_threadsafe(collected.set)
        return 7

    monkeypatch.setattr(worker_module, "_TERMINAL_CYCLE_COLLECTION_INTERVAL_SECONDS", 0.0)
    monkeypatch.setattr(worker_module, "_TERMINAL_CYCLE_COLLECTION_INFO_THRESHOLD_SECONDS", 0.0)
    monkeypatch.setattr(worker_module, "_terminal_cycle_collection_last_at", 0.0)
    monkeypatch.setattr(worker_module.gc, "collect", collect)
    with worker_module._terminal_cycle_collection_guard:
        worker_module._terminal_cycle_collection_scheduled_loops.discard(loop)

    caplog.set_level(logging.INFO, logger=worker_module.__name__)
    token = marker.set("run-context")
    try:
        worker_module._schedule_terminal_cycle_collection()
        worker_module._schedule_terminal_cycle_collection()
        await asyncio.wait_for(collected.wait(), timeout=1)

        async def wait_until_finished() -> None:
            while True:
                with worker_module._terminal_cycle_collection_guard:
                    if loop not in worker_module._terminal_cycle_collection_scheduled_loops:
                        return
                await asyncio.sleep(0)

        await asyncio.wait_for(wait_until_finished(), timeout=1)
    finally:
        marker.reset(token)

    assert len(seen) == 1
    assert seen[0][0] is None
    assert seen[0][1] != loop_thread_id
    assert "Terminal cyclic GC collected 7 object(s)" in caplog.text


@pytest.mark.anyio
async def test_run_agent_releases_terminal_runtime_callbacks():
    from deerflow.runtime.journal import RunJournal

    run_manager = RunManager()
    record = await run_manager.create("thread-runtime-release")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    captured: dict[str, Any] = {}

    class DummyAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, stream_mode, subgraphs
            captured["config"] = config
            callbacks = config.get("callbacks") or []
            captured["journal"] = next(callback for callback in callbacks if isinstance(callback, RunJournal))
            yield {"messages": []}

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(
            checkpointer=None,
            event_store=MemoryRunEventStore(),
        ),
        agent_factory=lambda **_kwargs: DummyAgent(),
        graph_input={},
        config={},
    )

    stream_config = captured["config"]
    journal = captured["journal"]
    assert "__pregel_runtime" not in stream_config["configurable"]
    assert "__run_journal" not in stream_config["context"]
    assert journal not in (stream_config.get("callbacks") or [])
    assert journal._closed is True
    assert journal._store is None
    assert journal._progress_reporter is None

    journal_ref = weakref.ref(journal)
    captured.clear()
    del journal
    await asyncio.sleep(0)
    assert journal_ref() is None


@pytest.mark.anyio
async def test_run_agent_threads_pre_existing_message_ids_into_runtime_context():
    run_manager = RunManager()
    record = await run_manager.create("thread-1")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    captured: dict[str, object] = {}

    class DummyCheckpointer:
        async def aget_tuple(self, _config):
            return SimpleNamespace(
                config={"configurable": {"checkpoint_id": "checkpoint-1"}},
                checkpoint={"channel_values": {}},
                metadata={},
                pending_writes=[],
            )

    class DummyAgent:
        async def aget_state(self, _config):
            return SimpleNamespace(
                values={
                    "messages": [
                        HumanMessage(id="h1", content="question"),
                        AIMessage(id="a1", content="answer"),
                    ]
                },
                config={
                    "configurable": {
                        "thread_id": "thread-1",
                        "checkpoint_ns": "",
                        "checkpoint_id": "checkpoint-1",
                    }
                },
                parent_config=None,
                metadata={},
                next=(),
                tasks=(),
                created_at=None,
            )

        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            captured["pre_existing_message_ids"] = config["context"][CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY]
            yield {"messages": []}

    def factory(*, config):
        return DummyAgent()

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(checkpointer=DummyCheckpointer()),
        agent_factory=factory,
        graph_input={},
        config={},
    )

    assert captured["pre_existing_message_ids"] == frozenset({"h1", "a1"})


@pytest.mark.anyio
async def test_run_agent_marks_rollback_unusable_when_capture_fails():
    """A failed pre-run capture must disable rollback entirely.

    Rollback now restores messages by forking the pre-run checkpoint through
    the graph (raw checkpoint blobs cannot reconstruct Delta-channel history),
    so there is no safe fallback when materialization fails: the worker must
    report ``snapshot_capture_failed=True`` with no rollback point rather
    than restoring an empty or partial message history.
    """
    run_manager = RunManager()
    record = await run_manager.create("thread-1")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )

    class DummyCheckpointer:
        async def aget_tuple(self, _config):
            return SimpleNamespace(
                config={
                    "configurable": {
                        "thread_id": "thread-1",
                        "checkpoint_ns": "",
                        "checkpoint_id": "checkpoint-1",
                    }
                },
                checkpoint={"id": "checkpoint-1", "channel_values": {}},
                metadata={"source": "loop"},
                pending_writes=[],
            )

    class DummyAgent:
        async def aget_state(self, _config):
            # Cancel after the worker crosses its startup barrier so this test
            # exercises the running rollback path, not pending cancellation.
            record.abort_action = "rollback"
            record.abort_event.set()
            raise RuntimeError("materialization failed")

        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            yield {"messages": []}

    with patch(
        "deerflow.runtime.runs.worker._rollback_to_pre_run_checkpoint",
        new_callable=AsyncMock,
    ) as rollback:
        await run_agent(
            bridge,
            run_manager,
            record,
            ctx=RunContext(checkpointer=DummyCheckpointer()),
            agent_factory=lambda *, config: DummyAgent(),
            graph_input={},
            config={},
        )

    rollback.assert_awaited_once()
    rollback_kwargs = rollback.await_args.kwargs
    assert rollback_kwargs["snapshot_capture_failed"] is True
    assert rollback_kwargs["rollback_point"] is None


@pytest.mark.anyio
async def test_run_agent_overrides_spoofed_pre_existing_message_ids_without_snapshot():
    run_manager = RunManager()
    record = await run_manager.create("thread-1")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    captured: dict[str, object] = {}

    class DummyAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            captured["pre_existing_message_ids"] = config["context"][CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY]
            yield {"messages": []}

    def factory(*, config):
        return DummyAgent()

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(checkpointer=None),
        agent_factory=factory,
        graph_input={},
        config={"context": {CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY: {"spoofed"}}},
    )

    assert captured["pre_existing_message_ids"] == frozenset()


@pytest.mark.anyio
async def test_run_agent_marks_llm_error_fallback_as_error_status():
    run_manager = RunManager()
    record = await run_manager.create("thread-1")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )

    class DummyAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            yield {
                "messages": [
                    AIMessage(
                        content="The configured LLM provider is temporarily unavailable after multiple retries.",
                        additional_kwargs={
                            "deerflow_error_fallback": True,
                            "error_type": "APIConnectionError",
                            "error_reason": "transient",
                            "error_detail": "Connection error.",
                        },
                    )
                ]
            }

    def factory(*, config):
        return DummyAgent()

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(checkpointer=None),
        agent_factory=factory,
        graph_input={},
        config={},
    )

    fetched = await run_manager.get(record.run_id)
    assert fetched is not None
    assert fetched.status == RunStatus.error
    assert fetched.error == "Connection error."
    bridge.publish_end.assert_awaited_once_with(record.run_id)


@pytest.mark.anyio
async def test_run_agent_rolls_back_failed_edit_replay_and_publishes_restored_values():
    run_manager = RunManager()
    record = await run_manager.create(
        "thread-1",
        metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-run"},
    )
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    before_messages = [HumanMessage(id="h1", content="original"), AIMessage(id="a1", content="answer")]

    class DummyCheckpointer:
        async def aget_tuple(self, _config):
            return SimpleNamespace(
                config={
                    "configurable": {
                        "thread_id": "thread-1",
                        "checkpoint_ns": "",
                        "checkpoint_id": "checkpoint-1",
                    }
                },
                checkpoint={"id": "checkpoint-1", "channel_values": {}},
                metadata={"source": "loop"},
                pending_writes=[],
            )

    class DummyAgent:
        async def aget_state(self, _config):
            return SimpleNamespace(
                values={"messages": before_messages},
                config={
                    "configurable": {
                        "thread_id": "thread-1",
                        "checkpoint_ns": "",
                        "checkpoint_id": "checkpoint-1",
                    }
                },
                parent_config=None,
                metadata={},
                next=(),
                tasks=(),
                created_at=None,
            )

        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            raise RuntimeError("edit replay failed")
            if False:
                yield  # pragma: no cover - keep this an async generator

    error_status_finalizing_states: list[bool] = []
    original_set_status = run_manager.set_status

    async def _set_status(run_id, status, **kwargs):
        if run_id == record.run_id and status == RunStatus.error:
            error_status_finalizing_states.append(record.finalizing)
        return await original_set_status(run_id, status, **kwargs)

    run_manager.set_status = _set_status  # type: ignore[method-assign]
    with patch(
        "deerflow.runtime.runs.worker._rollback_to_pre_run_checkpoint",
        new_callable=AsyncMock,
    ) as rollback:
        rollback.return_value = True
        await run_agent(
            bridge,
            run_manager,
            record,
            ctx=RunContext(checkpointer=DummyCheckpointer()),
            agent_factory=lambda *, config: DummyAgent(),
            graph_input={},
            config={},
        )

    fetched = await run_manager.get(record.run_id)
    assert fetched is not None
    assert fetched.status == RunStatus.error
    assert error_status_finalizing_states == [True]
    rollback.assert_awaited_once()
    publish_events = [call_args.args[1] for call_args in bridge.publish.await_args_list]
    assert "error" in publish_events
    assert "values" in publish_events
    assert publish_events.index("values") > publish_events.index("error")
    bridge.publish_end.assert_awaited_once_with(record.run_id)


@pytest.mark.anyio
async def test_failed_edit_replay_does_not_publish_restored_values_when_snapshot_capture_failed():
    run_manager = RunManager()
    record = await run_manager.create(
        "thread-1",
        metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-run"},
    )
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )

    class DummyCheckpointer:
        async def aget_tuple(self, _config):
            return SimpleNamespace(
                config={
                    "configurable": {
                        "thread_id": "thread-1",
                        "checkpoint_ns": "",
                        "checkpoint_id": "checkpoint-1",
                    }
                },
                checkpoint={"id": "checkpoint-1", "channel_values": {}},
                metadata={"source": "loop"},
                pending_writes=[],
            )

    class DummyAgent:
        async def aget_state(self, _config):
            raise RuntimeError("snapshot capture failed")

        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            raise RuntimeError("edit replay failed")
            if False:
                yield  # pragma: no cover - keep this an async generator

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(checkpointer=DummyCheckpointer()),
        agent_factory=lambda *, config: DummyAgent(),
        graph_input={},
        config={},
    )

    fetched = await run_manager.get(record.run_id)
    assert fetched is not None
    assert fetched.status == RunStatus.error
    publish_events = [call_args.args[1] for call_args in bridge.publish.await_args_list]
    assert "error" in publish_events
    assert "values" not in publish_events
    bridge.publish_end.assert_awaited_once_with(record.run_id)


@pytest.mark.anyio
async def test_run_agent_defaults_root_run_name_from_assistant_id():
    run_manager = RunManager()
    record = await run_manager.create("thread-1", assistant_id="lead_agent")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    captured: dict[str, object] = {}

    class DummyAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            captured["astream_run_name"] = config["run_name"]
            yield {"messages": []}

    def factory(*, config):
        captured["factory_run_name"] = config["run_name"]
        return DummyAgent()

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(checkpointer=None),
        agent_factory=factory,
        graph_input={},
        config={},
    )

    assert captured["factory_run_name"] == "lead_agent"
    assert captured["astream_run_name"] == "lead_agent"


@pytest.mark.anyio
async def test_run_agent_defaults_root_run_name_from_context_agent_name():
    run_manager = RunManager()
    record = await run_manager.create("thread-1", assistant_id="lead_agent")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    captured: dict[str, object] = {}

    class DummyAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            captured["astream_run_name"] = config["run_name"]
            yield {"messages": []}

    def factory(*, config):
        captured["factory_run_name"] = config["run_name"]
        return DummyAgent()

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(checkpointer=None),
        agent_factory=factory,
        graph_input={},
        config={"context": {"agent_name": "finalis"}},
    )

    assert captured["factory_run_name"] == "finalis"
    assert captured["astream_run_name"] == "finalis"


@pytest.mark.anyio
async def test_run_agent_defaults_root_run_name_from_configurable_agent_name():
    run_manager = RunManager()
    record = await run_manager.create("thread-1", assistant_id="lead_agent")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    captured: dict[str, object] = {}

    class DummyAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            captured["astream_run_name"] = config["run_name"]
            yield {"messages": []}

    def factory(*, config):
        captured["factory_run_name"] = config["run_name"]
        return DummyAgent()

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(checkpointer=None),
        agent_factory=factory,
        graph_input={},
        config={"configurable": {"agent_name": "finalis"}},
    )

    assert captured["factory_run_name"] == "finalis"
    assert captured["astream_run_name"] == "finalis"


@pytest.mark.anyio
async def test_rollback_forks_pre_run_checkpoint_without_deleting_thread(monkeypatch):
    checkpointer = FakeCheckpointer()
    mock_graph = _stub_mutation_graph(
        monkeypatch,
        restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}},
    )
    rollback_point = _make_rollback_point(
        pending_writes=[
            ("task-a", "messages", {"content": "first"}),
            ("task-a", "status", "done"),
            ("task-b", "events", {"type": "tool"}),
        ],
    )

    await _rollback_to_pre_run_checkpoint(
        accessor=_rollback_accessor(checkpointer),
        checkpointer=checkpointer,
        thread_id="thread-1",
        run_id="run-1",
        rollback_point=rollback_point,
        snapshot_capture_failed=False,
    )

    checkpointer.adelete_thread.assert_not_awaited()
    mock_graph.aupdate_state.assert_awaited_once()
    update_config, update_values = mock_graph.aupdate_state.await_args.args[:2]
    assert update_config["configurable"]["checkpoint_id"] == "ckpt-1"
    overwrite = update_values["messages"]
    assert isinstance(overwrite, Overwrite)
    assert overwrite.value == ["before"]
    assert mock_graph.aupdate_state.await_args.kwargs["as_node"] == "rollback_restore"
    assert checkpointer.aput_writes.await_args_list == [
        call(
            {"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}},
            [("messages", {"content": "first"}), ("status", "done")],
            task_id="task-a",
        ),
        call(
            {"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}},
            [("events", {"type": "tool"})],
            task_id="task-b",
        ),
    ]


@pytest.mark.anyio
async def test_rollback_skips_when_rollback_point_has_no_checkpoint_id(monkeypatch):
    """A rollback point without a checkpoint id cannot anchor a fork - skip safely."""
    checkpointer = FakeCheckpointer()
    mock_graph = _stub_mutation_graph(
        monkeypatch,
        restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}},
    )

    await _rollback_to_pre_run_checkpoint(
        accessor=_rollback_accessor(checkpointer),
        checkpointer=checkpointer,
        thread_id="thread-1",
        run_id="run-1",
        rollback_point=_make_rollback_point(checkpoint_id=None),
        snapshot_capture_failed=False,
    )

    checkpointer.adelete_thread.assert_not_awaited()
    mock_graph.aupdate_state.assert_not_awaited()
    checkpointer.aput_writes.assert_not_awaited()


class _ArtifactChannelState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]
    artifacts: Annotated[list[str], merge_artifacts]


def _build_message_and_artifact_graph(checkpointer: Any):
    async def _append(state: dict[str, Any]) -> dict[str, Any]:
        n = len(state.get("messages") or [])
        return {"messages": [HumanMessage(content=f"turn-{n}")], "artifacts": [f"a{n}"]}

    builder = StateGraph(_ArtifactChannelState)
    builder.add_node("append", _append)
    builder.set_entry_point("append")
    builder.set_finish_point("append")
    return builder.compile(checkpointer=checkpointer)


class _ExtensionFullState(TypedDict):
    """Full-mode schema with a channel base ThreadState does not know about.

    Mirrors a custom ``AgentMiddleware.state_schema`` contribution (memory,
    todos, uploads, …): absent from the base schema the mutation graph falls
    back to, so a base-schema restore silently drops it.
    """

    messages: Annotated[list[AnyMessage], add_messages]
    memory_notes: NotRequired[str]


class _ExtensionDeltaState(TypedDict):
    messages: Annotated[list[AnyMessage], DeltaChannel(merge_message_writes, snapshot_frequency=1000)]
    memory_notes: NotRequired[str]


@pytest.mark.anyio
@pytest.mark.parametrize(
    "mode,state_schema",
    [("full", _ExtensionFullState), ("delta", _ExtensionDeltaState)],
)
async def test_rollback_preserves_middleware_contributed_channels(mode, state_schema):
    """The rollback mutation graph must compile with the thread's effective schema.

    Regression: ``_rollback_to_pre_run_checkpoint`` used to build the mutation
    graph from the base ThreadState fallback, silently dropping channels
    contributed by custom middleware from the restored checkpoint.
    """
    checkpointer = InMemorySaver()

    async def _step(state: dict[str, Any]) -> dict[str, Any]:
        n = len(state.get("messages") or [])
        return {"messages": [HumanMessage(content=f"turn-{n}")], "memory_notes": f"note-{n}"}

    builder = StateGraph(state_schema)
    builder.add_node("step", _step)
    builder.set_entry_point("step")
    builder.set_finish_point("step")
    graph = builder.compile(checkpointer=checkpointer)

    accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode=mode)
    thread_config = {"configurable": {"thread_id": "thread-1"}}

    await graph.ainvoke({}, thread_config)
    rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config)
    assert rollback_point is not None

    await graph.ainvoke({}, thread_config)

    await _rollback_to_pre_run_checkpoint(
        accessor=accessor,
        checkpointer=checkpointer,
        thread_id="thread-1",
        run_id="run-1",
        rollback_point=rollback_point,
        snapshot_capture_failed=False,
    )

    latest_snapshot = await accessor.aget(thread_config)
    assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0"]
    assert latest_snapshot.values["memory_notes"] == "note-0"


@pytest.mark.anyio
async def test_rollback_restores_non_message_channels_via_fork_inheritance():
    """Rollback overwrites only messages; sibling reducer channels must be
    restored to their pre-run values by forking the pre-run checkpoint
    (non-updated channels inherit from the parent). Locks in the load-bearing
    assumption of the rollback design."""
    checkpointer = InMemorySaver()
    graph = _build_message_and_artifact_graph(checkpointer)
    accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="full")
    thread_config = {"configurable": {"thread_id": "thread-1"}}

    await graph.ainvoke({}, thread_config)
    rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config)
    assert rollback_point is not None

    await graph.ainvoke({}, thread_config)
    cancelled_snapshot = await accessor.aget(thread_config)
    assert cancelled_snapshot.values["artifacts"] == ["a0", "a1"]
    assert [message.content for message in cancelled_snapshot.values["messages"]] == ["turn-0", "turn-1"]

    await _rollback_to_pre_run_checkpoint(
        accessor=accessor,
        checkpointer=checkpointer,
        thread_id="thread-1",
        run_id="run-1",
        rollback_point=rollback_point,
        snapshot_capture_failed=False,
    )

    latest_snapshot = await accessor.aget(thread_config)
    assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0"]
    assert latest_snapshot.values["artifacts"] == ["a0"]


@pytest.mark.anyio
async def test_rollback_restored_checkpoint_becomes_latest_with_real_checkpointer():
    """Full-mode: the restored fork becomes latest, keeps pre-run values and
    pre-run pending writes, and drops writes attached to the cancelled run."""
    checkpointer = InMemorySaver()
    graph = _build_message_append_graph(_FullChannelState, checkpointer)
    accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="full")
    thread_config = {"configurable": {"thread_id": "thread-1"}}

    await graph.ainvoke({}, thread_config)
    pre_run_snapshot = await accessor.aget(thread_config)
    pre_run_checkpoint_id = pre_run_snapshot.config["configurable"]["checkpoint_id"]
    pre_run_config = {
        "configurable": {
            "thread_id": "thread-1",
            "checkpoint_ns": "",
            "checkpoint_id": pre_run_checkpoint_id,
        }
    }
    checkpointer.put_writes(pre_run_config, [("title", "pending-before")], task_id="task-before")

    rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config)
    assert rollback_point is not None
    assert rollback_point.pending_writes == (("task-before", "title", "pending-before"),)

    await graph.ainvoke({}, thread_config)
    cancelled_snapshot = await accessor.aget(thread_config)
    cancelled_checkpoint_id = cancelled_snapshot.config["configurable"]["checkpoint_id"]
    checkpointer.put_writes(
        {"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": cancelled_checkpoint_id}},
        [("title", "pending-after")],
        task_id="task-after",
    )

    await _rollback_to_pre_run_checkpoint(
        accessor=accessor,
        checkpointer=checkpointer,
        thread_id="thread-1",
        run_id="run-1",
        rollback_point=rollback_point,
        snapshot_capture_failed=False,
    )

    latest_snapshot = await accessor.aget(thread_config)
    restored_checkpoint_id = latest_snapshot.config["configurable"]["checkpoint_id"]
    assert restored_checkpoint_id not in (pre_run_checkpoint_id, cancelled_checkpoint_id)
    assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0"]

    latest_tuple = await checkpointer.aget_tuple(thread_config)
    assert latest_tuple.pending_writes == [("task-before", "title", "pending-before")]
    restored_tuple = await checkpointer.aget_tuple({"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": restored_checkpoint_id}})
    assert restored_tuple.parent_config["configurable"]["checkpoint_id"] == pre_run_checkpoint_id


@pytest.mark.anyio
async def test_rollback_deletes_thread_when_no_snapshot_exists():
    checkpointer = FakeCheckpointer()

    await _rollback_to_pre_run_checkpoint(
        accessor=_rollback_accessor(checkpointer),
        checkpointer=checkpointer,
        thread_id="thread-1",
        run_id="run-1",
        rollback_point=None,
        snapshot_capture_failed=False,
    )

    checkpointer.adelete_thread.assert_awaited_once_with("thread-1")
    checkpointer.aput_writes.assert_not_awaited()


@pytest.mark.anyio
async def test_rollback_raises_when_restore_config_has_no_checkpoint_id(monkeypatch):
    checkpointer = FakeCheckpointer()
    _stub_mutation_graph(
        monkeypatch,
        restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}},
    )

    with pytest.raises(RuntimeError, match="did not return checkpoint_id"):
        await _rollback_to_pre_run_checkpoint(
            accessor=_rollback_accessor(checkpointer),
            checkpointer=checkpointer,
            thread_id="thread-1",
            run_id="run-1",
            rollback_point=_make_rollback_point(pending_writes=[("task-a", "messages", "value")]),
            snapshot_capture_failed=False,
        )

    checkpointer.adelete_thread.assert_not_awaited()
    checkpointer.aput_writes.assert_not_awaited()


@pytest.mark.anyio
async def test_capture_rollback_point_normalizes_missing_checkpoint_ns():
    """Snapshots whose config omits ``checkpoint_ns`` must anchor the root namespace."""

    class _TupleCheckpointer:
        async def aget_tuple(self, _config):
            return SimpleNamespace(pending_writes=[])

    class _SnapshotAgent:
        async def aget_state(self, _config):
            return SimpleNamespace(
                values={"messages": [HumanMessage(id="h1", content="before")]},
                config={"configurable": {"thread_id": "thread-1", "checkpoint_id": "ckpt-1"}},
                metadata={"source": "loop"},
            )

    checkpointer = _TupleCheckpointer()
    accessor = CheckpointStateAccessor(graph=_SnapshotAgent(), checkpointer=checkpointer, mode="full")

    rollback_point = await _capture_rollback_point(accessor, checkpointer, {"configurable": {"thread_id": "thread-1"}})

    assert rollback_point is not None
    assert rollback_point.config["configurable"]["checkpoint_ns"] == ""
    assert rollback_point.config["configurable"]["checkpoint_id"] == "ckpt-1"
    assert [message.content for message in rollback_point.messages] == ["before"]
    assert rollback_point.metadata == {"source": "loop"}


@pytest.mark.anyio
async def test_capture_rollback_point_returns_none_without_checkpoint():
    """A thread with no checkpoints has no rollback point (delete/reset contract)."""

    class _EmptyAgent:
        async def aget_state(self, _config):
            return SimpleNamespace(
                values={},
                config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}},
                metadata={},
            )

    checkpointer = SimpleNamespace(aget_tuple=AsyncMock(return_value=None))
    accessor = CheckpointStateAccessor(graph=_EmptyAgent(), checkpointer=checkpointer, mode="full")

    assert await _capture_rollback_point(accessor, checkpointer, {"configurable": {"thread_id": "thread-1"}}) is None


@pytest.mark.anyio
async def test_rollback_raises_on_malformed_pending_write_not_a_tuple(monkeypatch):
    """pending_writes containing a non-3-tuple item should raise RuntimeError."""
    checkpointer = FakeCheckpointer()
    _stub_mutation_graph(
        monkeypatch,
        restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}},
    )

    with pytest.raises(RuntimeError, match="rollback failed: pending_write is not a 3-tuple"):
        await _rollback_to_pre_run_checkpoint(
            accessor=_rollback_accessor(checkpointer),
            checkpointer=checkpointer,
            thread_id="thread-1",
            run_id="run-1",
            rollback_point=_make_rollback_point(
                pending_writes=[
                    ("task-a", "messages", "valid"),  # valid
                    ["only", "two"],  # malformed: only 2 elements
                ],
            ),
            snapshot_capture_failed=False,
        )

    checkpointer.aput_writes.assert_not_awaited()


@pytest.mark.anyio
async def test_rollback_raises_on_malformed_pending_write_non_string_channel(monkeypatch):
    """pending_writes containing a non-string channel should raise RuntimeError."""
    checkpointer = FakeCheckpointer()
    _stub_mutation_graph(
        monkeypatch,
        restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}},
    )

    with pytest.raises(RuntimeError, match="rollback failed: pending_write has non-string channel"):
        await _rollback_to_pre_run_checkpoint(
            accessor=_rollback_accessor(checkpointer),
            checkpointer=checkpointer,
            thread_id="thread-1",
            run_id="run-1",
            rollback_point=_make_rollback_point(
                pending_writes=[
                    ("task-a", 123, "value"),  # malformed: channel is not a string
                ],
            ),
            snapshot_capture_failed=False,
        )

    checkpointer.aput_writes.assert_not_awaited()


@pytest.mark.anyio
async def test_rollback_propagates_aput_writes_failure(monkeypatch):
    """If aput_writes fails, the exception should propagate (not be swallowed)."""
    checkpointer = FakeCheckpointer()
    checkpointer.aput_writes.side_effect = RuntimeError("Database connection lost")
    _stub_mutation_graph(
        monkeypatch,
        restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}},
    )

    with pytest.raises(RuntimeError, match="Database connection lost"):
        await _rollback_to_pre_run_checkpoint(
            accessor=_rollback_accessor(checkpointer),
            checkpointer=checkpointer,
            thread_id="thread-1",
            run_id="run-1",
            rollback_point=_make_rollback_point(
                pending_writes=[
                    ("task-a", "messages", "value"),
                ],
            ),
            snapshot_capture_failed=False,
        )

    checkpointer.aput_writes.assert_awaited_once()


@pytest.mark.anyio
async def test_rollback_linearizes_delta_restore_onto_cancelled_head():
    """Delta rollback replaces state on the head instead of creating a fork.

    The cancelled path has already attached writes to the pre-run checkpoint,
    so forking that checkpoint would replay sibling writes. The restored
    checkpoint must instead descend from the cancelled head and replace the
    captured messages there.
    """
    checkpointer = InMemorySaver()
    graph = _build_message_append_graph(_DeltaChannelState, checkpointer)
    accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="delta")
    thread_config = {"configurable": {"thread_id": "thread-1"}}

    await graph.ainvoke({}, thread_config)
    await graph.ainvoke({}, thread_config)
    rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config)
    assert rollback_point is not None
    assert [message.content for message in rollback_point.messages] == ["turn-0", "turn-1"]
    pre_run_checkpoint_id = rollback_point.config["configurable"]["checkpoint_id"]

    await graph.ainvoke({}, thread_config)  # the run that gets cancelled
    cancelled_snapshot = await accessor.aget(thread_config)
    cancelled_checkpoint_id = cancelled_snapshot.config["configurable"]["checkpoint_id"]
    assert len(cancelled_snapshot.values["messages"]) == 3

    await _rollback_to_pre_run_checkpoint(
        accessor=accessor,
        checkpointer=checkpointer,
        thread_id="thread-1",
        run_id="run-1",
        rollback_point=rollback_point,
        snapshot_capture_failed=False,
    )

    latest_snapshot = await accessor.aget(thread_config)
    restored_checkpoint_id = latest_snapshot.config["configurable"]["checkpoint_id"]
    assert restored_checkpoint_id not in (pre_run_checkpoint_id, cancelled_checkpoint_id)
    assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0", "turn-1"]

    restored_tuple = await checkpointer.aget_tuple({"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": restored_checkpoint_id}})
    assert restored_tuple.parent_config["configurable"]["checkpoint_id"] == cancelled_checkpoint_id
    assert restored_tuple.metadata.get("source") == "update"
    # A non-snapshot Delta checkpoint must not persist the full message list.
    raw_messages = restored_tuple.checkpoint.get("channel_values", {}).get("messages")
    assert not isinstance(raw_messages, list)


@pytest.mark.anyio
@pytest.mark.parametrize(
    "mode,state_schema",
    [("full", _CompactionFullState), ("delta", _CompactionDeltaState)],
)
async def test_rollback_preserves_agent_binding_for_manual_compaction(monkeypatch, mode, state_schema):
    """A state-only rollback must keep the policy that produced its state."""
    import deerflow.config.agents_config as agents_config

    checkpointer = InMemorySaver()

    async def _finish(_state: dict[str, Any]) -> dict[str, Any]:
        return {}

    builder = StateGraph(state_schema)
    builder.add_node("finish", _finish)
    builder.set_entry_point("finish")
    builder.set_finish_point("finish")
    graph = builder.compile(checkpointer=checkpointer)
    accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode=mode)
    thread_config = {"configurable": {"thread_id": "thread-binding"}}
    seeded_config = {
        **thread_config,
        "metadata": {CHECKPOINT_AGENT_NAME_METADATA_KEY: "stateless-worker"},
    }
    original_messages = [
        HumanMessage(content="old question", id="h1"),
        AIMessage(content="old answer", id="a1"),
        HumanMessage(content="latest question", id="h2"),
    ]
    await graph.ainvoke({"messages": original_messages}, seeded_config)
    rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config)
    assert rollback_point is not None

    await graph.ainvoke(
        {"messages": [AIMessage(content="cancelled answer", id="a2")]},
        thread_config,
    )
    await _rollback_to_pre_run_checkpoint(
        accessor=accessor,
        checkpointer=checkpointer,
        thread_id="thread-binding",
        run_id="run-binding",
        rollback_point=rollback_point,
        snapshot_capture_failed=False,
    )

    restored = await accessor.aget(thread_config)
    assert restored.metadata[CHECKPOINT_AGENT_NAME_METADATA_KEY] == "stateless-worker"
    assert [message.id for message in restored.values["messages"]] == ["h1", "a1", "h2"]

    config_reads: list[tuple[str, str | None]] = []

    def _load_agent_config(name, *, user_id=None):
        config_reads.append((name, user_id))
        return SimpleNamespace(model=None, memory_enabled=False)

    monkeypatch.setattr(agents_config, "load_agent_config", _load_agent_config)
    captured: dict[str, Any] = {}

    class _CompactionMiddleware:
        async def acompact_state(self, state, runtime, *, force=False, raise_on_failure=False):
            del runtime, force, raise_on_failure
            return SimpleNamespace(
                summary_text="summary",
                messages_to_summarize=tuple(state["messages"][:-1]),
                preserved_messages=tuple(state["messages"][-1:]),
                total_tokens=42,
            )

    def _create_compaction_middleware(**kwargs):
        captured.update(kwargs)
        return _CompactionMiddleware()

    monkeypatch.setattr(context_compaction, "_create_compaction_middleware", _create_compaction_middleware)
    compaction_graph = build_state_mutation_graph("manual_compaction", mode, state_schema)
    compaction_accessor = CheckpointStateAccessor.bind(compaction_graph, checkpointer, mode=mode)

    result = await compact_thread_context(
        compaction_accessor,
        "thread-binding",
        app_config=SimpleNamespace(models=[]),
        user_id="user-1",
        agent_name="memory-enabled-impostor",
    )

    assert result.compacted is True
    assert captured["skip_memory_flush"] is True
    assert config_reads == [("stateless-worker", "user-1")]


@pytest.mark.anyio
async def test_rollback_restores_pre_run_pending_writes_for_delta_checkpoints():
    """Pre-run pending writes are re-attached to the restored checkpoint; writes
    attached to the cancelled run are not."""
    checkpointer = InMemorySaver()
    graph = _build_message_append_graph(_DeltaChannelState, checkpointer)
    accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="delta")
    thread_config = {"configurable": {"thread_id": "thread-1"}}

    await graph.ainvoke({}, thread_config)
    pre_run_snapshot = await accessor.aget(thread_config)
    pre_run_checkpoint_id = pre_run_snapshot.config["configurable"]["checkpoint_id"]
    checkpointer.put_writes(
        {"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": pre_run_checkpoint_id}},
        [("title", "in-flight-title")],
        task_id="task-in-flight",
    )

    rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config)
    assert rollback_point is not None
    assert rollback_point.pending_writes == (("task-in-flight", "title", "in-flight-title"),)

    await graph.ainvoke({}, thread_config)  # cancelled run
    cancelled_snapshot = await accessor.aget(thread_config)
    cancelled_checkpoint_id = cancelled_snapshot.config["configurable"]["checkpoint_id"]
    checkpointer.put_writes(
        {"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": cancelled_checkpoint_id}},
        [("title", "cancelled-title")],
        task_id="task-cancelled",
    )

    await _rollback_to_pre_run_checkpoint(
        accessor=accessor,
        checkpointer=checkpointer,
        thread_id="thread-1",
        run_id="run-1",
        rollback_point=rollback_point,
        snapshot_capture_failed=False,
    )

    latest_tuple = await checkpointer.aget_tuple(thread_config)
    assert latest_tuple.pending_writes == [("task-in-flight", "title", "in-flight-title")]
    assert [message.content for message in (await accessor.aget(thread_config)).values["messages"]] == ["turn-0"]


@pytest.mark.anyio
async def test_rollback_linearizes_delta_restore_sqlite_reopen(tmp_path):
    """Same linear restore contract against a disk-backed saver, verified after a
    close/reopen so only persisted bytes can satisfy the assertions."""
    db_path = tmp_path / "rollback.sqlite3"

    async with AsyncSqliteSaver.from_conn_string(str(db_path)) as checkpointer:
        await checkpointer.setup()
        graph = _build_message_append_graph(_DeltaChannelState, checkpointer)
        accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="delta")
        thread_config = {"configurable": {"thread_id": "thread-1"}}

        await graph.ainvoke({}, thread_config)
        await graph.ainvoke({}, thread_config)
        rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config)
        assert rollback_point is not None
        pre_run_checkpoint_id = rollback_point.config["configurable"]["checkpoint_id"]

        await graph.ainvoke({}, thread_config)  # cancelled run
        cancelled_snapshot = await accessor.aget(thread_config)
        cancelled_checkpoint_id = cancelled_snapshot.config["configurable"]["checkpoint_id"]

        await _rollback_to_pre_run_checkpoint(
            accessor=accessor,
            checkpointer=checkpointer,
            thread_id="thread-1",
            run_id="run-1",
            rollback_point=rollback_point,
            snapshot_capture_failed=False,
        )

        restored_snapshot = await accessor.aget(thread_config)
        restored_checkpoint_id = restored_snapshot.config["configurable"]["checkpoint_id"]
        assert restored_checkpoint_id not in (pre_run_checkpoint_id, cancelled_checkpoint_id)
        assert [message.content for message in restored_snapshot.values["messages"]] == ["turn-0", "turn-1"]
        restored_tuple = await checkpointer.aget_tuple({"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": restored_checkpoint_id}})
        assert restored_tuple.parent_config["configurable"]["checkpoint_id"] == cancelled_checkpoint_id

    async with AsyncSqliteSaver.from_conn_string(str(db_path)) as checkpointer:
        graph = _build_message_append_graph(_DeltaChannelState, checkpointer)
        accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="delta")
        thread_config = {"configurable": {"thread_id": "thread-1"}}

        latest_snapshot = await accessor.aget(thread_config)
        assert latest_snapshot.config["configurable"]["checkpoint_id"] == restored_checkpoint_id
        assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0", "turn-1"]

        restored_tuple = await checkpointer.aget_tuple({"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": restored_checkpoint_id}})
        assert restored_tuple.parent_config["configurable"]["checkpoint_id"] == cancelled_checkpoint_id
        raw_messages = restored_tuple.checkpoint.get("channel_values", {}).get("messages")
        assert not isinstance(raw_messages, list)


def test_agent_factory_supports_app_config_detects_supported_signature():
    def factory(*, config, app_config=None):
        return (config, app_config)

    assert _agent_factory_supports_app_config(factory) is True


def test_build_runtime_context_defaults_to_thread_and_run_id():
    ctx = _build_runtime_context("thread-1", "run-1", None)
    assert ctx == {"thread_id": "thread-1", "run_id": "run-1"}


@pytest.mark.parametrize("forged_reader", [lambda: None, {"allowed_ids": ["other-thread"]}])
@pytest.mark.parametrize("carrier", ["context", "configurable"])
def test_embedded_caller_cannot_supply_conversation_reader(forged_reader, carrier):
    key = "__conversation_reader"
    config = {carrier: {key: forged_reader}}

    runtime_context = _build_runtime_context("thread-1", "run-1", config.get("context"))
    _install_runtime_context(config, runtime_context)

    assert key not in runtime_context
    assert key not in config["context"]
    assert key not in config.get("configurable", {})


def test_host_conversation_reader_replaces_caller_value_in_both_contexts():
    key = "__conversation_reader"
    reader = AsyncMock()
    config = {"context": {key: AsyncMock()}, "configurable": {key: AsyncMock()}}

    runtime_context = _build_runtime_context("thread-1", "run-1", config["context"], conversation_reader=reader)
    _install_runtime_context(config, runtime_context)

    assert runtime_context[key] is reader
    assert config["context"][key] is reader
    assert key not in config["configurable"]


@pytest.mark.anyio
@pytest.mark.parametrize("outcome", ["success", "error", "cancel"])
async def test_run_agent_scopes_conversation_reader_to_the_active_run(outcome):
    key = "__conversation_reader"
    reader = AsyncMock(return_value="authorized conversation")
    run_manager = RunManager()
    record = await run_manager.create(f"thread-conversation-{outcome}")
    captured: dict[str, Any] = {}
    config = {"context": {key: AsyncMock()}}

    class DummyAgent:
        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, stream_mode, subgraphs
            captured["stream_config"] = config
            runtime_context = config["configurable"]["__pregel_runtime"].context
            captured["runtime_context"] = runtime_context
            assert config["context"][key] is reader
            assert runtime_context[key] is reader
            assert await runtime_context[key]("allowed-thread") == "authorized conversation"
            if outcome == "error":
                raise RuntimeError("model failed")
            if outcome == "cancel":
                record.abort_event.set()
            yield {"messages": []}

    def factory(*, config):
        captured["factory_context"] = config["context"]
        assert config["context"][key] is reader
        return DummyAgent()

    await run_agent(
        _lease_test_bridge(),
        run_manager,
        record,
        ctx=RunContext(checkpointer=None, conversation_reader=reader),
        agent_factory=factory,
        graph_input={},
        config=config,
    )
    await asyncio.sleep(0)

    reader.assert_awaited_once_with("allowed-thread")
    expected_status = {"success": RunStatus.success, "error": RunStatus.error, "cancel": RunStatus.interrupted}[outcome]
    assert record.status == expected_status
    for context in (config["context"], captured["factory_context"], captured["stream_config"]["context"], captured["runtime_context"]):
        assert key not in context


def test_build_runtime_context_merges_caller_context():
    """Regression for issue #2677: keys from ``config['context']`` (e.g. ``agent_name``)
    must be merged into the Runtime's context so that ``ToolRuntime.context`` — which
    is what ``setup_agent`` reads — can see them."""
    caller_context = {"agent_name": "my-agent", "is_bootstrap": True, "model_name": "gpt-4"}

    ctx = _build_runtime_context("thread-1", "run-1", caller_context)

    assert ctx["thread_id"] == "thread-1"
    assert ctx["run_id"] == "run-1"
    assert ctx["agent_name"] == "my-agent"
    assert ctx["is_bootstrap"] is True
    assert ctx["model_name"] == "gpt-4"


def test_build_runtime_context_caller_cannot_override_thread_id_or_run_id():
    """A malicious or buggy caller must not be able to overwrite the worker-assigned
    ``thread_id`` / ``run_id`` by stuffing them into ``config['context']``."""
    caller_context = {"thread_id": "spoofed", "run_id": "spoofed", "agent_name": "ok"}

    ctx = _build_runtime_context("real-thread", "real-run", caller_context)

    assert ctx["thread_id"] == "real-thread"
    assert ctx["run_id"] == "real-run"
    assert ctx["agent_name"] == "ok"


def test_build_runtime_context_ignores_caller_pre_existing_message_ids():
    caller_context = {CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY: {"spoofed"}}

    ctx = _build_runtime_context("thread-1", "run-1", caller_context)

    assert CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY not in ctx


def test_build_runtime_context_ignores_caller_sandbox_execution_identities():
    caller_context = {
        SANDBOX_LEASE_OWNER_CONTEXT_KEY: "forged-owner",
        SANDBOX_COMMAND_SCOPE_CONTEXT_KEY: "forged-scope",
    }

    ctx = _build_runtime_context("thread-1", "run-1", caller_context)

    assert SANDBOX_LEASE_OWNER_CONTEXT_KEY not in ctx
    assert SANDBOX_COMMAND_SCOPE_CONTEXT_KEY not in ctx


def test_build_runtime_context_ignores_caller_audit_attribution_and_recorders():
    caller_context = {
        "is_subagent": True,
        "agent_id": "forged-agent",
        "__run_loop_detection_recorder": object(),
        "__run_tool_promotion_recorder": object(),
        "__run_tool_progress_recorder": object(),
    }

    ctx = _build_runtime_context("thread-1", "run-1", caller_context)

    assert set(caller_context).isdisjoint(ctx)


def test_build_runtime_context_ignores_non_dict_caller_context():
    ctx = _build_runtime_context("thread-1", "run-1", "not-a-dict")
    assert ctx == {"thread_id": "thread-1", "run_id": "run-1"}


def test_agent_factory_supports_app_config_returns_false_when_signature_lookup_fails(monkeypatch):
    class BrokenCallable:
        def __call__(self, **kwargs):
            return kwargs

    monkeypatch.setattr("deerflow.runtime.runs.worker.inspect.signature", lambda _obj: (_ for _ in ()).throw(ValueError("boom")))

    assert _agent_factory_supports_app_config(BrokenCallable()) is False


# ---------------------------------------------------------------------------
# _extract_llm_error_fallback_message coverage
# ---------------------------------------------------------------------------


def test_try_extract_from_message_finds_fallback_on_message_object():
    msg = AIMessage(
        content="fallback",
        additional_kwargs={
            "deerflow_error_fallback": True,
            "error_detail": "Connection error.",
            "error_reason": "transient",
        },
    )
    assert _try_extract_from_message(msg) == "Connection error."


def test_try_extract_from_message_finds_fallback_on_dict():
    msg = {
        "content": "fallback",
        "additional_kwargs": {
            "deerflow_error_fallback": True,
            "error_detail": "Quota exceeded.",
        },
    }
    assert _try_extract_from_message(msg) == "Quota exceeded."


def test_try_extract_from_message_returns_none_for_normal_message():
    msg = AIMessage(content="hello")
    assert _try_extract_from_message(msg) is None


def test_extract_llm_error_fallback_message_large_state_chunk_no_fallback():
    """Normal-size state dict without fallback markers must not raise and should return None."""
    large_state = {
        "messages": [
            AIMessage(content="Hello!"),
            {"role": "user", "content": "Hi there"},
        ],
        "foo": "x" * 10_000,
        "bar": {"nested": {"deep": {"data": list(range(1000))}}},
        "baz": [{"id": i, "payload": "y" * 1000} for i in range(500)],
    }
    assert _extract_llm_error_fallback_message(large_state) is None


def test_extract_llm_error_fallback_message_finds_fallback_in_messages_list():
    state = {
        "messages": [
            AIMessage(content="Hello!"),
            AIMessage(
                content="Unavailable.",
                additional_kwargs={
                    "deerflow_error_fallback": True,
                    "error_detail": "Connection error.",
                },
            ),
        ],
        "other_state": "large_value" * 1000,
    }
    assert _extract_llm_error_fallback_message(state) == "Connection error."


def test_extract_llm_error_fallback_message_finds_fallback_in_raw_message():
    msg = AIMessage(
        content="Unavailable.",
        additional_kwargs={
            "deerflow_error_fallback": True,
            "error_reason": "quota",
        },
    )
    assert _extract_llm_error_fallback_message(msg) == "quota"


def test_extract_llm_error_fallback_message_finds_fallback_in_tuple():
    item = (
        "messages",
        AIMessage(
            content="Unavailable.",
            additional_kwargs={
                "deerflow_error_fallback": True,
                "error_detail": "Circuit open.",
            },
        ),
    )
    assert _extract_llm_error_fallback_message(item) == "Circuit open."


def test_extract_llm_error_fallback_message_returns_none_for_empty_values():
    assert _extract_llm_error_fallback_message({}) is None
    assert _extract_llm_error_fallback_message([]) is None
    assert _extract_llm_error_fallback_message(None) is None
    assert _extract_llm_error_fallback_message("string") is None


def test_extract_llm_error_fallback_message_finds_fallback_in_updates_mode():
    """stream_mode='updates' yields dicts keyed by node name (e.g. {'call_model': {...}}).
    Fallback marker is nested inside the node's state update, not at the top level."""
    update_chunk = {
        "call_model": {
            "messages": [
                AIMessage(
                    content="Unavailable.",
                    additional_kwargs={
                        "deerflow_error_fallback": True,
                        "error_detail": "Connection error.",
                    },
                )
            ]
        }
    }
    assert _extract_llm_error_fallback_message(update_chunk) == "Connection error."


def test_extract_llm_error_fallback_message_updates_mode_no_fallback():
    """Normal updates chunk without any fallback should return None safely."""
    update_chunk = {
        "__interrupt__": [
            {
                "value": "ask_human",
                "resumable": True,
                "ns": ["agent"],
                "when": "during",
            }
        ]
    }
    assert _extract_llm_error_fallback_message(update_chunk) is None


# ---------------------------------------------------------------------------
# pre_existing_ids filtering — stale fallback markers from prior runs
# ---------------------------------------------------------------------------


def test_try_extract_skips_message_with_pre_existing_id():
    """Fallback marker on a message whose id is in pre_existing_ids must be ignored."""
    msg = AIMessage(
        id="stale-1",
        content="Unavailable.",
        additional_kwargs={
            "deerflow_error_fallback": True,
            "error_detail": "Connection error.",
        },
    )
    assert _try_extract_from_message(msg, {"stale-1"}) is None
    # Without the filter, the same message would still surface the marker.
    assert _try_extract_from_message(msg) == "Connection error."


def test_try_extract_still_finds_fresh_message_when_others_are_stale():
    """A non-stale message with a fallback marker must still match."""
    msg = AIMessage(
        id="fresh-1",
        content="Unavailable.",
        additional_kwargs={
            "deerflow_error_fallback": True,
            "error_detail": "Connection error.",
        },
    )
    assert _try_extract_from_message(msg, {"stale-1", "stale-2"}) == "Connection error."


def test_try_extract_skips_dict_message_with_pre_existing_id():
    msg = {
        "id": "stale-2",
        "content": "Unavailable.",
        "additional_kwargs": {
            "deerflow_error_fallback": True,
            "error_detail": "Quota exceeded.",
        },
    }
    assert _try_extract_from_message(msg, {"stale-2"}) is None
    assert _try_extract_from_message(msg) == "Quota exceeded."


def test_extract_llm_error_fallback_message_skips_stale_history():
    """A state chunk replaying a stale fallback marker from a prior run must return None."""
    state = {
        "messages": [
            AIMessage(id="stale-1", content="Hi"),
            AIMessage(
                id="stale-fallback",
                content="Unavailable.",
                additional_kwargs={
                    "deerflow_error_fallback": True,
                    "error_detail": "Connection error.",
                },
            ),
        ]
    }
    assert _extract_llm_error_fallback_message(state, {"stale-1", "stale-fallback"}) is None


def test_extract_llm_error_fallback_message_returns_fresh_marker_alongside_stale_history():
    """Stale history is ignored, but a brand-new fallback in the same chunk is reported."""
    state = {
        "messages": [
            AIMessage(id="stale-1", content="Hi"),
            AIMessage(
                id="stale-fallback",
                content="Old failure.",
                additional_kwargs={
                    "deerflow_error_fallback": True,
                    "error_detail": "Old error.",
                },
            ),
            AIMessage(
                id="fresh-fallback",
                content="New failure.",
                additional_kwargs={
                    "deerflow_error_fallback": True,
                    "error_detail": "Fresh error.",
                },
            ),
        ]
    }
    assert _extract_llm_error_fallback_message(state, {"stale-1", "stale-fallback"}) == "Fresh error."


def test_extract_llm_error_fallback_message_default_filter_is_empty():
    """Passing no pre_existing_ids must preserve the original (pre-fix) behavior."""
    state = {
        "messages": [
            AIMessage(
                id="any",
                content="Unavailable.",
                additional_kwargs={
                    "deerflow_error_fallback": True,
                    "error_detail": "Connection error.",
                },
            )
        ]
    }
    assert _extract_llm_error_fallback_message(state) == "Connection error."


def test_collect_pre_existing_message_ids_pulls_ids_from_values():
    values = {
        "messages": [
            AIMessage(id="a", content="x"),
            AIMessage(id="b", content="y"),
            AIMessage(content="no-id-here"),  # ignored
        ]
    }
    assert _collect_pre_existing_message_ids(values) == {"a", "b"}


def test_collect_pre_existing_message_ids_handles_missing_pieces():
    assert _collect_pre_existing_message_ids(None) == set()
    assert _collect_pre_existing_message_ids({}) == set()
    assert _collect_pre_existing_message_ids({"messages": None}) == set()


@pytest.mark.anyio
async def test_run_agent_ignores_stale_llm_error_fallback_from_prior_run():
    """A stale fallback marker checkpointed by an earlier run on the same thread
    must NOT cause a successful current run to be reported as ``error``.

    This guards against the regression where one IndexError-driven failure (now
    classified transient and surfaced as a ``deerflow_error_fallback`` AIMessage)
    persisted in thread history and tripped ``RunStatus.error`` on every
    subsequent run that re-played the messages channel via ``stream_mode="values"``.
    """
    run_manager = RunManager()
    record = await run_manager.create("thread-1")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )

    stale_fallback = AIMessage(
        id="stale-fallback",
        content="Old failure.",
        additional_kwargs={
            "deerflow_error_fallback": True,
            "error_type": "IndexError",
            "error_reason": "transient",
            "error_detail": "list index out of range",
        },
    )

    class StaleHistoryCheckpointer:
        async def aget_tuple(self, config):
            checkpoint = empty_checkpoint()
            checkpoint["id"] = "ckpt-stale"
            checkpoint["channel_values"] = {"messages": [stale_fallback]}
            return SimpleNamespace(
                config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "ckpt-stale"}},
                checkpoint=checkpoint,
                metadata={},
                pending_writes=[],
            )

    class DummyAgent:
        async def aget_state(self, _config):
            return SimpleNamespace(
                values={"messages": [stale_fallback]},
                config={
                    "configurable": {
                        "thread_id": "thread-1",
                        "checkpoint_ns": "",
                        "checkpoint_id": "ckpt-stale",
                    }
                },
                parent_config=None,
                metadata={},
                next=(),
                tasks=(),
                created_at=None,
            )

        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            # Replay the prior fallback message (as LangGraph would when using
            # stream_mode="values") and then yield a fresh successful AIMessage.
            yield {
                "messages": [
                    stale_fallback,
                    AIMessage(id="fresh-ok", content="Hello — the run succeeded."),
                ]
            }

    def factory(*, config):
        return DummyAgent()

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(checkpointer=StaleHistoryCheckpointer()),
        agent_factory=factory,
        graph_input={},
        config={},
    )

    fetched = await run_manager.get(record.run_id)
    assert fetched is not None
    assert fetched.status == RunStatus.success, f"Stale fallback marker from prior run should not flip current run to error, got status={fetched.status} error={fetched.error!r}"
    bridge.publish_end.assert_awaited_once_with(record.run_id)


class _FakeCheckpointTuple:
    """Minimal stand-in for ``CheckpointTuple`` used by ``_ensure_interrupted_title``."""

    def __init__(self, *, checkpoint: dict, metadata: dict, config: dict | None = None):
        self.checkpoint = checkpoint
        self.metadata = metadata
        self.config = config or {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}}


class _TitleCheckpointer:
    """Captures ``aput`` arguments and exposes ``get_next_version`` like DB savers."""

    def __init__(self, *, tuple_value: _FakeCheckpointTuple | None, put_result: dict | None = None):
        self.aget_tuple = AsyncMock(return_value=tuple_value)
        self.aput = AsyncMock(return_value=put_result or {})

    def get_next_version(self, current, _channel):
        if current is None:
            return 1
        if isinstance(current, int):
            return current + 1
        if isinstance(current, str):
            try:
                return str(int(current) + 1)
            except ValueError:
                return f"{current}.1"
        return 1


@pytest.mark.anyio
async def test_interrupted_title_finalization_blocks_new_same_thread_run(monkeypatch):
    """A cancelled run must remain active while its title-only checkpoint is finalizing."""
    from deerflow.agents.middlewares.title_middleware import TitleMiddleware

    monkeypatch.setattr(
        TitleMiddleware,
        "_generate_title_result",
        lambda self, state, allow_partial_exchange=False: {"title": "Old Prompt"},
    )

    initial_checkpoint = {
        "id": "ckpt-old",
        "ts": "2026-06-29T00:00:00Z",
        "channel_values": {"messages": [{"type": "human", "content": "old prompt"}]},
        "channel_versions": {"messages": 1},
    }

    class _BlockingTitleCheckpointer:
        def __init__(self) -> None:
            self.latest_checkpoint = copy.deepcopy(initial_checkpoint)
            self.latest_metadata = {"source": "loop", "step": 1}
            self.title_write_started = asyncio.Event()
            self.release_title_write = asyncio.Event()

        async def aget_tuple(self, config):
            del config
            return _FakeCheckpointTuple(
                checkpoint=copy.deepcopy(self.latest_checkpoint),
                metadata=dict(self.latest_metadata),
                config={
                    "configurable": {
                        "thread_id": "thread-1",
                        "checkpoint_ns": "",
                        "checkpoint_id": self.latest_checkpoint["id"],
                    }
                },
            )

        async def aput(self, config, checkpoint, metadata, new_versions):
            del config, new_versions
            self.title_write_started.set()
            await self.release_title_write.wait()
            self.latest_checkpoint = copy.deepcopy(checkpoint)
            self.latest_metadata = dict(metadata)
            return {
                "configurable": {
                    "thread_id": "thread-1",
                    "checkpoint_ns": "",
                    "checkpoint_id": checkpoint["id"],
                }
            }

        def get_next_version(self, current, _channel):
            return (current or 0) + 1

    run_manager = RunManager()
    record = await run_manager.create("thread-1")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    checkpointer = _BlockingTitleCheckpointer()

    class _AbortingAgent:
        metadata = {"model_name": "fake-test-model"}
        checkpointer: Any | None = None
        store: Any | None = None
        interrupt_before_nodes = None
        interrupt_after_nodes = None

        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, config, stream_mode, subgraphs
            record.abort_event.set()
            if False:
                yield  # pragma: no cover

    def factory(*, config):
        del config
        return _AbortingAgent()

    task = asyncio.create_task(
        run_agent(
            bridge,
            run_manager,
            record,
            ctx=RunContext(checkpointer=checkpointer),
            agent_factory=factory,
            graph_input={"messages": [{"role": "user", "content": "old prompt"}]},
            config={},
        )
    )
    record.task = task

    try:
        await asyncio.wait_for(checkpointer.title_write_started.wait(), timeout=1.0)
        with pytest.raises(ConflictError, match="active run"):
            await run_manager.create_or_reject("thread-1", multitask_strategy="reject")
    finally:
        checkpointer.release_title_write.set()
        await task

    records = await run_manager.list_by_thread("thread-1")
    assert [item.run_id for item in records] == [record.run_id]
    assert checkpointer.latest_checkpoint["channel_values"]["messages"] == [{"type": "human", "content": "old prompt"}]


@pytest.mark.anyio
async def test_finalizing_run_only_blocks_reject_strategy():
    """A finalizing run must not break interrupt/rollback superseding semantics."""

    async def _seed_finalizing_run():
        run_manager = RunManager()
        record = await run_manager.create("thread-1")
        release_cleanup = asyncio.Event()
        cleanup_cancelled = asyncio.Event()

        async def _cleanup_task():
            try:
                await release_cleanup.wait()
            except asyncio.CancelledError:
                cleanup_cancelled.set()
                raise

        task = asyncio.create_task(_cleanup_task())
        record.task = task
        await run_manager.set_status(record.run_id, RunStatus.interrupted)
        await run_manager.set_finalizing(record.run_id, True)
        return run_manager, record, task, release_cleanup, cleanup_cancelled

    for strategy in ("interrupt", "rollback"):
        run_manager, record, task, release_cleanup, cleanup_cancelled = await _seed_finalizing_run()
        try:
            replacement = await run_manager.create_or_reject("thread-1", multitask_strategy=strategy)
            await asyncio.sleep(0)

            assert replacement.run_id != record.run_id
            assert record.status == RunStatus.interrupted
            assert record.finalizing is True
            assert not cleanup_cancelled.is_set()
            assert not task.done()
        finally:
            release_cleanup.set()
            with suppress(asyncio.CancelledError):
                await task

    run_manager, _record, task, release_cleanup, _cleanup_cancelled = await _seed_finalizing_run()
    try:
        with pytest.raises(ConflictError, match="active run"):
            await run_manager.create_or_reject("thread-1", multitask_strategy="reject")
    finally:
        release_cleanup.set()
        with suppress(asyncio.CancelledError):
            await task


@pytest.mark.anyio
async def test_admitted_pending_replacement_does_not_steal_interrupted_title_recovery(monkeypatch):
    """The old run must still write the fallback title before releasing a serialized replacement."""
    from deerflow.agents.middlewares.title_middleware import TitleMiddleware

    monkeypatch.setattr(
        TitleMiddleware,
        "_generate_title_result",
        lambda self, state, allow_partial_exchange=False: {"title": "Old Prompt"},
    )

    initial_checkpoint = {
        "id": "ckpt-old",
        "ts": "2026-06-29T00:00:00Z",
        "channel_values": {"messages": [{"type": "human", "content": "Old prompt"}]},
        "channel_versions": {"messages": 1},
    }

    run_manager = RunManager()
    old_record = await run_manager.create("thread-1")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    checkpointer = _TitleCheckpointer(
        tuple_value=_FakeCheckpointTuple(
            checkpoint=initial_checkpoint,
            metadata={"source": "loop", "step": 1},
            config={
                "configurable": {
                    "thread_id": "thread-1",
                    "checkpoint_ns": "",
                    "checkpoint_id": "ckpt-old",
                }
            },
        ),
    )

    old_title_gate_entered = asyncio.Event()
    release_old_title_gate = asyncio.Event()
    original_wait_for_prior_finalizing = run_manager.wait_for_prior_finalizing

    async def _wait_for_prior_finalizing(thread_id, run_id, **kwargs):
        if run_id == old_record.run_id and old_record.status == RunStatus.interrupted and old_record.finalizing:
            old_title_gate_entered.set()
            await release_old_title_gate.wait()
        return await original_wait_for_prior_finalizing(thread_id, run_id, **kwargs)

    run_manager.wait_for_prior_finalizing = _wait_for_prior_finalizing  # type: ignore[method-assign]

    class _AbortingAgent:
        metadata = {"model_name": "fake-test-model"}
        checkpointer: Any | None = None
        store: Any | None = None
        interrupt_before_nodes = None
        interrupt_after_nodes = None

        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, config, stream_mode, subgraphs
            old_record.abort_event.set()
            if False:
                yield  # pragma: no cover

    def factory(*, config):
        del config
        return _AbortingAgent()

    old_task = asyncio.create_task(
        run_agent(
            bridge,
            run_manager,
            old_record,
            ctx=RunContext(checkpointer=checkpointer),
            agent_factory=factory,
            graph_input={"messages": [{"role": "user", "content": "Old prompt"}]},
            config={},
        )
    )
    old_record.task = old_task

    try:
        await asyncio.wait_for(old_title_gate_entered.wait(), timeout=1.0)
        replacement_record = await run_manager.create_or_reject("thread-1", multitask_strategy="interrupt")
        assert replacement_record.status == RunStatus.pending

        release_old_title_gate.set()
        await old_task
    finally:
        release_old_title_gate.set()
        if not old_task.done():
            old_task.cancel()
            with suppress(asyncio.CancelledError):
                await old_task

    checkpointer.aput.assert_awaited_once()
    _, written_checkpoint, _, _ = checkpointer.aput.await_args.args
    assert written_checkpoint["channel_values"]["title"] == "Old Prompt"


@pytest.mark.anyio
async def test_interrupted_title_does_not_overwrite_checkpoint_from_admitted_replacement(monkeypatch):
    """A replacement run admitted by multitask interrupt must not lose its newer checkpoint."""
    from deerflow.agents.middlewares.title_middleware import TitleMiddleware

    monkeypatch.setattr(
        TitleMiddleware,
        "_generate_title_result",
        lambda self, state, allow_partial_exchange=False: {"title": "Old prompt"},
    )

    old_checkpoint = {
        "id": "ckpt-old",
        "ts": "2026-06-29T00:00:00Z",
        "channel_values": {"messages": [{"type": "human", "content": "Old prompt"}]},
        "channel_versions": {"messages": 1},
    }
    replacement_messages = [
        {"type": "human", "content": "Old prompt"},
        {"type": "human", "content": "Replacement prompt"},
    ]
    replacement_checkpoint = {
        "id": "ckpt-replacement",
        "ts": "2026-06-29T00:00:01Z",
        "channel_values": {"messages": replacement_messages},
        "channel_versions": {"messages": 2},
    }

    class _ReplacementRaceCheckpointer:
        def __init__(self) -> None:
            self.latest_checkpoint = copy.deepcopy(old_checkpoint)
            self.latest_metadata = {"source": "loop", "step": 1}
            self.title_write_started = asyncio.Event()
            self.replacement_checkpoint_written = asyncio.Event()

        async def aget_tuple(self, config):
            del config
            return _FakeCheckpointTuple(
                checkpoint=copy.deepcopy(self.latest_checkpoint),
                metadata=dict(self.latest_metadata),
                config={
                    "configurable": {
                        "thread_id": "thread-1",
                        "checkpoint_ns": "",
                        "checkpoint_id": self.latest_checkpoint["id"],
                    }
                },
            )

        async def aput(self, config, checkpoint, metadata, new_versions):
            del config, new_versions
            self.title_write_started.set()
            await self.replacement_checkpoint_written.wait()
            self.latest_checkpoint = copy.deepcopy(checkpoint)
            self.latest_metadata = dict(metadata)
            return {
                "configurable": {
                    "thread_id": "thread-1",
                    "checkpoint_ns": "",
                    "checkpoint_id": checkpoint["id"],
                }
            }

        def get_next_version(self, current, _channel):
            return (current or 0) + 1

    run_manager = RunManager()
    old_record = await run_manager.create("thread-1")
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    checkpointer = _ReplacementRaceCheckpointer()
    old_agent_started = asyncio.Event()

    class _BlockingAgent:
        metadata = {"model_name": "fake-test-model"}
        checkpointer: Any | None = None
        store: Any | None = None
        interrupt_before_nodes = None
        interrupt_after_nodes = None

        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, config, stream_mode, subgraphs
            old_agent_started.set()
            while True:
                await asyncio.sleep(0.05)
            if False:
                yield  # pragma: no cover

    def factory(*, config):
        del config
        return _BlockingAgent()

    old_task = asyncio.create_task(
        run_agent(
            bridge,
            run_manager,
            old_record,
            ctx=RunContext(checkpointer=checkpointer),
            agent_factory=factory,
            graph_input={"messages": [{"role": "user", "content": "Old prompt"}]},
            config={},
        )
    )
    old_record.task = old_task

    try:
        await asyncio.wait_for(old_agent_started.wait(), timeout=1.0)
        replacement_record = await run_manager.create_or_reject("thread-1", multitask_strategy="interrupt")
        assert replacement_record.run_id != old_record.run_id
        await run_manager.set_status(replacement_record.run_id, RunStatus.running)

        with suppress(asyncio.TimeoutError):
            await asyncio.wait_for(checkpointer.title_write_started.wait(), timeout=0.25)

        checkpointer.latest_checkpoint = copy.deepcopy(replacement_checkpoint)
        checkpointer.latest_metadata = {"source": "loop", "step": 2}
        checkpointer.replacement_checkpoint_written.set()
        await old_task
    finally:
        checkpointer.replacement_checkpoint_written.set()
        if not old_task.done():
            old_task.cancel()
            with suppress(asyncio.CancelledError):
                await old_task

    assert checkpointer.latest_checkpoint["channel_values"]["messages"] == replacement_messages


@pytest.mark.anyio
async def test_replacement_run_waits_for_prior_finalizing_run():
    """Replacement workers must not enter the graph while an older run is finalizing."""
    run_manager = RunManager()
    old_record = await run_manager.create("thread-1")
    replacement_record = await run_manager.create("thread-1")
    await run_manager.set_finalizing(old_record.run_id, True)

    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    replacement_started = asyncio.Event()

    class _ReplacementAgent:
        metadata = {"model_name": "fake-test-model"}
        checkpointer: Any | None = None
        store: Any | None = None
        interrupt_before_nodes = None
        interrupt_after_nodes = None

        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            del graph_input, config, stream_mode, subgraphs
            replacement_started.set()
            if False:
                yield  # pragma: no cover

    def factory(*, config):
        del config
        return _ReplacementAgent()

    task = asyncio.create_task(
        run_agent(
            bridge,
            run_manager,
            replacement_record,
            ctx=RunContext(checkpointer=None),
            agent_factory=factory,
            graph_input={"messages": [{"role": "user", "content": "Replacement prompt"}]},
            config={},
        )
    )
    replacement_record.task = task

    try:
        await asyncio.sleep(0.1)
        assert not replacement_started.is_set()

        await run_manager.set_finalizing(old_record.run_id, False)
        await asyncio.wait_for(replacement_started.wait(), timeout=1.0)
        await task
    finally:
        await run_manager.set_finalizing(old_record.run_id, False)
        if not task.done():
            task.cancel()
            with suppress(asyncio.CancelledError):
                await task


@pytest.mark.anyio
async def test_ensure_interrupted_title_reloads_latest_checkpoint_before_write():
    """If the checkpoint advances before the title write, preserve the newer messages."""
    from deerflow.config.title_config import TitleConfig

    old_checkpoint = {
        "id": "ckpt-old",
        "ts": "2026-06-29T00:00:00Z",
        "channel_values": {"messages": [{"type": "human", "content": "Old prompt"}]},
        "channel_versions": {"messages": 1},
    }
    new_messages = [{"type": "human", "content": "New prompt"}]
    new_checkpoint = {
        "id": "ckpt-new",
        "ts": "2026-06-29T00:00:01Z",
        "channel_values": {"messages": new_messages},
        "channel_versions": {"messages": 2},
    }

    class _AdvancingTitleCheckpointer:
        def __init__(self) -> None:
            self.read_count = 0
            self.aput = AsyncMock(return_value={})

        async def aget_tuple(self, config):
            del config
            self.read_count += 1
            checkpoint = old_checkpoint if self.read_count == 1 else new_checkpoint
            return _FakeCheckpointTuple(
                checkpoint=copy.deepcopy(checkpoint),
                metadata={"source": "loop", "step": self.read_count},
                config={
                    "configurable": {
                        "thread_id": "thread-1",
                        "checkpoint_ns": "",
                        "checkpoint_id": checkpoint["id"],
                    }
                },
            )

        def get_next_version(self, current, _channel):
            return (current or 0) + 1

    checkpointer = _AdvancingTitleCheckpointer()
    app_config = SimpleNamespace(title=TitleConfig(enabled=True, max_chars=40, max_words=20))

    title = await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=app_config)

    assert title == "New prompt"
    _, written_checkpoint, _, _ = checkpointer.aput.await_args.args
    assert written_checkpoint["channel_values"]["messages"] == new_messages
    assert written_checkpoint["channel_values"]["title"] == "New prompt"


@pytest.mark.anyio
async def test_ensure_interrupted_title_bumps_channel_version_and_declares_it_in_new_versions(monkeypatch):
    """Regression for #3859 review: DB-backed savers (Sqlite/Postgres) strip inline
    ``channel_values`` from ``put`` and only persist blobs for channels listed in
    ``new_versions``. The helper must therefore bump ``channel_versions["title"]``
    and pass ``{"title": next_version}`` so the fallback title actually survives
    a fresh ``aget_tuple`` after the worker's finally hook.
    """
    from deerflow.agents.middlewares.title_middleware import TitleMiddleware

    monkeypatch.setattr(
        TitleMiddleware,
        "_generate_title_result",
        lambda self, state, allow_partial_exchange=False: {"title": "Generated Title"},
    )

    initial_checkpoint = {
        "id": "ckpt-1",
        "ts": "2026-06-29T00:00:00Z",
        "channel_values": {"messages": [{"type": "human", "content": "hi"}]},
        "channel_versions": {"messages": 5},
    }
    checkpointer = _TitleCheckpointer(
        tuple_value=_FakeCheckpointTuple(
            checkpoint=initial_checkpoint,
            metadata={"source": "loop", "step": 7},
        ),
    )

    title = await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None)

    assert title == "Generated Title"
    checkpointer.aput.assert_awaited_once()
    write_config, written_checkpoint, written_metadata, new_versions = checkpointer.aput.await_args.args

    # The title channel must be declared in new_versions — without this, DB
    # savers drop the inline channel_values["title"] from the persisted blob.
    assert new_versions == {"title": 1}
    # Channel versions on the checkpoint itself must also reflect the bump,
    # so a subsequent aget_tuple reconstructs channel_values with the title.
    assert written_checkpoint["channel_versions"]["title"] == 1
    # Pre-existing channel versions must be preserved.
    assert written_checkpoint["channel_versions"]["messages"] == 5
    # The fallback title rides into channel_values for the (legacy / single-table)
    # savers that inline the snapshot.
    assert written_checkpoint["channel_values"]["title"] == "Generated Title"
    assert written_metadata["source"] == "update"
    assert written_metadata["step"] == 8
    assert written_metadata["writes"] == {"runtime_interrupt_title": {"title": "Generated Title"}}
    # The write is parented to the checkpoint it was derived from - without
    # the parent pointer the saver stores a root checkpoint, severing
    # Delta-channel replay ancestry and full-mode history walks.
    assert write_config == {"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "ckpt-1"}}


@pytest.mark.anyio
async def test_ensure_interrupted_title_writes_graph_input_fallback_without_checkpoint(monkeypatch):
    """When no checkpoint exists, graph_input should still seed the fallback title write."""
    from deerflow.agents.middlewares.title_middleware import TitleMiddleware

    captured_state: dict[str, Any] = {}

    def _generate(self, state, allow_partial_exchange=False):
        del self
        captured_state.update(state)
        assert allow_partial_exchange is True
        return {"title": "Graph Input Title"}

    monkeypatch.setattr(TitleMiddleware, "_generate_title_result", _generate)

    checkpointer = _TitleCheckpointer(tuple_value=None)

    title = await _ensure_interrupted_title(
        checkpointer=checkpointer,
        thread_id="thread-1",
        app_config=None,
        graph_input={"messages": [{"role": "user", "content": "Please name this thread"}]},
    )

    assert title == "Graph Input Title"
    assert captured_state["messages"] == [{"role": "user", "content": "Please name this thread"}]
    checkpointer.aput.assert_awaited_once()
    _, written_checkpoint, written_metadata, new_versions = checkpointer.aput.await_args.args
    assert written_checkpoint["channel_values"]["title"] == "Graph Input Title"
    assert written_metadata["writes"] == {"runtime_interrupt_title": {"title": "Graph Input Title"}}
    assert set(new_versions.keys()) == {"title"}


@pytest.mark.anyio
async def test_ensure_interrupted_title_bumps_existing_string_version(monkeypatch):
    """When the checkpointer lacks ``get_next_version`` and the prior title
    version is a string (some savers use UUID-shaped versions), the helper must
    still produce a strictly different value rather than overwriting in place.
    """
    from deerflow.agents.middlewares.title_middleware import TitleMiddleware

    monkeypatch.setattr(
        TitleMiddleware,
        "_generate_title_result",
        lambda self, state, allow_partial_exchange=False: {"title": "T"},
    )

    initial_checkpoint = {
        "id": "ckpt-1",
        "ts": "2026-06-29T00:00:00Z",
        "channel_values": {"messages": [{"type": "human", "content": "hi"}]},
        "channel_versions": {"title": "v3"},
    }

    class _NoGetNextVersion:
        def __init__(self):
            self.aget_tuple = AsyncMock(
                return_value=_FakeCheckpointTuple(
                    checkpoint=initial_checkpoint,
                    metadata={},
                ),
            )
            self.aput = AsyncMock(return_value={})

    checkpointer = _NoGetNextVersion()
    await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None)

    _, written_checkpoint, _, new_versions = checkpointer.aput.await_args.args
    bumped = written_checkpoint["channel_versions"]["title"]
    assert bumped != "v3", "title version must change so DB savers persist the update"
    assert new_versions == {"title": bumped}


@pytest.mark.anyio
async def test_ensure_interrupted_title_skips_when_title_already_set():
    """If the checkpoint already carries a title, no new checkpoint is written."""
    checkpointer = _TitleCheckpointer(
        tuple_value=_FakeCheckpointTuple(
            checkpoint={
                "id": "ckpt-1",
                "channel_values": {"messages": [], "title": "Already there"},
                "channel_versions": {"title": 1},
            },
            metadata={},
        ),
    )

    title = await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None)

    assert title == "Already there"
    checkpointer.aput.assert_not_awaited()


@pytest.mark.anyio
async def test_ensure_interrupted_title_round_trip_with_real_sqlite_checkpointer(tmp_path):
    """Full round-trip against a real ``AsyncSqliteSaver`` on a disk-backed DB.

    Mirrors what Gateway constructs in production via ``make_checkpointer`` when
    ``database.backend == "sqlite"``, then closes and re-opens the saver to
    simulate a fresh connection. The fallback title must survive that boundary —
    this is the scenario the #3874 review flagged as broken before the
    ``new_versions={"title": ...}`` fix.
    """
    from langchain_core.messages import HumanMessage
    from langgraph.checkpoint.base import empty_checkpoint
    from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver

    from deerflow.config.title_config import TitleConfig

    db_path = str(tmp_path / "ckpt.db")
    thread_cfg = {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}}

    # 1. Seed a first-turn checkpoint that has a human message and NO title —
    #    the same shape the agent leaves behind when interrupted mid-stream.
    async with AsyncSqliteSaver.from_conn_string(db_path) as writer:
        await writer.setup()
        ck = empty_checkpoint()
        ck["channel_values"] = {
            "messages": [HumanMessage(content="Why is the sky blue?").model_dump()],
        }
        ck["channel_versions"] = {"messages": 1}
        await writer.aput(thread_cfg, ck, {"source": "loop", "step": 1, "writes": {}}, {"messages": 1})

    # 2. Run the worker helper through a *fresh* saver instance — this is what
    #    the lifespan-owned checkpointer pool does for each request.
    title_config = TitleConfig(enabled=True, max_chars=40, max_words=20)
    app_config = SimpleNamespace(title=title_config)
    async with AsyncSqliteSaver.from_conn_string(db_path) as worker_saver:
        title = await _ensure_interrupted_title(
            checkpointer=worker_saver,
            thread_id="thread-1",
            app_config=app_config,
        )
    assert title, "fallback title must be generated from the seeded user message"

    # 3. Open ANOTHER fresh saver and confirm the title survives — this is the
    #    invariant the #3874 review was guarding: ``new_versions={}`` would
    #    cause DB savers to drop the title blob, so a fresh aget_tuple would
    #    read back without it.
    async with AsyncSqliteSaver.from_conn_string(db_path) as reader:
        tup = await reader.aget_tuple(thread_cfg)
    assert tup is not None
    persisted = tup.checkpoint.get("channel_values", {}).get("title")
    assert persisted == title


# ---------------------------------------------------------------------------
# _bump_channel_version — invariant: the returned version MUST differ from
# the prior value, no matter the checkpointer's versioning scheme.
# ---------------------------------------------------------------------------


class _CheckpointerWithIntVersion:
    """A checkpointer whose ``get_next_version`` increments integers (default LangGraph behavior)."""

    @staticmethod
    def get_next_version(current, _channel):
        return (current or 0) + 1


class _CheckpointerWithBrokenGetNextVersion:
    """A checkpointer whose ``get_next_version`` raises — must fall back, not propagate."""

    @staticmethod
    def get_next_version(current, _channel):
        raise RuntimeError("simulated saver bug")


def test_bump_channel_version_uses_checkpointer_get_next_version_when_available():
    """Happy path — saver's ``get_next_version`` result is preferred over our fallback."""
    assert _bump_channel_version(_CheckpointerWithIntVersion(), 5) == 6


def test_bump_channel_version_falls_back_on_broken_get_next_version():
    """A raising ``get_next_version`` must not propagate; the defensive path bumps from prior.

    Without this, a saver bug would leave ``new_versions={"title": v}`` no-op
    on DB savers — the very class of bug the #3874 review flagged.
    """
    bumped = _bump_channel_version(_CheckpointerWithBrokenGetNextVersion(), 7)
    assert bumped == 8


# ---------------------------------------------------------------------------
# _ensure_interrupted_title — additional defensive boundaries
# ---------------------------------------------------------------------------


@pytest.mark.anyio
async def test_ensure_interrupted_title_handles_none_messages_channel(monkeypatch):
    """A partially-initialized checkpoint with ``messages=None`` must not crash."""
    from deerflow.config.title_config import TitleConfig

    initial_checkpoint = {
        "id": "ckpt-1",
        "channel_values": {"messages": None},
        "channel_versions": {},
    }
    checkpointer = _TitleCheckpointer(
        tuple_value=_FakeCheckpointTuple(checkpoint=initial_checkpoint, metadata={}),
    )
    app_config = SimpleNamespace(title=TitleConfig(enabled=True, max_chars=40, max_words=20))

    assert await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=app_config) is None
    checkpointer.aput.assert_not_awaited()


@pytest.mark.anyio
async def test_ensure_interrupted_title_propagates_aput_error_to_caller(monkeypatch):
    """Exceptions from ``aput`` propagate — the caller (worker.run_agent finally block) is responsible for swallowing them.

    This test pins the contract: the helper itself does NOT silently eat saver errors,
    so a structural saver regression remains visible in the logs at the call site.
    """
    from deerflow.agents.middlewares.title_middleware import TitleMiddleware

    monkeypatch.setattr(
        TitleMiddleware,
        "_generate_title_result",
        lambda self, state, allow_partial_exchange=False: {"title": "Generated"},
    )

    initial_checkpoint = {
        "id": "ckpt-1",
        "channel_values": {"messages": [{"type": "human", "content": "hi"}]},
        "channel_versions": {"messages": 1},
    }
    checkpointer = _TitleCheckpointer(
        tuple_value=_FakeCheckpointTuple(checkpoint=initial_checkpoint, metadata={}),
    )
    checkpointer.aput.side_effect = RuntimeError("simulated DB write failure")

    with pytest.raises(RuntimeError, match="simulated DB write failure"):
        await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None)


@pytest.mark.anyio
async def test_ensure_interrupted_title_idempotent_across_repeated_calls(monkeypatch):
    """Second invocation against the now-titled checkpoint must not re-write.

    Regression anchor for the case where a brittle helper might re-trigger
    on subsequent finally-hook runs (e.g. retries) and rewrite the title.
    """
    from langgraph.checkpoint.memory import InMemorySaver

    from deerflow.agents.middlewares.title_middleware import TitleMiddleware

    monkeypatch.setattr(
        TitleMiddleware,
        "_generate_title_result",
        lambda self, state, allow_partial_exchange=False: {"title": "First Title"},
    )

    checkpointer = InMemorySaver()
    cfg = {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}}
    ck = empty_checkpoint()
    ck["channel_values"] = {"messages": [{"type": "human", "content": "hi"}]}
    ck["channel_versions"] = {"messages": 1}
    await checkpointer.aput(cfg, ck, {"source": "loop", "step": 1, "writes": {}}, {"messages": 1})

    first = await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None)
    assert first == "First Title"

    # Second call: title is now present, so the helper short-circuits without
    # rewriting — even if the middleware were to suggest a different title.
    monkeypatch.setattr(
        TitleMiddleware,
        "_generate_title_result",
        lambda self, state, allow_partial_exchange=False: {"title": "Different Title"},
    )
    second = await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None)
    assert second == "First Title"

    tup = await checkpointer.aget_tuple(cfg)
    assert tup.checkpoint["channel_values"]["title"] == "First Title"


@pytest.mark.anyio
async def test_ensure_interrupted_title_preserves_non_title_channel_versions(monkeypatch):
    """Bumping ``channel_versions["title"]`` must not modify other channels' versions.

    Regression anchor: an earlier draft built ``new_versions`` from
    ``dict(channel_versions)`` and would have erroneously declared every
    channel as "needs new blob" on DB savers.
    """
    from deerflow.agents.middlewares.title_middleware import TitleMiddleware

    monkeypatch.setattr(
        TitleMiddleware,
        "_generate_title_result",
        lambda self, state, allow_partial_exchange=False: {"title": "Generated"},
    )

    initial_checkpoint = {
        "id": "ckpt-1",
        "channel_values": {
            "messages": [{"type": "human", "content": "hi"}],
            "artifacts": [],
            "todos": None,
        },
        "channel_versions": {"messages": 5, "artifacts": 3, "todos": 1},
    }
    checkpointer = _TitleCheckpointer(
        tuple_value=_FakeCheckpointTuple(checkpoint=initial_checkpoint, metadata={}),
    )

    await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None)

    _, written_checkpoint, _, new_versions = checkpointer.aput.await_args.args
    # Only the title channel is declared in new_versions.
    assert set(new_versions.keys()) == {"title"}
    # Other channel versions are preserved verbatim on the written checkpoint.
    assert written_checkpoint["channel_versions"]["messages"] == 5
    assert written_checkpoint["channel_versions"]["artifacts"] == 3
    assert written_checkpoint["channel_versions"]["todos"] == 1


@pytest.mark.anyio
async def test_worker_finally_block_swallows_helper_exceptions(monkeypatch):
    """The worker's interrupted-title hook must remain non-fatal — any exception
    from the helper (DB saver bug, middleware bug, etc.) must not propagate past
    the run boundary or prevent the subsequent threads_meta sync block from
    running. This pins the integration of helper + finally try/except, not just
    the helper itself.
    """
    import deerflow.runtime.runs.worker as worker_module

    helper_called = asyncio.Event()

    async def _boom(*_args, **_kwargs):
        helper_called.set()
        raise RuntimeError("forced helper failure")

    monkeypatch.setattr(worker_module, "_ensure_interrupted_title", _boom)

    run_manager = RunManager()
    record = await run_manager.create("thread-1")

    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )

    class _MinimalCheckpointer:
        async def aget_tuple(self, config):
            return None

        async def aput(self, *args, **kwargs):
            return {}

    captured_status: dict[str, Any] = {}

    class _ThreadStore:
        async def update_display_name(self, thread_id, title):
            captured_status["display_name"] = (thread_id, title)

        async def update_status(self, thread_id, status):
            captured_status["status"] = (thread_id, status)

    class _AbortingAgent:
        def __init__(self) -> None:
            self.metadata = {"model_name": "fake-test-model"}
            self.checkpointer: Any | None = None
            self.store: Any | None = None
            self.interrupt_before_nodes = None
            self.interrupt_after_nodes = None

        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
            # Abort immediately so the run lands in the interrupted branch.
            record.abort_event.set()
            if False:
                yield  # pragma: no cover — make this an async generator
            return

    def factory(*, config):
        del config
        return _AbortingAgent()

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(checkpointer=_MinimalCheckpointer(), thread_store=_ThreadStore()),
        agent_factory=factory,
        graph_input={"messages": []},
        config={},
    )

    # The helper raised, but the run still reaches the threads_meta status sync
    # and ``publish_end`` — i.e. the SSE stream is closed cleanly and the row
    # reflects the run outcome.
    assert helper_called.is_set()
    assert captured_status.get("status") == ("thread-1", "interrupted")
    bridge.publish_end.assert_awaited_once_with(record.run_id)


@pytest.mark.anyio
async def test_worker_skips_execution_and_finalization_after_ownership_loss():
    """A fenced worker closes its stream without starting or finalizing work."""
    run_manager = RunManager()
    record = await run_manager.create("thread-lease-lost")
    record.ownership_lost = True
    record.abort_event.set()
    record.status = RunStatus.error

    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )
    thread_store = SimpleNamespace(
        update_display_name=AsyncMock(),
        update_status=AsyncMock(),
    )
    on_run_completed = AsyncMock()
    agent_factory = MagicMock(side_effect=AssertionError("fenced worker started the agent"))

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(
            checkpointer=None,
            thread_store=thread_store,
            on_run_completed=on_run_completed,
        ),
        agent_factory=agent_factory,
        graph_input={"messages": []},
        config={},
    )

    agent_factory.assert_not_called()
    thread_store.update_display_name.assert_not_awaited()
    thread_store.update_status.assert_not_awaited()
    on_run_completed.assert_not_awaited()
    bridge.publish_end.assert_awaited_once_with(record.run_id)


@pytest.mark.anyio
async def test_worker_discards_buffered_journal_events_after_ownership_loss(monkeypatch):
    """A fenced worker detaches its journal without appending buffered events."""

    class TrackingRunEventStore(MemoryRunEventStore):
        def __init__(self) -> None:
            super().__init__()
            self.put_batch_calls = 0

        async def put_batch(self, events):
            self.put_batch_calls += 1
            return await super().put_batch(events)

    journals: list[RunJournal] = []

    class BufferedRunJournal(RunJournal):
        def __init__(self, *args, **kwargs) -> None:
            super().__init__(*args, **kwargs)
            self.record_middleware("buffered", name="test", hook="after", action="record", changes={})
            journals.append(self)

    monkeypatch.setattr("deerflow.runtime.journal.RunJournal", BufferedRunJournal)

    event_store = TrackingRunEventStore()
    run_manager = RunManager()
    record = await run_manager.create("thread-lease-lost-buffered")
    record.ownership_lost = True
    record.abort_event.set()
    record.status = RunStatus.error
    bridge = SimpleNamespace(
        publish=AsyncMock(),
        publish_end=AsyncMock(),
        cleanup=AsyncMock(),
    )

    await run_agent(
        bridge,
        run_manager,
        record,
        ctx=RunContext(checkpointer=None, event_store=event_store),
        agent_factory=MagicMock(side_effect=AssertionError("fenced worker started the agent")),
        graph_input={"messages": []},
        config={},
    )

    assert event_store.put_batch_calls == 0
    assert len(journals) == 1
    assert journals[0]._closed is True
    assert journals[0]._store is None
    assert journals[0]._buffer == []