deer-flow/backend/tests/test_dynamic_context_middleware.py

"""Tests for DynamicContextMiddleware.

Verifies that memory and current date are injected as a <system-reminder> into
the first HumanMessage exactly once per session (frozen-snapshot pattern).
"""

import asyncio
import hashlib
from types import SimpleNamespace
from unittest import mock

import pytest
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage, SystemMessage
from langgraph.graph.message import add_messages

from deerflow.agents.middlewares.dynamic_context_middleware import (
    _DYNAMIC_CONTEXT_REMINDER_KEY,
    DynamicContextMiddleware,
)
from deerflow.runtime.context_keys import CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY

_SYSTEM_REMINDER_TAG = "<system-reminder>"


def _make_middleware(**kwargs) -> DynamicContextMiddleware:
    return DynamicContextMiddleware(**kwargs)


def _fake_runtime(journal=None, *, pre_existing_message_ids=(), user_id: str | None = None):
    context = {"__run_journal": journal} if journal is not None else {}
    if pre_existing_message_ids:
        context[CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] = frozenset(pre_existing_message_ids)
    if user_id is not None:
        context["user_id"] = user_id
    return SimpleNamespace(context=context)


def _reminder_msg(content: str, msg_id: str) -> HumanMessage:
    """Build a pre-PR HumanMessage reminder — simulates historical checkpoints.

    Uses HumanMessage (DEPRECATED format) to exercise the backward-compat
    path in ``is_dynamic_context_reminder``.  New reminders are SystemMessage.
    """
    return HumanMessage(
        content=content,
        id=msg_id,
        additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
    )


class _FakeRequest:
    """Minimal ModelRequest stand-in for the wrap_model_call hooks."""

    def __init__(self, messages, runtime):
        self.messages = list(messages)
        self.runtime = runtime

    def override(self, **kwargs):
        return _FakeRequest(kwargs.get("messages", self.messages), self.runtime)


def _drive_first_model_call(mw: DynamicContextMiddleware, messages, runtime):
    """Run the first model-request assembly (wrap_model_call) and capture the request."""
    captured: dict = {}

    def _handler(request):
        captured["messages"] = list(request.messages)
        return "response"

    result = mw.wrap_model_call(_FakeRequest(messages, runtime), _handler)
    assert result == "response"
    return captured["messages"]


def _date_reminder_msg(date_str: str, msg_id: str) -> SystemMessage:
    """Build a persisted date reminder in the current production shape.

    A date SystemMessage whose ``reminder_date`` additional_kwargs carries the
    authoritative date — what ``DynamicContextMiddleware`` now writes to state.
    """
    content = f"<system-reminder>\n<current_date>{date_str}</current_date>\n</system-reminder>"
    return SystemMessage(
        content=content,
        id=msg_id,
        additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True, "reminder_date": date_str},
    )


# ---------------------------------------------------------------------------
# Basic injection
# ---------------------------------------------------------------------------


def test_injects_system_reminder_into_first_human_message():
    mw = _make_middleware()
    state = {"messages": [HumanMessage(content="Hello", id="msg-1")]}

    with mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=""), mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, _fake_runtime())

    assert result is not None
    updated_msgs = result["messages"]
    assert len(updated_msgs) == 2

    reminder_msg = updated_msgs[0]
    assert isinstance(reminder_msg, SystemMessage)
    assert reminder_msg.id == "msg-1"  # takes the original ID (position swap)
    assert reminder_msg.additional_kwargs.get(_DYNAMIC_CONTEXT_REMINDER_KEY) is True
    assert _SYSTEM_REMINDER_TAG in reminder_msg.content
    assert "<current_date>2026-05-08, Friday</current_date>" in reminder_msg.content
    assert "Hello" not in reminder_msg.content  # reminder only — no user text

    user_msg = updated_msgs[1]
    assert isinstance(user_msg, HumanMessage)
    assert user_msg.id == "msg-1__user"  # derived ID
    assert user_msg.content == "Hello"


def test_memory_included_when_present():
    mw = _make_middleware()
    state = {"messages": [HumanMessage(content="Hi", id="msg-1")]}

    with (
        mock.patch(
            "deerflow.agents.lead_agent.prompt._get_memory_context",
            return_value="<memory>\nUser prefers Python.\n</memory>",
        ),
        mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt,
    ):
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, _fake_runtime())

    # Memory is a separate HumanMessage — not merged into SystemMessage (OWASP LLM01)
    msgs = result["messages"]
    assert len(msgs) == 3  # date SystemMessage + memory HumanMessage + user HumanMessage

    assert isinstance(msgs[0], SystemMessage)
    assert "<current_date>2026-05-08, Friday</current_date>" in msgs[0].content
    assert "User prefers Python." not in msgs[0].content  # memory NOT in system role

    assert isinstance(msgs[1], HumanMessage)
    assert "User prefers Python." in msgs[1].content

    assert msgs[2].content == "Hi"


def test_memory_opt_out_keeps_date_reminder_without_reading_memory():
    mw = _make_middleware(memory_enabled=False)
    state = {"messages": [HumanMessage(content="Hi", id="msg-1")]}

    with (
        mock.patch(
            "deerflow.agents.lead_agent.prompt._get_memory_context",
            side_effect=AssertionError("disabled custom agent must not read memory"),
        ),
        mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt,
    ):
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, _fake_runtime())

    assert result is not None
    assert len(result["messages"]) == 2
    assert isinstance(result["messages"][0], SystemMessage)
    assert "<current_date>2026-05-08, Friday</current_date>" in result["messages"][0].content
    assert result["messages"][1].content == "Hi"


