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).
"""

from types import SimpleNamespace
from unittest import mock

from langchain_core.messages import AIMessage, HumanMessage

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

_SYSTEM_REMINDER_TAG = "<system-reminder>"


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


def _fake_runtime():
    return SimpleNamespace(context={})


def _reminder_msg(content: str, msg_id: str) -> HumanMessage:
    """Build a reminder HumanMessage the way the middleware would produce it."""
    return HumanMessage(
        content=content,
        id=msg_id,
        additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
    )


# ---------------------------------------------------------------------------
# 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, HumanMessage)
    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())

    # Reminder is the first returned message; user query is the second
    reminder_content = result["messages"][0].content
    assert "User prefers Python." in reminder_content
    assert "<current_date>2026-05-08, Friday</current_date>" in reminder_content
    assert result["messages"][1].content == "Hi"


# ---------------------------------------------------------------------------
# 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()
    reminder_content = "<system-reminder>\n<current_date>2026-05-08, Friday</current_date>\n</system-reminder>"
    state = {
        "messages": [
            _reminder_msg(reminder_content, "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_injects_only_into_first_human_message_not_later_ones():
    """Reminder targets the first HumanMessage; subsequent messages are not touched."""
    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 + original first query)
    assert len(msgs) == 2
    assert msgs[0].id == "msg-1"  # reminder takes first 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-1__user"  # original content with derived ID
    assert msgs[1].content == "First"
    # "Second" (msg-2) is not in the returned update — it is left unchanged
    assert all(m.id != "msg-2" 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_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


# ---------------------------------------------------------------------------
# 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()
    reminder_content = "<system-reminder>\n<current_date>2026-05-08, Friday</current_date>\n</system-reminder>"
    state = {
        "messages": [
            _reminder_msg(reminder_content, "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

    # 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()
    reminder_content = "<system-reminder>\n<current_date>2026-05-08, Friday</current_date>\n</system-reminder>"
    state = {
        "messages": [
            _reminder_msg(reminder_content, "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_no_second_midnight_injection_once_date_updated():
    """After a midnight update is persisted, the same-day path skips re-injection."""
    mw = _make_middleware()
    date_update_content = "<system-reminder>\n<current_date>2026-05-09, Saturday</current_date>\n</system-reminder>"
    state = {
        "messages": [
            _reminder_msg(
                "<system-reminder>\n<current_date>2026-05-08, Friday</current_date>\n</system-reminder>",
                "msg-1",
            ),
            HumanMessage(content="Hello", id="msg-1__user"),
            AIMessage(content="Response"),
            _reminder_msg(date_update_content, "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