github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-08-11 15:28:37 +00:00
Code Issues Packages Projects Releases Wiki Activity
deer-flow/backend/tests/test_memory_staleness_review.py
Tianye Song 54f3c43fe3
fix(security): html-escape fact content in memory prompt sections (#4028) 
* fix(security): html-escape fact content in memory prompt sections

Raw memory fact content was injected verbatim into prompt XML — a fact
containing a literal `"` could break the `"..."` delimiter, and a
closing tag like `</consolidation_candidates>` could prematurely end
the XML block, both potentially confusing the model.

Apply `html.escape()` to `content` in `_build_staleness_section` and
`_build_consolidation_section`, and to `cat` in the consolidation
section's XML attribute. Tests added for both sections covering special
characters, XML tag injection, and attribute injection.

Follow-up to #3996 as noted by reviewer willem-bd.

* fix(security): address reviewer follow-ups on html-escaping PR

- Escape `cat` in _build_staleness_section for symmetry with the
  consolidation section (both sections now consistently html-escape
  all LLM-derived category values that appear in the prompt)
- Add comment at current_memory=json.dumps() documenting the conscious
  accept: json.dumps leaves < > & unescaped; lower-risk than
  staleness/consolidation (read-only context, not delete/merge
  instructions); fix at fact-content insert time if revisited
- Add test for category escaping in the staleness section

* fix(security): reference tracking issue #4044 in conscious-accept comment

* style: compress conscious-accept comment to two lines
2026-07-11 08:54:45 +08:00

722 lines
27 KiB
Python
Raw Blame History

"""Tests for the staleness review feature in the memory updater.
Covers:
- Candidate selection (age threshold, protected categories)
- Trigger conditions (min candidates, enabled flag)
- Prompt section formatting
- Staleness removal in _apply_updates (safety cap, observability)
- Normalization of staleFactsToRemove from LLM responses
- Integration with _prepare_update_prompt
"""
from datetime import UTC, datetime, timedelta
from unittest.mock import MagicMock, patch
from deerflow.agents.memory.updater import (
MemoryUpdater,
_build_staleness_section,
_normalize_memory_update_data,
_parse_fact_datetime,
_select_stale_candidates,
)
from deerflow.config.memory_config import MemoryConfig
# ── Helpers ────────────────────────────────────────────────────────────────
def _memory_config(**overrides: object) -> MemoryConfig:
config = MemoryConfig()
for key, value in overrides.items():
setattr(config, key, value)
return config
def _make_fact(
fact_id: str,
content: str = "test content",
category: str = "knowledge",
confidence: float = 0.9,
days_ago: int = 100,
) -> dict:
created = (datetime.now(UTC) - timedelta(days=days_ago)).isoformat().replace("+00:00", "Z")
return {
"id": fact_id,
"content": content,
"category": category,
"confidence": confidence,
"createdAt": created,
"source": "thread-test",
}
def _make_memory(facts: list[dict] | None = None) -> dict:
return {
"version": "1.0",
"lastUpdated": "",
"user": {
"workContext": {"summary": "", "updatedAt": ""},
"personalContext": {"summary": "", "updatedAt": ""},
"topOfMind": {"summary": "", "updatedAt": ""},
},
"history": {
"recentMonths": {"summary": "", "updatedAt": ""},
"earlierContext": {"summary": "", "updatedAt": ""},
"longTermBackground": {"summary": "", "updatedAt": ""},
},
"facts": facts or [],
}
# ── _parse_fact_datetime ──────────────────────────────────────────────────
class TestParseFactDatetime:
def test_z_suffix(self):
result = _parse_fact_datetime("2025-06-01T12:00:00Z")
assert result is not None
assert result.year == 2025
assert result.month == 6
def test_offset_format(self):
result = _parse_fact_datetime("2025-06-01T12:00:00+00:00")
assert result is not None
assert result.year == 2025
def test_empty_string(self):
assert _parse_fact_datetime("") is None
def test_invalid_format(self):
assert _parse_fact_datetime("not-a-date") is None
def test_naive_datetime_gets_utc(self):
"""Naive datetime (no tzinfo) should be treated as UTC, not cause TypeError."""
result = _parse_fact_datetime("2025-06-01T12:00:00")
assert result is not None
assert result.tzinfo is not None
assert result.utcoffset().total_seconds() == 0
# ── _select_stale_candidates ──────────────────────────────────────────────
class TestSelectStaleCandidates:
def test_old_facts_selected(self):
memory = _make_memory(
[
_make_fact("fact_old", days_ago=100),
_make_fact("fact_new", days_ago=10),
]
)
config = _memory_config(staleness_age_days=90)
candidates = _select_stale_candidates(memory, config)
assert len(candidates) == 1
assert candidates[0]["id"] == "fact_old"
def test_protected_category_excluded(self):
memory = _make_memory(
[
_make_fact("fact_correction", category="correction", days_ago=200),
_make_fact("fact_knowledge", category="knowledge", days_ago=200),
]
)
config = _memory_config(staleness_age_days=90, staleness_protected_categories=["correction"])
candidates = _select_stale_candidates(memory, config)
assert len(candidates) == 1
assert candidates[0]["id"] == "fact_knowledge"
def test_custom_protected_categories(self):
memory = _make_memory(
[
_make_fact("fact_goal", category="goal", days_ago=200),
]
)
config = _memory_config(staleness_age_days=90, staleness_protected_categories=["goal"])
candidates = _select_stale_candidates(memory, config)
assert len(candidates) == 0
def test_no_facts(self):
memory = _make_memory([])
config = _memory_config(staleness_age_days=90)
assert _select_stale_candidates(memory, config) == []
def test_all_recent(self):
memory = _make_memory(
[
_make_fact("fact_a", days_ago=10),
_make_fact("fact_b", days_ago=30),
]
)
config = _memory_config(staleness_age_days=90)
assert _select_stale_candidates(memory, config) == []
# ── Trigger conditions via _select_stale_candidates + config ─────────────
class TestStalenessTriggerConditions:
"""The old _should_run_staleness_review was removed; trigger logic is now
inlined in _prepare_update_prompt. We verify the gating conditions here
through _select_stale_candidates + config flags directly."""
def test_disabled_means_no_section(self):
memory = _make_memory([_make_fact(f"f{i}", days_ago=100) for i in range(5)])
config = _memory_config(staleness_review_enabled=False, staleness_age_days=90, staleness_min_candidates=3)
candidates = _select_stale_candidates(memory, config)
# Even though candidates exist, the caller checks enabled flag first
assert config.staleness_review_enabled is False
assert len(candidates) >= config.staleness_min_candidates
def test_below_min_candidates(self):
memory = _make_memory([_make_fact("fact_only", days_ago=100)])
config = _memory_config(staleness_review_enabled=True, staleness_age_days=90, staleness_min_candidates=3)
candidates = _select_stale_candidates(memory, config)
assert len(candidates) < config.staleness_min_candidates
def test_at_min_candidates(self):
memory = _make_memory([_make_fact(f"fact_{i}", days_ago=100) for i in range(3)])
config = _memory_config(staleness_review_enabled=True, staleness_age_days=90, staleness_min_candidates=3)
candidates = _select_stale_candidates(memory, config)
assert len(candidates) >= config.staleness_min_candidates
def test_above_min_candidates(self):
memory = _make_memory([_make_fact(f"fact_{i}", days_ago=100) for i in range(10)])
config = _memory_config(staleness_review_enabled=True, staleness_age_days=90, staleness_min_candidates=3)
candidates = _select_stale_candidates(memory, config)
assert len(candidates) >= config.staleness_min_candidates
# ── _build_staleness_section ──────────────────────────────────────────────
class TestBuildStalenessSection:
def test_empty_candidates(self):
assert _build_staleness_section([], 90) == ""
def test_includes_fact_details(self):
candidates = [
_make_fact("fact_vue", "User uses Vue.js", "knowledge", 0.95, days_ago=120),
]
section = _build_staleness_section(candidates, 90)
assert "fact_vue" in section
assert "User uses Vue.js" in section
assert "0.95" in section
assert "90 days" in section
def test_multiple_facts(self):
candidates = [
_make_fact("fact_a", "Fact A", "knowledge", 0.9, days_ago=100),
_make_fact("fact_b", "Fact B", "preference", 0.8, days_ago=150),
]
section = _build_staleness_section(candidates, 90)
assert "fact_a" in section
assert "fact_b" in section
assert "<stale_facts>" in section
def test_html_special_chars_in_content_are_escaped(self):
"""Fact content with XML tags or quotes is HTML-escaped so it cannot
break the surrounding prompt structure."""
candidates = [
_make_fact("fact_x", 'Like <b>bold</b> & "quotes"', "knowledge", 0.9, days_ago=100),
]
section = _build_staleness_section(candidates, 90)
assert "<b>" not in section
assert "&lt;b&gt;" in section
assert "&amp;" in section
assert "&quot;" in section
def test_closing_tag_in_content_is_escaped(self):
"""A closing </stale_facts> tag embedded in content must not prematurely
end the prompt XML block."""
candidates = [
_make_fact("fact_y", "</stale_facts><injected>bad</injected>", "knowledge", 0.8, days_ago=100),
]
section = _build_staleness_section(candidates, 90)
assert "</stale_facts><injected>" not in section
assert "&lt;/stale_facts&gt;" in section
def test_special_chars_in_category_are_escaped(self):
"""A category name with XML tags or quotes is HTML-escaped, consistent
with how category is handled in the consolidation section."""
candidates = [
_make_fact("fact_z", "content", 'pref<"erences>', 0.8, days_ago=100),
]
section = _build_staleness_section(candidates, 90)
assert 'pref<"erences>' not in section
assert "pref&lt;&quot;erences&gt;" in section
# ── _apply_updates with staleness removals ─────────────────────────────────
class TestApplyUpdatesStaleness:
def test_stale_facts_removed(self):
updater = MemoryUpdater()
current_memory = _make_memory(
[
_make_fact("fact_keep", "User knows Python", days_ago=100),
_make_fact("fact_stale", "User uses Vue.js", days_ago=120),
]
)
update_data = {
"user": {},
"history": {},
"newFacts": [],
"factsToRemove": [],
"staleFactsToRemove": [
{"id": "fact_stale", "reason": "User switched to React"},
],
}
with patch(
"deerflow.agents.memory.updater.get_memory_config",
return_value=_memory_config(max_facts=100, staleness_max_removals_per_cycle=10),
):
result = updater._apply_updates(current_memory, update_data)
assert len(result["facts"]) == 1
assert result["facts"][0]["id"] == "fact_keep"
def test_stale_candidate_without_id_does_not_raise(self):
"""A legacy / hand-edited fact that lacks an ``id`` must not crash the
staleness apply path.
Regression: ``candidate_ids`` was built with a direct ``f["id"]``
access over ``_select_stale_candidates`` output, but every other fact
access in the module uses ``f.get("id")``. An aged, non-protected fact
with no ``id`` key (common in legacy / migrated ``memory.json``) is a
valid staleness candidate, so it reached ``f["id"]`` and raised
``KeyError: 'id'``, aborting the whole memory-update cycle.
"""
updater = MemoryUpdater()
aged = (datetime.now(UTC) - timedelta(days=120)).isoformat().replace("+00:00", "Z")
# An aged, non-protected fact deliberately missing the "id" key.
idless_fact = {"content": "User uses Vue.js", "category": "knowledge", "confidence": 0.8, "createdAt": aged}
current_memory = _make_memory([_make_fact("fact_keep", days_ago=100), idless_fact])
update_data = {
"user": {},
"history": {},
"newFacts": [],
"factsToRemove": [],
"staleFactsToRemove": [
{"id": "fact_keep", "reason": "outdated"},
],
}
with patch(
"deerflow.agents.memory.updater.get_memory_config",
return_value=_memory_config(max_facts=100, staleness_max_removals_per_cycle=10),
):
# Must not raise KeyError: 'id'.
result = updater._apply_updates(current_memory, update_data)
# The id-less fact survives (it can never be targeted by the id-based
# removal set), and the id-based removal of fact_keep still applies.
contents = {f.get("content") for f in result["facts"]}
assert "User uses Vue.js" in contents
def test_safety_cap_limits_removals(self):
updater = MemoryUpdater()
# 5 stale facts, but cap is 2 → only 2 lowest-confidence should be removed
current_memory = _make_memory(
[
_make_fact("fact_high", confidence=0.95, days_ago=100),
_make_fact("fact_mid", confidence=0.80, days_ago=100),
_make_fact("fact_low1", confidence=0.70, days_ago=100),
_make_fact("fact_low2", confidence=0.65, days_ago=100),
_make_fact("fact_low3", confidence=0.60, days_ago=100),
]
)
update_data = {
"user": {},
"history": {},
"newFacts": [],
"factsToRemove": [],
"staleFactsToRemove": [
{"id": "fact_high", "reason": "outdated"},
{"id": "fact_mid", "reason": "outdated"},
{"id": "fact_low1", "reason": "outdated"},
{"id": "fact_low2", "reason": "outdated"},
{"id": "fact_low3", "reason": "outdated"},
],
}
with patch(
"deerflow.agents.memory.updater.get_memory_config",
return_value=_memory_config(max_facts=100, staleness_max_removals_per_cycle=2),
):
result = updater._apply_updates(current_memory, update_data)
# 5 - 2 = 3 facts remain; the 2 lowest-confidence removed
assert len(result["facts"]) == 3
remaining_ids = {f["id"] for f in result["facts"]}
assert "fact_high" in remaining_ids
assert "fact_mid" in remaining_ids
assert "fact_low1" in remaining_ids
def test_empty_stale_removals_no_effect(self):
updater = MemoryUpdater()
current_memory = _make_memory(
[
_make_fact("fact_a", days_ago=100),
]
)
update_data = {
"user": {},
"history": {},
"newFacts": [],
"factsToRemove": [],
"staleFactsToRemove": [],
}
with patch(
"deerflow.agents.memory.updater.get_memory_config",
return_value=_memory_config(max_facts=100),
):
result = updater._apply_updates(current_memory, update_data)
assert len(result["facts"]) == 1
def test_missing_stale_removals_key_no_effect(self):
"""When LLM doesn't return staleFactsToRemove, existing behavior is preserved."""
updater = MemoryUpdater()
current_memory = _make_memory([_make_fact("fact_a")])
update_data = {
"user": {},
"history": {},
"newFacts": [],
"factsToRemove": [],
# no staleFactsToRemove key
}
with patch(
"deerflow.agents.memory.updater.get_memory_config",
return_value=_memory_config(max_facts=100),
):
result = updater._apply_updates(current_memory, update_data)
assert len(result["facts"]) == 1
def test_contradiction_and_staleness_removals_combined(self):
"""Both factsToRemove and staleFactsToRemove work together."""
updater = MemoryUpdater()
current_memory = _make_memory(
[
_make_fact("fact_keep", days_ago=10),
_make_fact("fact_contradicted", days_ago=10),
_make_fact("fact_stale", days_ago=200),
]
)
update_data = {
"user": {},
"history": {},
"newFacts": [],
"factsToRemove": ["fact_contradicted"],
"staleFactsToRemove": [{"id": "fact_stale", "reason": "old"}],
}
with patch(
"deerflow.agents.memory.updater.get_memory_config",
return_value=_memory_config(max_facts=100, staleness_max_removals_per_cycle=10),
):
result = updater._apply_updates(current_memory, update_data)
assert len(result["facts"]) == 1
assert result["facts"][0]["id"] == "fact_keep"
def test_protected_category_fact_refused_at_apply(self):
"""Regression: LLM hallucinating a correction-category fact id in
staleFactsToRemove must be silently rejected at the apply layer,
even though it appears in the serialized prompt JSON."""
updater = MemoryUpdater()
current_memory = _make_memory(
[
_make_fact("fact_stale", category="knowledge", days_ago=200),
_make_fact("fact_correction", category="correction", days_ago=200),
]
)
update_data = {
"user": {},
"history": {},
"newFacts": [],
"factsToRemove": [],
"staleFactsToRemove": [
{"id": "fact_stale", "reason": "outdated"},
{"id": "fact_correction", "reason": "LLM slip"},
],
}
with patch(
"deerflow.agents.memory.updater.get_memory_config",
return_value=_memory_config(
max_facts=100,
staleness_review_enabled=True,
staleness_age_days=90,
staleness_min_candidates=1,
staleness_max_removals_per_cycle=10,
staleness_protected_categories=["correction"],
),
):
result = updater._apply_updates(current_memory, update_data)
# fact_stale removed, fact_correction kept (protected)
assert len(result["facts"]) == 1
assert result["facts"][0]["id"] == "fact_correction"
def test_non_aged_fact_refused_at_apply(self):
"""Regression: LLM returning a fresh (non-aged) fact id in
staleFactsToRemove must be silently rejected."""
updater = MemoryUpdater()
current_memory = _make_memory(
[
_make_fact("fact_stale", days_ago=200),
_make_fact("fact_fresh", days_ago=10),
]
)
update_data = {
"user": {},
"history": {},
"newFacts": [],
"factsToRemove": [],
"staleFactsToRemove": [
{"id": "fact_stale", "reason": "outdated"},
{"id": "fact_fresh", "reason": "LLM hallucination"},
],
}
with patch(
"deerflow.agents.memory.updater.get_memory_config",
return_value=_memory_config(
max_facts=100,
staleness_review_enabled=True,
staleness_age_days=90,
staleness_min_candidates=1,
staleness_max_removals_per_cycle=10,
staleness_protected_categories=["correction"],
),
):
result = updater._apply_updates(current_memory, update_data)
# fact_stale removed, fact_fresh kept (not in candidate set)
assert len(result["facts"]) == 1
assert result["facts"][0]["id"] == "fact_fresh"
def test_guardrail_runs_when_staleness_review_disabled(self):
"""Regression: guardrail must reject invalid ids even when
staleness_review_enabled=False, so the protection is independent
of the feature flag and model behavior."""
updater = MemoryUpdater()
current_memory = _make_memory(
[
_make_fact("fact_stale", days_ago=200),
_make_fact("fact_fresh", days_ago=5),
]
)
update_data = {
"user": {},
"history": {},
"newFacts": [],
"factsToRemove": [],
"staleFactsToRemove": [
{"id": "fact_stale", "reason": "LLM hallucination"},
{"id": "fact_fresh", "reason": "LLM hallucination"},
],
}
with patch(
"deerflow.agents.memory.updater.get_memory_config",
return_value=_memory_config(
max_facts=100,
staleness_review_enabled=False,
staleness_age_days=90,
staleness_min_candidates=3,
staleness_max_removals_per_cycle=10,
staleness_protected_categories=["correction"],
),
):
result = updater._apply_updates(current_memory, update_data)
# Guardrail runs regardless of feature flag:
# fact_stale is a valid candidate (200 days old) → removed
# fact_fresh is not a candidate (5 days old) → kept
assert len(result["facts"]) == 1
assert result["facts"][0]["id"] == "fact_fresh"
# ── _normalize_memory_update_data with staleFactsToRemove ─────────────────
class TestNormalizeStaleFactsToRemove:
def test_valid_entries(self):
data = {
"user": {},
"history": {},
"newFacts": [],
"factsToRemove": [],
"staleFactsToRemove": [
{"id": "fact_a", "reason": "User moved offices"},
{"id": "fact_b", "reason": "Tech stack changed"},
],
}
result = _normalize_memory_update_data(data)
assert len(result["staleFactsToRemove"]) == 2
assert result["staleFactsToRemove"][0]["id"] == "fact_a"
assert result["staleFactsToRemove"][1]["reason"] == "Tech stack changed"
def test_missing_key(self):
data = {"user": {}, "history": {}, "newFacts": [], "factsToRemove": []}
result = _normalize_memory_update_data(data)
assert result["staleFactsToRemove"] == []
def test_non_list_ignored(self):
data = {
"user": {},
"history": {},
"newFacts": [],
"factsToRemove": [],
"staleFactsToRemove": "not a list",
}
result = _normalize_memory_update_data(data)
assert result["staleFactsToRemove"] == []
def test_non_dict_entries_skipped(self):
data = {
"user": {},
"history": {},
"newFacts": [],
"factsToRemove": [],
"staleFactsToRemove": ["just a string", 42, {"id": "fact_ok", "reason": "valid"}],
}
result = _normalize_memory_update_data(data)
assert len(result["staleFactsToRemove"]) == 1
assert result["staleFactsToRemove"][0]["id"] == "fact_ok"
def test_empty_id_skipped(self):
data = {
"user": {},
"history": {},
"newFacts": [],
"factsToRemove": [],
"staleFactsToRemove": [{"id": "", "reason": "no id"}],
}
result = _normalize_memory_update_data(data)
assert result["staleFactsToRemove"] == []
def test_non_string_reason_defaulted(self):
data = {
"user": {},
"history": {},
"newFacts": [],
"factsToRemove": [],
"staleFactsToRemove": [{"id": "fact_a", "reason": 123}],
}
result = _normalize_memory_update_data(data)
assert result["staleFactsToRemove"][0]["reason"] == ""
def test_missing_reason_defaulted(self):
data = {
"user": {},
"history": {},
"newFacts": [],
"factsToRemove": [],
"staleFactsToRemove": [{"id": "fact_a"}],
}
result = _normalize_memory_update_data(data)
assert result["staleFactsToRemove"][0]["reason"] == ""
# ── Integration: _prepare_update_prompt ────────────────────────────────────
class TestPrepareUpdatePromptStaleness:
def test_staleness_section_included_when_triggered(self):
updater = MemoryUpdater()
old_facts = [_make_fact(f"fact_{i}", days_ago=100) for i in range(5)]
memory = _make_memory(old_facts)
msg = MagicMock()
msg.type = "human"
msg.content = "Hello, I'm using React now"
config = _memory_config(
enabled=True,
staleness_review_enabled=True,
staleness_age_days=90,
staleness_min_candidates=3,
)
with (
patch("deerflow.agents.memory.updater.get_memory_config", return_value=config),
patch("deerflow.agents.memory.updater.get_memory_data", return_value=memory),
):
result = updater._prepare_update_prompt(
messages=[msg],
agent_name=None,
correction_detected=False,
reinforcement_detected=False,
)
assert result is not None
_, prompt = result
assert "Staleness Review" in prompt
assert "<stale_facts>" in prompt
def test_staleness_section_omitted_when_not_triggered(self):
updater = MemoryUpdater()
memory = _make_memory([]) # no facts at all
msg = MagicMock()
msg.type = "human"
msg.content = "Hello"
config = _memory_config(
enabled=True,
staleness_review_enabled=True,
staleness_age_days=90,
staleness_min_candidates=3,
)
with (
patch("deerflow.agents.memory.updater.get_memory_config", return_value=config),
patch("deerflow.agents.memory.updater.get_memory_data", return_value=memory),
):
result = updater._prepare_update_prompt(
messages=[msg],
agent_name=None,
correction_detected=False,
reinforcement_detected=False,
)
assert result is not None
_, prompt = result
assert "Staleness Review" not in prompt
assert "<stale_facts>" not in prompt
def test_staleness_section_omitted_when_disabled(self):
updater = MemoryUpdater()
old_facts = [_make_fact(f"fact_{i}", days_ago=200) for i in range(10)]
memory = _make_memory(old_facts)
msg = MagicMock()
msg.type = "human"
msg.content = "Hello"
config = _memory_config(
enabled=True,
staleness_review_enabled=False,
)
with (
patch("deerflow.agents.memory.updater.get_memory_config", return_value=config),
patch("deerflow.agents.memory.updater.get_memory_data", return_value=memory),
):
result = updater._prepare_update_prompt(
messages=[msg],
agent_name=None,
correction_detected=False,
reinforcement_detected=False,
)
assert result is not None
_, prompt = result
assert "Staleness Review" not in prompt
Reference in New Issue View Git Blame Copy Permalink