def test_memory_opt_out_removes_frozen_checkpoint_memory_but_keeps_date():
    """Changing an existing agent to stateless must take effect immediately.

    The date reminder and original user turn remain valid framework context, but
    a server-tagged ``__memory`` message from an earlier checkpoint must not be
    sent to the model after the agent opts out. An untagged client message with a
    similar ID is not middleware-owned and must be preserved.
    """
    date = "2026-05-08, Friday"
    frozen_memory = HumanMessage(
        content="<memory>User prefers Python.</memory>",
        id="msg-1__memory",
        additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
    )
    client_message = HumanMessage(content="client data", id="client__memory")
    state = {
        "messages": [
            _date_reminder_msg(date, "msg-1"),
            frozen_memory,
            HumanMessage(content="Hello", id="msg-1__user"),
            client_message,
            HumanMessage(content="Continue", id="msg-2"),
        ]
    }

    with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value.strftime.return_value = date
        result = _make_middleware(memory_enabled=False).before_agent(state, _fake_runtime())

    assert result is not None
    assert result["messages"] == [RemoveMessage(id="msg-1__memory")]

    updated = add_messages(state["messages"], result["messages"])
    assert "msg-1__memory" not in {message.id for message in updated}
    assert next(message for message in updated if message.id == "msg-1").content.endswith(f"<current_date>{date}</current_date>\n</system-reminder>")
    assert next(message for message in updated if message.id == "client__memory").content == "client data"


async def _wait_forever(*_args, **_kwargs):
    await asyncio.Event().wait()


@pytest.mark.asyncio
async def test_memory_opt_out_async_timeout_still_removes_frozen_memory(monkeypatch):
    """A date/injection timeout cannot delay an already-active opt-out."""
    date = "2026-05-08, Friday"
    state = {
        "messages": [
            _date_reminder_msg(date, "msg-1"),
            HumanMessage(
                content="<memory>User prefers Python.</memory>",
                id="msg-1__memory",
                additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
            ),
            HumanMessage(content="Continue", id="msg-2"),
        ]
    }
    monkeypatch.setattr(
        "deerflow.agents.middlewares.dynamic_context_middleware.asyncio.to_thread",
        _wait_forever,
    )
    monkeypatch.setattr(
        "deerflow.agents.middlewares.dynamic_context_middleware._INJECT_TIMEOUT_SECONDS",
        0.01,
    )

    result = await _make_middleware(memory_enabled=False).abefore_agent(state, _fake_runtime())

    assert result == {"messages": [RemoveMessage(id="msg-1__memory")]}


def test_memory_lookup_uses_runtime_user_id():
    mw = _make_middleware()
    state = {"messages": [HumanMessage(content="Hi", id="msg-1")]}

    with (
        mock.patch(
            "deerflow.agents.lead_agent.prompt._get_memory_context",
            return_value="",
        ) as get_memory_context,
        mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt,
    ):
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        mw.before_agent(state, _fake_runtime(user_id="runtime-user"))

    get_memory_context.assert_called_once_with(
        None,
        app_config=None,
        user_id="runtime-user",
        query="Hi",
    )


def test_first_run_records_exact_effective_memory():
    """The run's context event fires at the first model-request assembly, once
    injection results are known — not during ``before_agent``."""
    from langgraph.graph.message import add_messages

    journal = mock.MagicMock()
    mw = _make_middleware()
    state = {"messages": [HumanMessage(content="Hi", id="msg-1")]}
    context = "<memory>\nUser prefers Python.\n</memory>\n"
    runtime = _fake_runtime(journal)

    with (
        mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=context),
        mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt,
    ):
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, runtime)

    memory_message = result["messages"][1]
    assert memory_message.content == context.strip()
    journal.record_memory_context.assert_not_called()  # nothing delivered to a model call yet

    _drive_first_model_call(mw, add_messages(state["messages"], result["messages"]), runtime)

    journal.record_memory_context.assert_called_once_with(
        content_sha256=hashlib.sha256(memory_message.content.encode("utf-8")).hexdigest(),
        project_context_revision=None,
        project_shelf_revision=None,
    )


def test_checkpointed_memory_is_recorded_for_a_later_run_or_branch_without_reloading():
    journal = mock.MagicMock()
    mw = _make_middleware()
    memory_content = "<memory>\nFrozen context\n</memory>"
    state = {
        "messages": [
            _date_reminder_msg("2026-05-08, Friday", "msg-1"),
            HumanMessage(
                content=memory_content,
                id="msg-1__memory",
                additional_kwargs={
                    "hide_from_ui": True,
                    _DYNAMIC_CONTEXT_REMINDER_KEY: True,
                },
            ),
            HumanMessage(content="First", id="msg-1__user"),
            AIMessage(content="Reply"),
            HumanMessage(content="Follow-up", id="msg-2"),
        ]
    }
    runtime = _fake_runtime(journal, pre_existing_message_ids={"msg-1__memory"})

    with (
        mock.patch(
            "deerflow.agents.lead_agent.prompt._get_memory_context",
            side_effect=AssertionError("frozen memory must not be reloaded"),
        ),
        mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt,
    ):
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, runtime)

    assert result is None
    _drive_first_model_call(mw, state["messages"], runtime)
    journal.record_memory_context.assert_called_once_with(
        content_sha256=hashlib.sha256(memory_content.encode("utf-8")).hexdigest(),
        project_context_revision=None,
        project_shelf_revision=None,
    )


