feat(memory): add guaranteed injection for correction facts with graceful fallback (#3592)

* 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 <willem.jiang@gmail.com>

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),
     )

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():

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
 

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:

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

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

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

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