From b990da785f0ffccfe50ef4fb4a772f8099c7fdde Mon Sep 17 00:00:00 2001 From: Tianye Song <162393000+sontianye@users.noreply.github.com> Date: Tue, 23 Jun 2026 08:10:12 +0800 Subject: [PATCH] feat(memory): add guaranteed injection for correction facts with graceful fallback (#3592) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(memory): add guaranteed injection for correction facts with graceful fallback When the token budget is tight, high-value facts (e.g. user corrections) can be silently evicted by lower-priority regular facts. This change: - Introduces configurable 'guaranteed_categories' (default: [correction]) whose facts draw from a separate 'guaranteed_token_budget', ensuring they are never dropped due to budget pressure. - Adds a graceful fallback to confidence-only ranking when the guaranteed-category path raises an unexpected exception. - Refactors fact selection into a header-agnostic helper (_select_fact_lines) with explicit token accounting in the caller, eliminating double-counting of separators. - Emits a single 'Facts:' header regardless of whether both guaranteed and regular facts are present. - Extends the final safety truncation limit to account for the additional guaranteed budget so guaranteed facts survive end-to-end. * refactor(memory): address review feedback on guaranteed injection - Restore strict break-on-overflow in `_select_fact_lines` to preserve the caller's confidence-ordered ranking; add a regression test locking in the invariant that a shorter lower-confidence fact never slips ahead of a skipped higher-confidence one. - Account for the inter-group `\n` separator between guaranteed and regular fact blocks in the regular budget (1-token precision fix). - Clarify docstrings on `format_memory_for_injection` and `MemoryConfig.guaranteed_token_budget` to distinguish the common *displacement* case (total stays within `max_tokens`) from the rarer *additive* case (safety-truncation ceiling raised when guaranteed lines alone would overflow). * fix(memory): address P1 safety truncation + P2s from review - Structure-aware safety truncation: Facts block is now a protected suffix so guaranteed-category facts can never be silently discarded by a prefix-cut on overflow. Only the preceding (user/history) sections are eligible for truncation. - Extend the same protected-suffix treatment to the except/fallback path by returning fact lines alongside the formatted section from _fallback_format_facts, avoiding string parsing. - Single inter-section separator: facts section no longer embeds its own leading \n\n; the final "\n\n".join(sections) is the single source of truth for section-to-section spacing. - Bare string for guaranteed_categories now raises TypeError instead of silently iterating single characters. - Category-less / malformed facts no longer default-promote into the guaranteed "context" pool — only facts with an explicit category field qualify. - Lift valid_facts pre-filter outside the try so the fallback path reuses it instead of re-doing validation work. - MemoryConfigResponse + DeerFlowClient.get_memory_config now expose guaranteed_categories / guaranteed_token_budget. - config.example.yaml: document the two new fields and bump config_version from 12 to 13. - Add regression tests for every finding. --------- Co-authored-by: Willem Jiang --- backend/app/gateway/routers/memory.py | 12 + .../deerflow/agents/lead_agent/prompt.py | 2 + .../harness/deerflow/agents/memory/prompt.py | 335 ++++++++++-- backend/packages/harness/deerflow/client.py | 2 + .../harness/deerflow/config/memory_config.py | 25 + backend/tests/test_lead_agent_prompt.py | 9 +- backend/tests/test_memory_prompt_injection.py | 481 ++++++++++++++++++ config.example.yaml | 14 + 8 files changed, 833 insertions(+), 47 deletions(-) diff --git a/backend/app/gateway/routers/memory.py b/backend/app/gateway/routers/memory.py index 6275303e5..6d1c356c8 100644 --- a/backend/app/gateway/routers/memory.py +++ b/backend/app/gateway/routers/memory.py @@ -123,6 +123,14 @@ class MemoryConfigResponse(BaseModel): injection_enabled: bool = Field(..., description="Whether memory injection is enabled") max_injection_tokens: int = Field(..., description="Maximum tokens for memory injection") token_counting: str = Field(..., description="Token counting strategy for memory injection ('tiktoken' or 'char')") + guaranteed_categories: list[str] = Field( + ..., + description="Fact categories that bypass the regular injection budget (always injected from a reserved allowance)", + ) + guaranteed_token_budget: int = Field( + ..., + description="Token ceiling for guaranteed-category facts (displaces regular lines in the common case; additive only when guaranteed alone overflows max_injection_tokens)", + ) class MemoryStatusResponse(BaseModel): @@ -350,6 +358,8 @@ async def get_memory_config_endpoint() -> MemoryConfigResponse: injection_enabled=config.injection_enabled, max_injection_tokens=config.max_injection_tokens, token_counting=config.token_counting, + guaranteed_categories=config.guaranteed_categories, + guaranteed_token_budget=config.guaranteed_token_budget, ) @@ -379,6 +389,8 @@ async def get_memory_status(http_request: Request) -> MemoryStatusResponse: injection_enabled=config.injection_enabled, max_injection_tokens=config.max_injection_tokens, token_counting=config.token_counting, + guaranteed_categories=config.guaranteed_categories, + guaranteed_token_budget=config.guaranteed_token_budget, ), data=MemoryResponse(**memory_data), ) diff --git a/backend/packages/harness/deerflow/agents/lead_agent/prompt.py b/backend/packages/harness/deerflow/agents/lead_agent/prompt.py index df9ac5333..f94a07014 100644 --- a/backend/packages/harness/deerflow/agents/lead_agent/prompt.py +++ b/backend/packages/harness/deerflow/agents/lead_agent/prompt.py @@ -610,6 +610,8 @@ def _get_memory_context(agent_name: str | None = None, *, app_config: AppConfig memory_data, max_tokens=config.max_injection_tokens, use_tiktoken=(config.token_counting == "tiktoken"), + guaranteed_categories=getattr(config, "guaranteed_categories", None), + guaranteed_token_budget=getattr(config, "guaranteed_token_budget", 500), ) if not memory_content.strip(): diff --git a/backend/packages/harness/deerflow/agents/memory/prompt.py b/backend/packages/harness/deerflow/agents/memory/prompt.py index 69fe48aff..ac5620236 100644 --- a/backend/packages/harness/deerflow/agents/memory/prompt.py +++ b/backend/packages/harness/deerflow/agents/memory/prompt.py @@ -316,7 +316,115 @@ def _coerce_confidence(value: Any, default: float = 0.0) -> float: return max(0.0, min(1.0, confidence)) -def format_memory_for_injection(memory_data: dict[str, Any], max_tokens: int = 2000, *, use_tiktoken: bool = True) -> str: +def _format_fact_line(fact: dict[str, Any]) -> str | None: + """Build a single formatted fact line, or return ``None`` for invalid facts. + + Extracted as a shared helper so the guaranteed-injection and regular-injection + paths produce identical line formatting. + """ + content_value = fact.get("content") + if not isinstance(content_value, str): + return None + content = content_value.strip() + if not content: + return None + category = str(fact.get("category", "context")).strip() or "context" + confidence = _coerce_confidence(fact.get("confidence"), default=0.0) + source_error = fact.get("sourceError") + if category == "correction" and isinstance(source_error, str) and source_error.strip(): + return f"- [{category} | {confidence:.2f}] {content} (avoid: {source_error.strip()})" + return f"- [{category} | {confidence:.2f}] {content}" + + +def _select_fact_lines( + ranked_facts: list[dict[str, Any]], + *, + token_budget: int, + use_tiktoken: bool, +) -> tuple[list[str], int]: + """Greedily select formatted fact lines within a *line-only* token budget. + + This function is intentionally **header-agnostic**: it counts only the + fact lines themselves (including ``\\n`` separators between lines). The + caller is responsible for reserving tokens for the ``"Facts:\\n"`` header + and any inter-section ``"\\n\\n"`` separator *before* calling this + function, and passing the remaining capacity as *token_budget*. + + Stops at the first fact that would exceed the budget so the caller's + pre-sorted order (typically confidence-descending) is preserved strictly: + a shorter lower-ranked fact can never slip ahead of a skipped + higher-ranked one. + + Args: + ranked_facts: Facts pre-sorted by the caller's preferred ranking. + token_budget: Maximum tokens available for fact lines only. + use_tiktoken: Whether to use tiktoken for counting. + + Returns: + ``(selected_lines, consumed_tokens)`` — *consumed_tokens* is the + exact token cost of the returned lines (including inter-line + ``\\n`` separators, but *not* a leading header). + """ + lines: list[str] = [] + consumed = 0 + for fact in ranked_facts: + formatted = _format_fact_line(fact) + if formatted is None: + continue + line_text = ("\n" + formatted) if lines else formatted + line_tokens = _count_tokens(line_text, use_tiktoken=use_tiktoken) + if consumed + line_tokens > token_budget: + break + lines.append(formatted) + consumed += line_tokens + return lines, consumed + + +def _fallback_format_facts( + valid_facts: list[dict[str, Any]], + *, + preceding_section_cost: int, + max_tokens: int, + use_tiktoken: bool, +) -> tuple[str, list[str]] | tuple[None, None]: + """Confidence-only ranking used when the primary path raises an exception. + + Returns a tuple ``(section_text, fact_lines)`` where ``section_text`` is the + formatted ``"Facts:\\n..."`` section string (without any leading inter-section + separator — the caller owns that), and ``fact_lines`` are the individual lines + that make up the facts block. Both elements are ``None`` if no facts survive. + + Returning the lines separately lets the caller track them for the + structure-aware safety truncation so fallback facts enjoy the same + protected-suffix treatment as facts emitted by the primary path. + + *valid_facts* is the already-filtered fact list built by the primary path so + the fallback does not redo validation work. *preceding_section_cost* is the + tokens already consumed by user-context / history sections (used to derive + the remaining budget). + """ + ranked = sorted(valid_facts, key=lambda f: _coerce_confidence(f.get("confidence"), default=0.0), reverse=True) + + header = "Facts:\n" + overhead = _count_tokens(header, use_tiktoken=use_tiktoken) + line_budget = max_tokens - preceding_section_cost - overhead + if line_budget <= 0: + return None, None + + lines, _ = _select_fact_lines(ranked, token_budget=line_budget, use_tiktoken=use_tiktoken) + if not lines: + return None, None + return header + "\n".join(lines), lines + + +def format_memory_for_injection( + memory_data: dict[str, Any], + max_tokens: int = 2000, + *, + use_tiktoken: bool = True, + guaranteed_categories: list[str] | None = None, + guaranteed_token_budget: int = 500, +) -> str: """Format memory data for injection into system prompt. Args: @@ -325,6 +433,18 @@ def format_memory_for_injection(memory_data: dict[str, Any], max_tokens: int = 2 use_tiktoken: When ``False``, all token counting uses the network-free character-based estimate instead of tiktoken (see ``memory.token_counting`` config). Defaults to ``True``. + guaranteed_categories: Fact categories that must always be injected + regardless of the regular token budget. These facts draw from a + separate *guaranteed_token_budget*. When ``None`` or empty, all + facts compete for the same budget (original behaviour). + guaranteed_token_budget: Token ceiling for the guaranteed section. + In the common case the guaranteed lines *displace* regular lines + within *max_tokens* (the total output stays ≤ ``max_tokens``); + the budget becomes truly additive only when the guaranteed lines + alone would push the assembled output past *max_tokens*, at which + point the safety-truncation ceiling is raised to + ``max_tokens + guaranteed_actual_usage`` to protect them. + Ignored when *guaranteed_categories* is ``None`` or empty. Returns: Formatted memory string for system prompt injection. @@ -332,7 +452,16 @@ def format_memory_for_injection(memory_data: dict[str, Any], max_tokens: int = 2 if not memory_data: return "" - sections = [] + # Reject a bare string explicitly: iterating a ``str`` yields single + # characters, which would silently produce a meaningless frozenset of + # letters and turn the guarantee off without any warning. Config-layer + # callers go through Pydantic (which enforces ``list[str]``), so this + # only guards the public helper surface. + if isinstance(guaranteed_categories, str): + raise TypeError("guaranteed_categories must be an iterable of strings, not a bare str") + effective_guaranteed: frozenset[str] = frozenset(c.strip() for c in guaranteed_categories if isinstance(c, str) and c.strip()) if guaranteed_categories else frozenset() + + sections: list[str] = [] # Format user context user_data = memory_data.get("user", {}) @@ -374,67 +503,181 @@ def format_memory_for_injection(memory_data: dict[str, Any], max_tokens: int = 2 if history_sections: sections.append("History:\n" + "\n".join(f"- {s}" for s in history_sections)) - # Format facts (sorted by confidence; include as many as token budget allows) + # ── Facts ──────────────────────────────────────────────────────────────── + # + # Design notes + # ~~~~~~~~~~~~ + # • A single ``"Facts:\\n"`` header is emitted at most once. + # • Guaranteed-category facts are selected first from their own + # *guaranteed_token_budget* and placed at the front of the Facts block, + # so they cannot be evicted by regular facts. In the common case the + # total output still fits within *max_tokens* (guaranteed lines displace + # regular ones); the budget becomes truly additive only when the + # guaranteed lines alone push the output past *max_tokens*, in which + # case the safety-truncation ceiling is raised accordingly. + # • Regular facts draw from *max_tokens* only. + # • All token accounting (header, separators, lines) is performed here + # in the caller; the ``_select_fact_lines`` helper is header-agnostic. + # • When the primary path raises any exception, ``_fallback_format_facts`` + # performs a single-pass confidence-only ranking. facts_data = memory_data.get("facts", []) + guaranteed_line_tokens = 0 # used later for the effective truncation limit if isinstance(facts_data, list) and facts_data: - ranked_facts = sorted( - (f for f in facts_data if isinstance(f, dict) and isinstance(f.get("content"), str) and f.get("content").strip()), - key=lambda fact: _coerce_confidence(fact.get("confidence"), default=0.0), - reverse=True, - ) - - # Compute token count for existing sections once, then account - # incrementally for each fact line to avoid full-string re-tokenization. + # Token cost of sections built above (user context, history). base_text = "\n\n".join(sections) base_tokens = _count_tokens(base_text, use_tiktoken=use_tiktoken) if base_text else 0 - # Account for the separator between existing sections and the facts section. + + # Pre-filter valid facts *before* entering the try so the except + # path can pass the same list straight into the fallback without + # redoing validation work on the hot prompt-injection path. + valid_facts = [f for f in facts_data if isinstance(f, dict) and isinstance(f.get("content"), str) and f.get("content", "").strip()] + + # Initialise the facts-block markers *before* the try so the + # structure-aware truncation at the bottom of the function can + # reason about them regardless of whether the primary path or + # the except/fallback path produced the final Facts section. facts_header = "Facts:\n" - separator_tokens = _count_tokens("\n\n" + facts_header, use_tiktoken=use_tiktoken) if base_text else _count_tokens(facts_header, use_tiktoken=use_tiktoken) - running_tokens = base_tokens + separator_tokens + all_fact_lines: list[str] = [] - fact_lines: list[str] = [] - for fact in ranked_facts: - content_value = fact.get("content") - if not isinstance(content_value, str): - continue - content = content_value.strip() - if not content: - continue - category = str(fact.get("category", "context")).strip() or "context" - confidence = _coerce_confidence(fact.get("confidence"), default=0.0) - source_error = fact.get("sourceError") - if category == "correction" and isinstance(source_error, str) and source_error.strip(): - line = f"- [{category} | {confidence:.2f}] {content} (avoid: {source_error.strip()})" + try: + # Partition valid facts into guaranteed vs regular groups. + # Use the *raw* category field (no ``or "context"`` default) so + # a category-less legacy fact is never silently promoted into + # a guaranteed pool whose operator configured + # ``guaranteed_categories=["context"]``. Missing-category facts + # always fall through to the regular path. + def _confidence_key(fact: dict[str, Any]) -> float: + return _coerce_confidence(fact.get("confidence"), default=0.0) + + if effective_guaranteed: + + def _category_match(fact: dict[str, Any]) -> bool: + raw = fact.get("category") + if not isinstance(raw, str): + return False + cat = raw.strip() + return bool(cat) and cat in effective_guaranteed + + guaranteed = sorted( + [f for f in valid_facts if _category_match(f)], + key=_confidence_key, + reverse=True, + ) + regular = sorted( + [f for f in valid_facts if not _category_match(f)], + key=_confidence_key, + reverse=True, + ) else: - line = f"- [{category} | {confidence:.2f}] {content}" + guaranteed = [] + regular = sorted(valid_facts, key=_confidence_key, reverse=True) - # Each additional line is preceded by a newline (except the first). - line_text = ("\n" + line) if fact_lines else line - line_tokens = _count_tokens(line_text, use_tiktoken=use_tiktoken) + # ── Phase 1: select guaranteed lines ────────────────────────── + header_cost = _count_tokens(facts_header, use_tiktoken=use_tiktoken) - if running_tokens + line_tokens <= max_tokens: - fact_lines.append(line) - running_tokens += line_tokens - else: - break + guaranteed_lines: list[str] = [] + if guaranteed: + guaranteed_line_budget = guaranteed_token_budget + guaranteed_lines, guaranteed_line_tokens = _select_fact_lines( + guaranteed, + token_budget=guaranteed_line_budget, + use_tiktoken=use_tiktoken, + ) - if fact_lines: - sections.append("Facts:\n" + "\n".join(fact_lines)) + # ── Phase 2: select regular lines ──────────────────────────── + # Regular facts compete for *max_tokens* (the main budget). + # Subtract everything already accounted for: + # base sections + inter-section separator + header + # + guaranteed lines + the inter-group ``\n`` that joins the + # regular block to the guaranteed block (when both are present). + regular_lines: list[str] = [] + if regular: + inter_group_newline_tokens = _count_tokens("\n", use_tiktoken=use_tiktoken) if guaranteed_lines else 0 + used_before_regular = base_tokens + header_cost + guaranteed_line_tokens + inter_group_newline_tokens + regular_line_budget = max_tokens - used_before_regular + if regular_line_budget > 0: + regular_lines, _ = _select_fact_lines( + regular, + token_budget=regular_line_budget, + use_tiktoken=use_tiktoken, + ) + + # ── Emit a single "Facts:" section ─────────────────────────── + # Leading inter-section separator is NOT embedded here; the + # final ``"\n\n".join(sections)`` is the single source of truth + # for section-to-section spacing, preventing the prior + # double-``\n\n`` bug. + all_fact_lines = guaranteed_lines + regular_lines + if all_fact_lines: + section_text = facts_header + "\n".join(all_fact_lines) + sections.append(section_text) + + except Exception: + # ── Fallback: confidence-only ranking, single budget ───────── + # Any unexpected error in the partition / guaranteed path must + # not prevent memory injection entirely. Fall back to the + # original single-pass confidence ranking. Re-use the + # pre-filtered ``valid_facts`` so we don't redo validation work + # on the hot fallback path. + logger.warning( + "Memory injection: guaranteed-category path failed, falling back to confidence-only ranking", + exc_info=True, + ) + fallback, fallback_lines = _fallback_format_facts( + valid_facts, + preceding_section_cost=base_tokens, + max_tokens=max_tokens, + use_tiktoken=use_tiktoken, + ) + if fallback: + sections.append(fallback) + # Surface the fallback's lines to ``all_fact_lines`` so the + # structure-aware truncation below treats fallback facts as a + # protected suffix too. Without this, a large user-context + # prefix could silently clip fallback facts via the original + # prefix-cut. + all_fact_lines = fallback_lines if not sections: return "" result = "\n\n".join(sections) - # Use accurate token counting with tiktoken (or the char-based estimate - # when use_tiktoken is False). token_count = _count_tokens(result, use_tiktoken=use_tiktoken) - if token_count > max_tokens: - # Truncate to fit within token limit - # Estimate characters to remove based on token ratio - char_per_token = len(result) / token_count - target_chars = int(max_tokens * char_per_token * 0.95) # 95% to leave margin - result = result[:target_chars] + "\n..." + effective_limit = max_tokens + guaranteed_line_tokens + if token_count > effective_limit: + # Structure-aware truncation: the ``Facts:\n...`` block is treated as + # a *protected suffix* so guaranteed-category facts — the very facts + # this PR exists to preserve — can never be silently discarded by a + # prefix-cut on overflow. Only the preceding (user-context / history) + # sections are eligible for truncation; if they alone exceed the + # budget available after reserving the Facts block, they are clipped + # from the tail. When *guaranteed_line_tokens* is zero (no + # guaranteed categories configured or no facts survived), the + # equation collapses to the original prefix-truncation against + # ``max_tokens``, so backward compatibility is preserved. + facts_block = (facts_header + "\n".join(all_fact_lines)) if all_fact_lines else "" + facts_block_tokens = _count_tokens(facts_block, use_tiktoken=use_tiktoken) + separator_tokens = _count_tokens("\n\n", use_tiktoken=use_tiktoken) + budget_for_non_facts = max( + 0, + effective_limit - facts_block_tokens - (separator_tokens if facts_block else 0), + ) + + # Build the preceding (non-facts) portion from *sections* excluding + # the trailing Facts block. + preceding_sections = sections[:-1] if all_fact_lines else sections + preceding = "\n\n".join(preceding_sections) + + if preceding: + preceding_tokens = _count_tokens(preceding, use_tiktoken=use_tiktoken) + if preceding_tokens > budget_for_non_facts: + char_per_token = len(preceding) / max(preceding_tokens, 1) + target_chars = int(budget_for_non_facts * char_per_token * 0.95) + preceding = preceding[:target_chars].rstrip() + "\n..." + result = (preceding + "\n\n" + facts_block) if facts_block else preceding + else: + result = facts_block return result diff --git a/backend/packages/harness/deerflow/client.py b/backend/packages/harness/deerflow/client.py index b0c0b8b13..5ecf1b15a 100644 --- a/backend/packages/harness/deerflow/client.py +++ b/backend/packages/harness/deerflow/client.py @@ -1142,6 +1142,8 @@ class DeerFlowClient: "injection_enabled": config.injection_enabled, "max_injection_tokens": config.max_injection_tokens, "token_counting": config.token_counting, + "guaranteed_categories": config.guaranteed_categories, + "guaranteed_token_budget": config.guaranteed_token_budget, } def get_memory_status(self) -> dict: diff --git a/backend/packages/harness/deerflow/config/memory_config.py b/backend/packages/harness/deerflow/config/memory_config.py index 9a2c12952..f14e777e9 100644 --- a/backend/packages/harness/deerflow/config/memory_config.py +++ b/backend/packages/harness/deerflow/config/memory_config.py @@ -73,6 +73,31 @@ class MemoryConfig(BaseModel): "CJK-aware character-based estimate and never touches tiktoken." ), ) + guaranteed_categories: list[str] = Field( + default_factory=lambda: ["correction"], + description=( + "Fact categories that are always injected into the prompt regardless " + "of the regular token budget. These facts are allocated from a " + "separate reserved budget (``guaranteed_token_budget``). " + "This ensures high-value facts such as explicit user corrections " + "are never silently dropped when the token budget is tight." + ), + ) + guaranteed_token_budget: int = Field( + default=500, + ge=50, + le=2000, + description=( + "Token ceiling for guaranteed-category facts. " + "Guaranteed facts are selected first from this budget and placed at " + "the front of the Facts block so they cannot be evicted by regular " + "facts. In the common case the total output still fits within " + "``max_injection_tokens`` (guaranteed lines displace regular ones); " + "the budget becomes additive only when guaranteed lines alone push " + "the output past ``max_injection_tokens``, in which case the " + "safety-truncation ceiling is raised accordingly." + ), + ) # Global configuration instance diff --git a/backend/tests/test_lead_agent_prompt.py b/backend/tests/test_lead_agent_prompt.py index e72e180e7..82582acd7 100644 --- a/backend/tests/test_lead_agent_prompt.py +++ b/backend/tests/test_lead_agent_prompt.py @@ -204,7 +204,14 @@ def test_get_memory_context_uses_explicit_app_config_without_global_config(monke captured["user_id"] = user_id return {"facts": []} - def fake_format_memory_for_injection(memory_data, *, max_tokens, use_tiktoken=True): + def fake_format_memory_for_injection( + memory_data, + *, + max_tokens, + use_tiktoken=True, + guaranteed_categories=None, + guaranteed_token_budget=500, + ): captured["memory_data"] = memory_data captured["max_tokens"] = max_tokens captured["use_tiktoken"] = use_tiktoken diff --git a/backend/tests/test_memory_prompt_injection.py b/backend/tests/test_memory_prompt_injection.py index c2b58b61f..b103c83b0 100644 --- a/backend/tests/test_memory_prompt_injection.py +++ b/backend/tests/test_memory_prompt_injection.py @@ -2,6 +2,8 @@ import math +import pytest + from deerflow.agents.memory.prompt import _coerce_confidence, format_memory_for_injection @@ -173,3 +175,482 @@ def test_format_memory_includes_long_term_background() -> None: assert "Background: Core expertise in distributed systems" in result assert "Recent: Recent activity summary" in result assert "Earlier: Earlier context summary" in result + + +# --------------------------------------------------------------------------- +# Guaranteed-category injection tests +# --------------------------------------------------------------------------- + + +def test_guaranteed_correction_injected_when_budget_tight(monkeypatch) -> None: + """Correction facts must be injected even when the regular budget is exhausted.""" + # Deterministic char-based counting. + monkeypatch.setattr( + "deerflow.agents.memory.prompt._count_tokens", + lambda text, encoding_name="cl100k_base", *, use_tiktoken=True: len(text), + ) + + memory_data = { + "user": {}, + "history": {}, + "facts": [ + # Many high-confidence regular facts that will eat the budget. + {"content": "Regular fact A " * 20, "category": "knowledge", "confidence": 0.95}, + {"content": "Regular fact B " * 20, "category": "knowledge", "confidence": 0.90}, + {"content": "Regular fact C " * 20, "category": "knowledge", "confidence": 0.85}, + # A correction fact with lower confidence. + {"content": "Use make dev, not npm start", "category": "correction", "confidence": 0.7}, + ], + } + + # Tight budget that cannot fit all facts. + result = format_memory_for_injection( + memory_data, + max_tokens=200, + guaranteed_categories=["correction"], + guaranteed_token_budget=100, + ) + + # The correction fact MUST appear regardless of budget pressure. + assert "Use make dev, not npm start" in result + + +def test_guaranteed_facts_sorted_by_confidence() -> None: + """Guaranteed facts should be sorted by confidence descending.""" + memory_data = { + "user": {}, + "history": {}, + "facts": [ + {"content": "Low conf correction", "category": "correction", "confidence": 0.6}, + {"content": "High conf correction", "category": "correction", "confidence": 0.95}, + {"content": "Regular fact", "category": "knowledge", "confidence": 0.8}, + ], + } + + result = format_memory_for_injection( + memory_data, + max_tokens=2000, + guaranteed_categories=["correction"], + guaranteed_token_budget=500, + ) + + assert "High conf correction" in result + assert "Low conf correction" in result + assert result.index("High conf correction") < result.index("Low conf correction") + + +def test_guaranteed_budget_isolation() -> None: + """Guaranteed facts draw from their own budget, not the regular budget.""" + memory_data = { + "user": {}, + "history": {}, + "facts": [ + {"content": "Correction one", "category": "correction", "confidence": 0.9}, + {"content": "Regular knowledge", "category": "knowledge", "confidence": 0.8}, + ], + } + + result = format_memory_for_injection( + memory_data, + max_tokens=2000, + guaranteed_categories=["correction"], + guaranteed_token_budget=500, + ) + + # Both facts should appear (separate budgets). + assert "Correction one" in result + assert "Regular knowledge" in result + + +def test_no_guaranteed_categories_backward_compatible() -> None: + """When guaranteed_categories is None, behaviour matches the original.""" + memory_data = { + "user": {}, + "history": {}, + "facts": [ + {"content": "High conf", "category": "knowledge", "confidence": 0.95}, + {"content": "Low conf", "category": "context", "confidence": 0.4}, + ], + } + + # No guaranteed_categories passed → original behaviour. + result = format_memory_for_injection(memory_data, max_tokens=2000) + + assert "High conf" in result + assert result.index("High conf") < result.index("Low conf") + + +def test_empty_guaranteed_list_backward_compatible() -> None: + """An empty guaranteed_categories list should behave like None.""" + memory_data = { + "user": {}, + "history": {}, + "facts": [ + {"content": "Correction fact", "category": "correction", "confidence": 0.9}, + {"content": "Regular fact", "category": "knowledge", "confidence": 0.8}, + ], + } + + result = format_memory_for_injection( + memory_data, + max_tokens=2000, + guaranteed_categories=[], + ) + + assert "Correction fact" in result + assert "Regular fact" in result + + +def test_fallback_on_ranking_error(monkeypatch) -> None: + """If the guaranteed path raises, fall back to confidence-only ranking.""" + + memory_data = { + "user": {}, + "history": {}, + "facts": [ + {"content": "Fact A", "category": "knowledge", "confidence": 0.9}, + {"content": "Fact B", "category": "correction", "confidence": 0.8}, + ], + } + + # Force _select_fact_lines to raise on the *first* call (the guaranteed + # path) but succeed on subsequent calls (the fallback path). + call_count = {"n": 0} + prompt_module = __import__("deerflow.agents.memory.prompt", fromlist=["_select_fact_lines"]) + original_select = prompt_module._select_fact_lines + + def flaky_select(*args, **kwargs): + call_count["n"] += 1 + if call_count["n"] == 1: + raise RuntimeError("simulated error in guaranteed path") + return original_select(*args, **kwargs) + + monkeypatch.setattr( + "deerflow.agents.memory.prompt._select_fact_lines", + flaky_select, + ) + + result = format_memory_for_injection( + memory_data, + max_tokens=2000, + guaranteed_categories=["correction"], + guaranteed_token_budget=500, + ) + + # Both facts should still appear via the fallback path. + assert "Fact A" in result + assert "Fact B" in result + + +def test_guaranteed_respects_its_own_budget_limit(monkeypatch) -> None: + """Even guaranteed facts are capped by guaranteed_token_budget.""" + monkeypatch.setattr( + "deerflow.agents.memory.prompt._count_tokens", + lambda text, encoding_name="cl100k_base", *, use_tiktoken=True: len(text), + ) + + # Many correction facts that together exceed the guaranteed budget. + # Formatted line example: "- [correction | 0.95] CorrA xxxxxxxxxxxxxxxx" + # Each line is ~50 chars; with "Facts:\n" header (7 chars), two lines + # need ~107 chars, exceeding the 80-char guaranteed budget. + memory_data = { + "user": {}, + "history": {}, + "facts": [ + {"content": "CorrA " + "x" * 20, "category": "correction", "confidence": 0.95}, + {"content": "CorrB " + "x" * 20, "category": "correction", "confidence": 0.90}, + {"content": "CorrC " + "x" * 20, "category": "correction", "confidence": 0.85}, + {"content": "Short regular", "category": "knowledge", "confidence": 0.8}, + ], + } + + result = format_memory_for_injection( + memory_data, + max_tokens=2000, + guaranteed_categories=["correction"], + guaranteed_token_budget=80, # Small guaranteed budget — fits 1 fact line only. + ) + + # At least the highest-confidence correction should appear. + assert "CorrA" in result + # The regular fact should also appear (it has its own budget). + assert "Short regular" in result + + +def test_guaranteed_fact_with_source_error_rendered() -> None: + """Guaranteed correction facts should still render sourceError.""" + memory_data = { + "facts": [ + { + "content": "Use uv, not pip.", + "category": "correction", + "confidence": 0.95, + "sourceError": "Agent suggested pip install.", + }, + {"content": "Likes Python", "category": "preference", "confidence": 0.8}, + ], + } + + result = format_memory_for_injection( + memory_data, + max_tokens=2000, + guaranteed_categories=["correction"], + guaranteed_token_budget=500, + ) + + assert "Use uv, not pip." in result + assert "avoid: Agent suggested pip install." in result + assert "Likes Python" in result + + +def test_single_facts_header_when_both_guaranteed_and_regular() -> None: + """When both guaranteed and regular facts exist, emit exactly one 'Facts:' header.""" + memory_data = { + "user": {"workContext": {"summary": "Dev"}}, # non-empty preceding section + "history": {}, + "facts": [ + {"content": "Correction fact", "category": "correction", "confidence": 0.95}, + {"content": "Knowledge fact", "category": "knowledge", "confidence": 0.80}, + ], + } + + result = format_memory_for_injection( + memory_data, + max_tokens=2000, + guaranteed_categories=["correction"], + guaranteed_token_budget=500, + ) + + # Exactly one "Facts:" header. + assert result.count("Facts:") == 1, f"Expected exactly one 'Facts:' header, got:\n{result}" + # Both facts appear under the single header. + assert "Correction fact" in result + assert "Knowledge fact" in result + # Guaranteed fact comes first (higher confidence + guaranteed). + assert result.index("Correction fact") < result.index("Knowledge fact") + + +def test_strict_confidence_order_when_high_confidence_fact_overflows(monkeypatch) -> None: + """Within a single budget, a higher-confidence fact that exceeds the + remaining budget must NOT be skipped in favour of a shorter, lower- + confidence fact ranked after it. + + This locks in the strict confidence-ordered selection semantics. + """ + monkeypatch.setattr( + "deerflow.agents.memory.prompt._count_tokens", + lambda text, encoding_name="cl100k_base", *, use_tiktoken=True: len(text), + ) + + memory_data = { + "user": {}, + "history": {}, + "facts": [ + # Higher-confidence but long enough to exceed the remaining budget. + {"content": "Long high-confidence fact " + "x" * 50, "category": "knowledge", "confidence": 0.95}, + # Lower-confidence but short — would fit if we kept scanning past + # the over-budget high-confidence fact above. + {"content": "Short low", "category": "knowledge", "confidence": 0.50}, + ], + } + + # Budget large enough only for ~one short fact, not the long one. + result = format_memory_for_injection(memory_data, max_tokens=70, guaranteed_categories=None) + + # The high-confidence fact does not fit, and the low-confidence fact + # MUST NOT slip in ahead of it. + assert "Short low" not in result, "Lower-confidence fact should not be selected when a higher-confidence fact ranked before it was skipped (strict ordering)." + + +# ── Regression tests for willem-bd's review on PR #3592 ────────────────── + + +def test_structure_aware_truncation_preserves_guaranteed_on_overflow(monkeypatch) -> None: + """[P1] When user context overflows, the trailing ``Facts:\\n...`` block + is treated as a protected suffix and only the preceding sections are + clipped — guaranteed-category facts can never be silently discarded by + a prefix-cut on overflow. + + Locks in the fix for willem-bd's P1 finding on PR #3592. + """ + monkeypatch.setattr( + "deerflow.agents.memory.prompt._count_tokens", + lambda text, encoding_name="cl100k_base", *, use_tiktoken=True: len(text), + ) + + memory_data = { + # Oversized preceding section that would otherwise push Facts past the + # effective truncation ceiling. + "user": {"workContext": {"summary": "X" * 4000}}, + "facts": [ + { + "content": "CRITICAL: never use pip", + "category": "correction", + "confidence": 1.0, + "sourceError": "pip is deprecated", + }, + {"content": "B", "category": "knowledge", "confidence": 0.5}, + ], + } + + result = format_memory_for_injection( + memory_data, + max_tokens=200, + guaranteed_categories=["correction"], + guaranteed_token_budget=500, + use_tiktoken=False, + ) + + # Guaranteed correction must survive even when preceding sections are huge. + assert "never use pip" in result, f"Guaranteed correction was silently truncated away:\n{result[-200:]}" + assert "pip is deprecated" in result + # The protected suffix shape: Facts block is at the tail. + assert result.rstrip().endswith("(avoid: pip is deprecated)") + + +def test_single_inter_section_separator_between_user_and_facts() -> None: + """[P2] Exactly one ``\\n\\n`` separator between ``User Context:`` and + ``Facts:`` — never four newlines. + + Locks in the fix for willem-bd's P2 separator finding on PR #3592. + """ + memory_data = { + "user": {"workContext": {"summary": "Python developer"}}, + "history": {}, + "facts": [ + {"content": "fact A", "category": "knowledge", "confidence": 0.9}, + { + "content": "fact B", + "category": "correction", + "confidence": 0.8, + "sourceError": "avoid X", + }, + ], + } + + result = format_memory_for_injection(memory_data, max_tokens=2000) + + assert "\n\n\n\n" not in result, f"Found four consecutive newlines between sections:\n{result[:200]!r}" + # Exactly one \n\n between User Context: and Facts:. + idx_user = result.index("User Context:") + idx_facts = result.index("Facts:") + between = result[idx_user:idx_facts] + assert between.count("\n\n") == 1, f"Expected exactly one \\n\\n between sections, got:\n{between!r}" + + +def test_bare_string_guaranteed_categories_raises_type_error() -> None: + """[P2] Passing a bare ``str`` for *guaranteed_categories* must raise + ``TypeError`` instead of silently iterating single characters and + disabling the guarantee. + + Locks in the fix for willem-bd's P2 bare-string finding on PR #3592. + """ + memory_data = { + "facts": [ + {"content": "CRITICAL", "category": "correction", "confidence": 0.8}, + ], + } + with pytest.raises(TypeError, match="iterable"): + format_memory_for_injection( + memory_data, + guaranteed_categories="correction", # type: ignore[arg-type] + ) + + +def test_categoryless_fact_not_promoted_into_guaranteed_context_pool(monkeypatch) -> None: + """[P2] A fact with a missing/empty ``category`` field is *never* + silently promoted into a ``guaranteed_categories=["context"]`` pool — + only facts with an *explicit* ``category == "context"`` qualify. + + Strategy: set a guaranteed budget tight enough to fit only the short + *explicit* ``context`` fact. If the legacy (no-category) fact were + silently promoted into the guaranteed pool, it would claim the budget + first (higher confidence) and push the explicit one out into the + regular pool where, under a tight ``max_tokens``, it would be lost. + If the fix holds, the explicit fact owns the guaranteed pool alone + and survives. + + Locks in the fix for willem-bd's P2 category-less finding on PR #3592. + """ + monkeypatch.setattr( + "deerflow.agents.memory.prompt._count_tokens", + lambda text, encoding_name="cl100k_base", *, use_tiktoken=True: len(text), + ) + + memory_data = { + "facts": [ + # Long legacy fact with NO category field. + { + "content": "legacy " + "x" * 80, + "confidence": 0.95, + }, + # Short explicit context fact. + { + "content": "explicit ctx", + "category": "context", + "confidence": 0.9, + }, + ], + } + + # Guaranteed budget sized for the short explicit fact only. + result = format_memory_for_injection( + memory_data, + max_tokens=200, + guaranteed_categories=["context"], + guaranteed_token_budget=40, + use_tiktoken=False, + ) + + # The explicit context fact must survive in the guaranteed pool. + assert "explicit ctx" in result, f"Explicit 'context' fact was evicted — legacy no-category fact was silently promoted into the guaranteed pool.\n{result!r}" + + +def test_fallback_uses_prefiltered_valid_facts(monkeypatch) -> None: + """[P2] When the primary path raises after ``valid_facts`` has been + built, the fallback operates on the pre-filtered list (no raw-content + facts leak through) and still produces a valid ``Facts:`` section. + + Locks in the fix for willem-bd's P2 fallback-duplication finding on + PR #3592. + """ + monkeypatch.setattr( + "deerflow.agents.memory.prompt._count_tokens", + lambda text, encoding_name="cl100k_base", *, use_tiktoken=True: len(text), + ) + + call_count = {"select": 0} + original_select = __import__("deerflow.agents.memory.prompt", fromlist=["_select_fact_lines"])._select_fact_lines + + def raising_select(*args, **kwargs): + call_count["select"] += 1 + if call_count["select"] == 1: + raise RuntimeError("primary path failure") + return original_select(*args, **kwargs) + + monkeypatch.setattr("deerflow.agents.memory.prompt._select_fact_lines", raising_select) + + memory_data = { + "facts": [ + {"content": "valid fact", "category": "knowledge", "confidence": 0.9}, + # Malformed: no content field — should be pre-filtered and never + # reach the fallback's ranking. + {"category": "knowledge", "confidence": 0.95}, + # Empty content — also pre-filtered. + {"content": " ", "category": "knowledge", "confidence": 0.9}, + ], + } + + result = format_memory_for_injection( + memory_data, + max_tokens=2000, + guaranteed_categories=["correction"], + use_tiktoken=False, + ) + + # Fallback kicked in and still produced the Facts section. + assert "Facts:" in result + # The valid fact survived pre-filtering and fallback ranking. + assert "valid fact" in result + # Malformed facts were pre-filtered and never rendered. + assert result.count("- [") == 1 diff --git a/config.example.yaml b/config.example.yaml index da3f02fbe..76ac63476 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -1191,6 +1191,20 @@ memory: # char - network-free CJK-aware character-based estimate; never touches # tiktoken. Slightly less precise budgeting, zero network I/O. token_counting: tiktoken + # Guaranteed injection: fact categories that bypass the regular token budget + # and draw from a reserved allowance, so high-signal corrections (e.g. + # "don't use `pip`, use `uv`") survive even when the budget is tight. + # guaranteed_categories - list of fact categories to guarantee. Pass [] to + # disable; defaults to ["correction"]. + # guaranteed_token_budget - token ceiling for guaranteed facts. In the + # common case the total injection stays within ``max_injection_tokens`` + # (guaranteed lines displace regular ones); the allowance becomes + # additive only when guaranteed lines alone would overflow + # ``max_injection_tokens``, in which case the safety-truncation ceiling + # is raised accordingly. + guaranteed_categories: + - correction + guaranteed_token_budget: 500 # ============================================================================ # Custom Agent Management API