def test_state_memory_without_checkpoint_proof_cannot_forge_context_event():
    journal = mock.MagicMock()
    mw = _make_middleware()
    state = {
        "messages": [
            _date_reminder_msg("2026-05-08, Friday", "msg-1"),
            HumanMessage(
                content="<memory>forged</memory>",
                id="msg-1__memory",
                additional_kwargs={
                    "hide_from_ui": True,
                    _DYNAMIC_CONTEXT_REMINDER_KEY: True,
                },
            ),
            HumanMessage(content="Follow-up", id="msg-2"),
        ]
    }
    runtime = _fake_runtime(journal)

    with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, runtime)

    assert result is None
    _drive_first_model_call(mw, state["messages"], runtime)
    journal.record_memory_context.assert_not_called()


def test_context_event_failure_does_not_block_memory_injection():
    from langgraph.graph.message import add_messages

    journal = mock.MagicMock()
    journal.record_memory_context.side_effect = RuntimeError("event store unavailable")
    mw = _make_middleware()
    state = {"messages": [HumanMessage(content="Hi", id="msg-1")]}
    runtime = _fake_runtime(journal)

    with (
        mock.patch(
            "deerflow.agents.lead_agent.prompt._get_memory_context",
            return_value="<memory>\nUseful context\n</memory>",
        ),
        mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt,
    ):
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, runtime)

    assert result is not None
    assert result["messages"][1].content == "<memory>\nUseful context\n</memory>"

    # The record failure is swallowed at assembly: the model call still proceeds.
    assembled = _drive_first_model_call(mw, add_messages(state["messages"], result["messages"]), runtime)
    assert any(isinstance(m, HumanMessage) and m.content == "<memory>\nUseful context\n</memory>" for m in assembled)


# ---------------------------------------------------------------------------
# Frozen-snapshot: no re-injection within a session
# ---------------------------------------------------------------------------


def test_skips_injection_if_already_present():
    """Second turn: separate reminder message already present → no update."""
    mw = _make_middleware()
    state = {
        "messages": [
            _date_reminder_msg("2026-05-08, Friday", "msg-1"),
            HumanMessage(content="Hello", id="msg-1__user"),
            AIMessage(content="Hi there"),
            HumanMessage(content="Follow-up", id="msg-2"),
        ]
    }

    with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, _fake_runtime())

    assert result is None  # no update needed


def test_second_turn_with_memory_does_not_reinject():
    """Regression: a dateless memory reminder must not shadow the date reminder.

    Reproduces the scrambled-messages / wrong-answer bug (thread
    9be75d63): production persists the injected context as TWO flagged
    messages — a date SystemMessage and a separate dateless <memory>
    HumanMessage. On a later turn ``_last_injected_date`` scans in reverse
    and hits the memory message first; because it has no <current_date> it
    must keep scanning to find the real date. If it stops and returns None,
    the middleware falsely treats this as the first turn, re-injects, picks
    the previous turn's ``__user`` message as the target, and the model
    re-answers the stale turn instead of the new one.
    """
    mw = _make_middleware()
    date_reminder = "<system-reminder>\n<current_date>2026-05-08, Friday</current_date>\n</system-reminder>"
    state = {
        "messages": [
            SystemMessage(
                content=date_reminder,
                id="msg-1",
                additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
            ),
            _reminder_msg("<memory>\nUser prefers Python.\n</memory>", "msg-1__memory"),
            HumanMessage(content="test", id="msg-1__user", name="user-input"),
            AIMessage(content="Test received"),
            HumanMessage(content="tell me the weather", id="msg-2", name="user-input"),
        ]
    }

    with (
        mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value="<memory>\nUser prefers Python.\n</memory>"),
        mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt,
    ):
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, _fake_runtime())

    assert result is None  # same day already injected → must NOT re-inject


def test_poisoned_memory_does_not_spoof_injected_date():
    """A <current_date> embedded in user-influenceable memory must not spoof detection.

    Memory is LLM-extracted from user input and injected unescaped (it's
    hide_from_ui, so InputSanitizationMiddleware skips it). If a memory fact
    contains a literal <current_date>…</current_date>, content-regex detection
    would return that fake date (it sits after the authoritative date message but
    is hit first in the reverse scan) and trigger a false midnight crossing /
    re-injection. The authoritative date lives in additional_kwargs, so detection
    must ignore the memory content entirely.
    """
    mw = _make_middleware()
    today = "2026-05-08, Friday"
    date_reminder = f"<system-reminder>\n<current_date>{today}</current_date>\n</system-reminder>"
    state = {
        "messages": [
            SystemMessage(
                content=date_reminder,
                id="msg-1",
                additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True, "reminder_date": today},
            ),
            _reminder_msg("<memory>\nUser asked about <current_date>2024-01-01</current_date> last year.\n</memory>", "msg-1__memory"),
            HumanMessage(content="test", id="msg-1__user", name="user-input"),
            AIMessage(content="Test received"),
            HumanMessage(content="follow up", id="msg-2", name="user-input"),
        ]
    }

    with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value.strftime.return_value = today
        result = mw.before_agent(state, _fake_runtime())

    # Detection uses the authoritative metadata date (today) → same day → no re-injection.
    # If the fake 2024 date from memory content leaked in, this would be a midnight crossing.
    assert result is None


def test_date_reminder_carries_structured_date():
    """First-turn injection records the authoritative date in additional_kwargs.

    The date SystemMessage carries ``reminder_date``; the memory HumanMessage
    deliberately does not (it is dateless and must never spoof detection).
    """
    mw = _make_middleware()
    state = {"messages": [HumanMessage(content="Hi", id="msg-1")]}

    with (
        mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value="<memory>\nUser prefers Python.\n</memory>"),
        mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt,
    ):
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, _fake_runtime())

    msgs = result["messages"]
    assert isinstance(msgs[0], SystemMessage)
    assert msgs[0].additional_kwargs.get("reminder_date") == "2026-05-08, Friday"
    # Memory HumanMessage must not carry the authoritative date
    assert isinstance(msgs[1], HumanMessage)
    assert "reminder_date" not in msgs[1].additional_kwargs


def test_legacy_systemmessage_reminder_without_key_detected():
    """Backward-compat: pre-reminder_date checkpoints kept the date in content only.

    A date SystemMessage with the date in content but no ``reminder_date`` key
    must still be detected (via the SystemMessage-scoped content fallback) so
    in-flight conversations from before the upgrade do not re-inject.
    """
    mw = _make_middleware()
    state = {
        "messages": [
            SystemMessage(
                content="<system-reminder>\n<current_date>2026-05-08, Friday</current_date>\n</system-reminder>",
                id="msg-1",
                additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},  # no reminder_date
            ),
            HumanMessage(content="Hello", id="msg-1__user"),
            AIMessage(content="Hi there"),
            HumanMessage(content="Follow-up", id="msg-2"),
        ]
    }

    with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, _fake_runtime())

    assert result is None  # same day detected from content → no re-injection


def test_first_turn_fallback_targets_the_latest_user_message():
    """Fallback injection (history without any reminder) attaches to the LAST user target.

    A history that reaches the ``last_date is None`` branch while already
    holding multiple turns only exists because an earlier injection never
    happened (e.g. the async degraded path skipped it) — the genuinely first
    turn has exactly one message.  The reminder must therefore attach to the
    latest user message: the ID-swap's ``{id}__user`` copy is appended by
    ``add_messages``, so attaching to an earlier message would move that stale
    prompt to the tail, ahead of the current question, and the model would
    answer it as the current turn.
    """
    mw = _make_middleware()
    state = {
        "messages": [
            HumanMessage(content="First", id="msg-1"),
            AIMessage(content="Reply"),
            HumanMessage(content="Second", id="msg-2"),
        ]
    }

    with mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=""), mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, _fake_runtime())

    assert result is not None
    msgs = result["messages"]
    # Only the two injected messages are returned (reminder + latest user query)
    assert len(msgs) == 2
    assert msgs[0].id == "msg-2"  # reminder takes the latest message's ID
    assert msgs[0].additional_kwargs.get(_DYNAMIC_CONTEXT_REMINDER_KEY) is True
    assert _SYSTEM_REMINDER_TAG in msgs[0].content
    assert msgs[1].id == "msg-2__user"  # latest user content with derived ID
    assert msgs[1].content == "Second"
    # "First" (msg-1) is not in the returned update — it is left unchanged
    assert all(m.id != "msg-1" for m in msgs)


# ---------------------------------------------------------------------------
# Edge cases
# ---------------------------------------------------------------------------


def test_no_messages_returns_none():
    mw = _make_middleware()
    result = mw.before_agent({"messages": []}, _fake_runtime())
    assert result is None


def test_no_human_message_returns_none():
    mw = _make_middleware()
    state = {"messages": [AIMessage(content="assistant only")]}
    with mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=""):
        result = mw.before_agent(state, _fake_runtime())
    assert result is None


def test_list_content_message_handled_as_separate_reminder():
    """List-content (e.g. multi-modal) messages remain intact; reminder is a separate message."""
    mw = _make_middleware()
    original_content = [{"type": "text", "text": "Hello"}]
    state = {"messages": [HumanMessage(content=original_content, id="msg-1")]}

    with mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=""), mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, _fake_runtime())

    assert result is not None
    msgs = result["messages"]
    assert len(msgs) == 2
    # Reminder is a plain string message with the flag set
    assert isinstance(msgs[0].content, str)
    assert msgs[0].additional_kwargs.get(_DYNAMIC_CONTEXT_REMINDER_KEY) is True
    assert _SYSTEM_REMINDER_TAG in msgs[0].content
    # Original list-content message is untouched
    assert msgs[1].content == original_content


def test_reminder_uses_original_id_user_message_uses_derived_id():
    """Reminder takes original ID (position swap); user message gets {id}__user."""
    mw = _make_middleware()
    original_id = "original-id-abc"
    state = {"messages": [HumanMessage(content="Hello", id=original_id)]}

    with mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=""), mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, _fake_runtime())

    assert result["messages"][0].id == original_id
    assert result["messages"][1].id == f"{original_id}__user"


def test_user_message_copy_preserves_response_metadata():
    """Re-id'd user message keeps fields a hand-built HumanMessage would drop.

    API callers can attach response_metadata to the incoming turn; rebuilding
    the message field-by-field silently discarded it, so the copy is made with
    ``model_copy`` (same pattern as the pii_redaction / uploads middlewares)
    while still deriving the ``{id}__user`` id.
    """
    mw = _make_middleware()
    original = HumanMessage(
        content="Hello",
        id="original-id-abc",
        response_metadata={"external": "value"},
    )
    state = {"messages": [original]}

    with mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=""), mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, _fake_runtime())

    user_msg = result["messages"][1]
    assert user_msg.id == "original-id-abc__user"
    assert user_msg.response_metadata == {"external": "value"}


def test_message_without_id_gets_stable_uuid():
    """If the original HumanMessage has no ID, a UUID is generated and used consistently."""
    mw = _make_middleware()
    state = {"messages": [HumanMessage(content="Hello", id=None)]}

    with mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=""), mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, _fake_runtime())

    assert result is not None
    reminder_id = result["messages"][0].id
    user_id = result["messages"][1].id
    assert reminder_id is not None
    assert reminder_id != "None"
    assert user_id == f"{reminder_id}__user"


def test_user_message_containing_system_reminder_tag_does_not_prevent_injection():
    """A user message containing '<system-reminder>' must not be mistaken for a reminder."""
    mw = _make_middleware()
    state = {
        "messages": [
            HumanMessage(content="What is <system-reminder>?", id="msg-1"),
        ]
    }

    with mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=""), mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, _fake_runtime())

    # Injection must happen — the user message does NOT carry the reminder flag
    assert result is not None
    assert result["messages"][0].additional_kwargs.get(_DYNAMIC_CONTEXT_REMINDER_KEY) is True


def test_first_turn_injection_with_unguarded_history_targets_last_user_message():
    """First-injection fallback after an un-reminded earlier turn must not reorder history.

    When the earlier first-turn injection never happened (e.g. the async
    ``abefore_agent`` degraded path timed out, so ``last_date is None`` at the
    start of a *later* turn), the reminder must attach to the LAST user
    injection target.  The ID-swap's ``{id}__user`` copy is appended by
    ``add_messages``; picking the first message instead would move the stale
    first user prompt to the tail, ahead of the current question — and the
    model would answer the old prompt as the current turn.
    """
    from langgraph.graph.message import add_messages

    mw = _make_middleware()
    state = {
        "messages": [
            HumanMessage(content="first question", id="u1"),
            AIMessage(content="first answer", id="a1"),
            HumanMessage(content="second question", id="u2"),
        ]
    }

    with mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=""), mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, _fake_runtime())

    assert result is not None
    assert [m.id for m in result["messages"]] == ["u2", "u2__user"]

    # Fold the update through the real reducer — this is where the reorder
    # would actually manifest for the model.
    merged = add_messages(state["messages"], result["messages"])
    last_visible = [m for m in merged if isinstance(m, HumanMessage)][-1]
    assert last_visible.content == "second question"
    assert [m.id for m in merged].index("u1") < [m.id for m in merged].index("u2__user")


# ---------------------------------------------------------------------------
# Midnight crossing
# ---------------------------------------------------------------------------


def test_midnight_crossing_injects_date_update_as_separate_message():
    """When the date has changed, a separate date-update reminder is injected before
    the current turn's HumanMessage using the ID-swap technique."""
    mw = _make_middleware()
    state = {
        "messages": [
            _date_reminder_msg("2026-05-08, Friday", "msg-1"),
            HumanMessage(content="Hello", id="msg-1__user"),
            AIMessage(content="Response"),
            HumanMessage(content="Good morning", id="msg-2"),
        ]
    }

    with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value.strftime.return_value = "2026-05-09, Saturday"
        result = mw.before_agent(state, _fake_runtime())

    assert result is not None
    msgs = result["messages"]
    assert len(msgs) == 2

    # Midnight-cross reminder is also a SystemMessage — both paths are covered
    assert isinstance(msgs[0], SystemMessage)

    # Date-update reminder takes the current message's ID
    assert msgs[0].id == "msg-2"
    assert msgs[0].additional_kwargs.get(_DYNAMIC_CONTEXT_REMINDER_KEY) is True
    assert _SYSTEM_REMINDER_TAG in msgs[0].content
    assert "<current_date>2026-05-09, Saturday</current_date>" in msgs[0].content
    assert "Good morning" not in msgs[0].content  # reminder only

    # Original user text appended with derived ID
    assert msgs[1].id == "msg-2__user"
    assert msgs[1].content == "Good morning"


def test_midnight_crossing_id_swap():
    """Date-update reminder uses original ID; user message uses {id}__user."""
    mw = _make_middleware()
    state = {
        "messages": [
            _date_reminder_msg("2026-05-08, Friday", "msg-1"),
            HumanMessage(content="Next day message", id="msg-2"),
        ]
    }

    with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value.strftime.return_value = "2026-05-09, Saturday"
        result = mw.before_agent(state, _fake_runtime())

    assert result["messages"][0].id == "msg-2"
    assert result["messages"][1].id == "msg-2__user"


def test_memory_message_carries_reminder_key_for_title_eligibility():
    """Regression: memory HumanMessage must carry _DYNAMIC_CONTEXT_REMINDER_KEY.

    Without it, title_middleware._is_user_message_for_title counts the memory
    block as a second user message and skips title generation entirely.
    Similarly, summarization_middleware._preserve_required_context
    would not rescue the memory block from summary compression.
    """
    from deerflow.agents.middlewares.dynamic_context_middleware import is_dynamic_context_reminder

    mw = _make_middleware()
    state = {"messages": [HumanMessage(content="Hi", id="msg-1")]}

    with (
        mock.patch(
            "deerflow.agents.lead_agent.prompt._get_memory_context",
            return_value="<memory>\nUser prefers Python.\n</memory>",
        ),
        mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt,
    ):
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result = mw.before_agent(state, _fake_runtime())

    msgs = result["messages"]
    # Memory message must be recognized as a dynamic-context reminder
    memory_msg = msgs[1]
    assert isinstance(memory_msg, HumanMessage)
    assert memory_msg.id == "msg-1__memory"
    assert is_dynamic_context_reminder(memory_msg) is True

    # Only the actual user message is title-eligible
    from deerflow.agents.middlewares.title_middleware import TitleMiddleware

    title_eligible = [m for m in msgs if TitleMiddleware._is_user_message_for_title(m)]
    assert len(title_eligible) == 1
    assert title_eligible[0].content == "Hi"


def test_no_second_midnight_injection_once_date_updated():
    """After a midnight update is persisted, the same-day path skips re-injection."""
    mw = _make_middleware()
    state = {
        "messages": [
            _date_reminder_msg("2026-05-08, Friday", "msg-1"),
            HumanMessage(content="Hello", id="msg-1__user"),
            AIMessage(content="Response"),
            _date_reminder_msg("2026-05-09, Saturday", "msg-2"),
            HumanMessage(content="Good morning", id="msg-2__user"),
            AIMessage(content="Good morning!"),
            HumanMessage(content="Third turn", id="msg-3"),
        ]
    }

    with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value.strftime.return_value = "2026-05-09, Saturday"
        result = mw.before_agent(state, _fake_runtime())

    assert result is None  # same day as last injected date → no update


# ---------------------------------------------------------------------------
# ID-swap recursive-injection guard (issue #3725)
# ---------------------------------------------------------------------------


def test_user_suffix_message_is_not_injection_target():
    """Regression guard: HumanMessage whose ID ends with ``__user`` must not be
    treated as an injection target.

    After the ID-swap in ``_make_reminder_and_user_messages``, the original user
    text becomes ``HumanMessage(id=X__user)``. If the middleware processes this
    message again, it would perform another ID-swap → ``X__user__user`` → … →
    unbounded suffix growth and ghost-message re-execution (issue #3725).
    """
    from deerflow.agents.middlewares.dynamic_context_middleware import _is_user_injection_target

    # A __user-suffix message is NOT a valid injection target
    user_swap_msg = HumanMessage(content="Hello", id="msg-1__user")
    assert _is_user_injection_target(user_swap_msg) is False

    # A __memory-suffix message is already tagged as a reminder, also rejected
    memory_swap_msg = HumanMessage(
        content="<memory>prefs</memory>",
        id="msg-1__memory",
        additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
    )
    assert _is_user_injection_target(memory_swap_msg) is False

    # A normal HumanMessage without __user suffix IS a valid target
    normal_msg = HumanMessage(content="Hello", id="msg-1")
    assert _is_user_injection_target(normal_msg) is True


def test_legacy_summary_message_is_not_injection_target():
    from deerflow.agents.middlewares.dynamic_context_middleware import _is_user_injection_target

    summary_msg = HumanMessage(content="Here is a summary of the conversation", name="summary")

    assert _is_user_injection_target(summary_msg) is False


def test_endswith_not_substring_prevents_false_positive():
    """``endswith("__user")`` must NOT reject messages whose ID merely contains
    ``__user`` somewhere in the middle (e.g. ``user__question-123``).

    A substring check (``"__user" in id``) would incorrectly reject such IDs.
    """
    from deerflow.agents.middlewares.dynamic_context_middleware import _is_user_injection_target

    # ID contains "__user" in the middle — should NOT be rejected
    middle_match = HumanMessage(content="question", id="user__question-123")
    assert _is_user_injection_target(middle_match) is True

    # ID ends with "__user" — should be rejected
    suffix_match = HumanMessage(content="question", id="msg-1__user")
    assert _is_user_injection_target(suffix_match) is False

    # Nested suffix "__user__user" — should also be rejected (recursive case)
    recursive_match = HumanMessage(content="question", id="msg-1__user__user")
    assert _is_user_injection_target(recursive_match) is False


def test_no_recursive_id_swap_in_full_middleware_flow():
    """End-to-end guard: after the first ID-swap, a second call to ``before_agent``
    must NOT produce a second swap on the ``__user`` message.

    This reproduces the exact scenario from issue #3725: a session with an
    existing ID-swap triplet receives a new HumanMessage, and the middleware
    must only inject into the new message — not re-process the ``__user`` peer.

    The state_v2 reminder deliberately omits the parseable date from both
    content and additional_kwargs so ``_last_injected_date`` returns None.
    This forces the first-turn injection path to actually reach
    ``_is_user_injection_target``, which must reject ``msg-1__user`` and
    select ``msg-2`` instead — exercising the endswith("__user") guard
    end-to-end rather than relying on the same-day short-circuit.
    """
    mw = _make_middleware()

    # First call: inject into HumanMessage(id="msg-1")
    state_v1 = {"messages": [HumanMessage(content="Hello", id="msg-1")]}

    with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt, mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=""):
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result_v1 = mw.before_agent(state_v1, _fake_runtime())

    assert result_v1 is not None
    msgs_v1 = result_v1["messages"]
    assert len(msgs_v1) == 2
    assert msgs_v1[0].id == "msg-1"  # reminder takes original ID
    assert msgs_v1[1].id == "msg-1__user"  # user content gets derived ID

    # Simulate state after first turn: ID-swap triplet (without parseable date
    # so _last_injected_date returns None → first-turn path is exercised)
    # + AI reply + new user message.
    state_v2 = {
        "messages": [
            SystemMessage(
                content="<system-reminder>\nplaceholder\n</system-reminder>",
                id="msg-1",
                additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
            ),
            HumanMessage(content="Hello", id="msg-1__user"),
            AIMessage(content="Hi there"),
            HumanMessage(content="Follow-up", id="msg-2"),
        ]
    }

    # Second call: _last_injected_date returns None (no parseable date),
    # so _inject enters first-turn path and must skip msg-1__user via the
    # endswith("__user") guard, then inject into msg-2.
    with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt, mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=""):
        mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
        result_v2 = mw.before_agent(state_v2, _fake_runtime())

    # The guard must route injection to msg-2, not msg-1__user.
    assert result_v2 is not None
    msgs_v2 = result_v2["messages"]
    assert msgs_v2[0].id == "msg-2"  # reminder takes new message's ID
    assert msgs_v2[1].id == "msg-2__user"  # user content gets derived ID


# ---------------------------------------------------------------------------
# Date timezone formatting
# ---------------------------------------------------------------------------


def test_format_current_date_defaults_to_server_local_without_env(monkeypatch):
    """Without DEER_FLOW_DATE_TIMEZONE the formatter keeps the legacy server-local behavior."""
    from datetime import datetime

    from deerflow.agents.middlewares.dynamic_context_middleware import _format_current_date

    with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value = datetime(2026, 5, 8, 9, 0)
        monkeypatch.delenv("DEER_FLOW_DATE_TIMEZONE", raising=False)

        assert _format_current_date() == "2026-05-08, Friday"
        mock_dt.now.assert_called_once_with()


def test_format_current_date_honors_configured_timezone(monkeypatch):
    """A UTC instant must be rendered in the IANA zone named by DEER_FLOW_DATE_TIMEZONE."""
    from datetime import UTC, datetime

    from deerflow.agents.middlewares.dynamic_context_middleware import _format_current_date

    # 2026-09-02 20:30 UTC is 2026-09-03 04:30 in Asia/Shanghai: a UTC-only
    # formatter would report the wrong day for a Shanghai user.
    fixed_utc = datetime(2026, 9, 2, 20, 30, 0, tzinfo=UTC)

    def fake_now(tz=None):
        # datetime.now(tz) semantics: the fixed instant expressed in *tz*.
        return fixed_utc.astimezone(tz) if tz is not None else fixed_utc.replace(tzinfo=None)

    with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.side_effect = fake_now
        monkeypatch.setenv("DEER_FLOW_DATE_TIMEZONE", "Asia/Shanghai")

        assert _format_current_date() == "2026-09-03, Thursday"


def test_format_current_date_invalid_timezone_falls_back(monkeypatch, caplog):
    """An unparseable IANA name must warn and degrade to the server-local timezone."""
    from datetime import datetime

    from deerflow.agents.middlewares.dynamic_context_middleware import _format_current_date

    with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
        mock_dt.now.return_value = datetime(2026, 5, 8, 9, 0)
        monkeypatch.setenv("DEER_FLOW_DATE_TIMEZONE", "Not/A_Zone")

        assert _format_current_date() == "2026-05-08, Friday"
    assert "DEER_FLOW_DATE_TIMEZONE" in caplog.text


def _declared_date_timezone_policies():
    from deerflow.agents.middlewares.dynamic_context_middleware import (
        DynamicContextMiddleware,
        SubagentDateContextMiddleware,
    )

    return [
        DynamicContextMiddleware().release_policy_parameters(),
        SubagentDateContextMiddleware().release_policy_parameters(),
    ]


# The request-only client block is model-visible, so its delivery and limits are
# part of DynamicContext's assembly identity.
_CLIENT_CONTEXT_POLICY = {
    "client_context_delivery": "framed_request_message_after_system_prompt",
    "client_context_limits": {"max_items": 32, "max_name_length": 80, "max_value_length": 200, "max_integer": 2**53 - 1},
    "client_context_pii_redaction": False,
}


def _expected_policies(tz_name: str):
    """The DynamicContext declaration also carries the shelf index's effective
    rendering caps (no app config here ⇒ ProjectsConfig defaults); the
    SubagentDateContext declaration is the timezone alone."""
    from deerflow.config.projects_config import ProjectsConfig

    defaults = ProjectsConfig()
    return [
        {
            "current_date_timezone": tz_name,
            "memory_enabled": True,
            "shelf_index_max_entries": defaults.shelf_index_max_entries,
            "shelf_index_max_bytes": defaults.shelf_index_max_bytes,
            **_CLIENT_CONTEXT_POLICY,
        },
        {"current_date_timezone": tz_name},
    ]


def test_date_middlewares_declare_configured_timezone(monkeypatch):
    """Assembly identity must reflect the zone the injected date follows."""
    monkeypatch.setenv("DEER_FLOW_DATE_TIMEZONE", "Asia/Shanghai")
    assert _declared_date_timezone_policies() == _expected_policies("Asia/Shanghai")


def test_date_middlewares_declare_utc_timezone(monkeypatch):
    monkeypatch.setenv("DEER_FLOW_DATE_TIMEZONE", "UTC")
    assert _declared_date_timezone_policies() == _expected_policies("UTC")


def test_dynamic_context_release_policy_includes_memory_opt_out(monkeypatch):
    from deerflow.agents.middlewares.dynamic_context_middleware import DynamicContextMiddleware

    monkeypatch.setenv("DEER_FLOW_DATE_TIMEZONE", "UTC")

    from deerflow.config.projects_config import ProjectsConfig

    defaults = ProjectsConfig()
    assert DynamicContextMiddleware(memory_enabled=False).release_policy_parameters() == {
        "current_date_timezone": "UTC",
        "memory_enabled": False,
        "shelf_index_max_entries": defaults.shelf_index_max_entries,
        "shelf_index_max_bytes": defaults.shelf_index_max_bytes,
        **_CLIENT_CONTEXT_POLICY,
    }


def test_date_middlewares_declare_resolved_local_zone_without_env(monkeypatch):
    """Without the knob the declaration resolves the actual local zone, so two
    hosts that render different dates still get different assembly fingerprints."""
    from deerflow.agents.middlewares.dynamic_context_middleware import _effective_date_timezone_name

    monkeypatch.delenv("DEER_FLOW_DATE_TIMEZONE", raising=False)
    expected_zone = _effective_date_timezone_name()
    assert expected_zone

    assert _declared_date_timezone_policies() == _expected_policies(expected_zone)


def test_date_middlewares_declare_resolved_local_zone_for_invalid_env(monkeypatch, caplog):
    """An invalid IANA name degrades to server-local and is declared as such."""
    from deerflow.agents.middlewares.dynamic_context_middleware import _effective_date_timezone_name

    monkeypatch.setenv("DEER_FLOW_DATE_TIMEZONE", "Not/A_Zone")
    assert _declared_date_timezone_policies() == _expected_policies(_effective_date_timezone_name())
    assert "DEER_FLOW_DATE_TIMEZONE" in caplog.text


def test_server_local_timezone_name_reads_tz_env(monkeypatch):
    """A POSIX TZ env var naming a real zone resolves to its IANA key."""
    from deerflow.agents.middlewares.dynamic_context_middleware import _server_local_timezone_name

    monkeypatch.setenv("TZ", "Asia/Shanghai")
    assert _server_local_timezone_name() == "Asia/Shanghai"


def test_server_local_timezone_name_reads_direct_macos_symlink_target(monkeypatch):
    """macOS /etc/localtime points at the unversioned zoneinfo dir; the direct
    symlink target must be read instead of a fully resolved path."""
    import deerflow.agents.middlewares.dynamic_context_middleware as module

    monkeypatch.delenv("TZ", raising=False)
    monkeypatch.setattr(module.os, "readlink", lambda _path: "/var/db/timezone/zoneinfo/Asia/Shanghai")

    assert module._server_local_timezone_name() == "Asia/Shanghai"


def test_server_local_timezone_name_reads_apple_versioned_symlink_target(monkeypatch):
    """Apple's canonical versioned zoneinfo path must still yield the zone key."""
    import deerflow.agents.middlewares.dynamic_context_middleware as module

    monkeypatch.delenv("TZ", raising=False)
    monkeypatch.setattr(
        module.os,
        "readlink",
        lambda _path: "/private/var/db/timezone/tz/2026c.1.0/zoneinfo/America/New_York",
    )

    assert module._server_local_timezone_name() == "America/New_York"


def test_server_local_timezone_name_normalizes_relative_symlink_target(monkeypatch):
    """A relative /etc/localtime target is resolved against /etc."""
    import deerflow.agents.middlewares.dynamic_context_middleware as module

    monkeypatch.delenv("TZ", raising=False)
    monkeypatch.setattr(module.os, "readlink", lambda _path: "../usr/share/zoneinfo/Etc/UTC")

    assert module._server_local_timezone_name() == "Etc/UTC"


def test_effective_timezone_sentinel_uses_offset_when_local_zone_is_not_resolvable(monkeypatch):
    """Without a recoverable IANA key the declaration pins a stable sentinel."""
    import deerflow.agents.middlewares.dynamic_context_middleware as module

    monkeypatch.setattr(module, "_server_local_timezone_name", lambda: None)
    monkeypatch.setattr(module, "_server_local_utc_offset_minutes", lambda: 8 * 60)
    monkeypatch.delenv("DEER_FLOW_DATE_TIMEZONE", raising=False)

    assert module._effective_date_timezone_name() == "server-local(+08:00)"

    monkeypatch.setattr(module, "_server_local_utc_offset_minutes", lambda: -5 * 60 - 30)
    assert module._effective_date_timezone_name() == "server-local(-05:30)"