diff --git a/agent/context_compressor.py b/agent/context_compressor.py index e9471aaa4335..cf7ccdfcfd36 100644 --- a/agent/context_compressor.py +++ b/agent/context_compressor.py @@ -32,6 +32,12 @@ call_llm, ) from agent.context_engine import ContextEngine, sanitize_memory_context +from agent.context_compressor_message_helpers import ( + _DB_PERSISTED_MARKER, + _fresh_compaction_message_copy, + _strip_persistence_markers, + _template_visible_role, +) from agent.error_classifier import FailoverReason, classify_api_error from agent.model_metadata import ( MINIMUM_CONTEXT_LENGTH, @@ -155,7 +161,6 @@ def _is_summary_access_or_quota_error(exc: Exception) -> bool: # micro markers: a batch marker's content is NOT contained in the micro # rolling summary, so dropping or rewriting one destroys history. MICRO_COMPACT_MARKER_KEY = "_micro_compact_marker" -_DB_PERSISTED_MARKER = "_db_persisted" PROACTIVE_PRUNE_REARM_MODEL_CONFIG_KEY = "_proactive_prune_rearm_tokens" _NO_USER_TASK_SENTINEL = "None. This session contains no user-authored turns." @@ -170,73 +175,6 @@ def _is_summary_access_or_quota_error(exc: Exception) -> bool: ) -def _fresh_compaction_message_copy(msg: Dict[str, Any]) -> Dict[str, Any]: - """Copy a message for compaction assembly without persistence markers. - - Live cached-gateway transcripts stamp ``_db_persisted`` during incremental - flushes. Shallow ``.copy()`` propagates that marker into the post-rotation - compressed list, so ``_flush_messages_to_session_db`` skips every row when - writing to the new child session (#57491). - - This strips at the copy site (clearest intent, and cheap), but the - authoritative guarantee is the single terminal sweep in ``compress()`` - (``_strip_persistence_markers``): no message may leave ``compress()`` - carrying ``_db_persisted`` regardless of how many intermediate copy sites - a future refactor adds. - """ - fresh = msg.copy() - fresh.pop(_DB_PERSISTED_MARKER, None) - return fresh - - -def _template_visible_role(message: Any) -> Optional[str]: - """Role as counted by strict chat-template alternation checks. - - Mistral-family templates (Devstral, Mistral Small 3.x, Magistral) - enforce user/assistant alternation at render time but EXEMPT the tool - flow from the check: ``tool`` results and assistant messages carrying - ``tool_calls`` are skipped. A summary role chosen against the *literal* - neighbouring roles can therefore still violate alternation as the - template sees it. The canonical failure: the protected head ends - ``[user, assistant(tool_calls), tool]``, so the literal last role is - ``tool`` and the summary is pinned to ``role="user"`` -- but the last - role the template counts is ``user``, the template sees user -> user, - and llama.cpp / Mistral-hosted backends reject the ENTIRE request with - a Jinja alternation error (HTTP 500). Because the summary persists in - the stored conversation, every retry replays the same poisoned history - and the session is unrecoverable. - - Returns ``None`` for messages the alternation check skips. - """ - if not isinstance(message, dict): - return None - role = message.get("role") - if role == "tool": - return None - if role == "assistant" and message.get("tool_calls"): - return None - return role - - -def _strip_persistence_markers(messages: List[Dict[str, Any]]) -> None: - """Enforce the compaction invariant: no assembled message carries a - session-store persistence marker. - - ``compress()`` copies protected head/tail messages out of the live - cached-gateway transcript, which stamps ``_db_persisted`` on every message - over the life of the session. If any copied dict keeps that marker, the - rotation flush to the child session skips it and the compacted transcript is - lost from ``state.db`` (#57491). Stripping at each copy site is necessary - but *positional* — a copy site added after the assembly loops would re-leak. - This single terminal sweep makes the guarantee structural instead: run it - once on the fully-assembled list so the invariant holds no matter where the - copies happened. Mutates in place (the dicts are compaction-local copies). - """ - for msg in messages: - if isinstance(msg, dict): - msg.pop(_DB_PERSISTED_MARKER, None) - - # Appended to every standalone summary message (and to the merged-into-tail # prefix) so the model has an unambiguous "summary ends here" boundary. # Without it, weak models read the verbatim "## Active Task" quote as fresh diff --git a/agent/context_compressor_message_helpers.py b/agent/context_compressor_message_helpers.py new file mode 100644 index 000000000000..6e622974687b --- /dev/null +++ b/agent/context_compressor_message_helpers.py @@ -0,0 +1,72 @@ +"""Message-role and persistence-marker helpers for context compaction.""" + +from typing import Any, Dict, List, Optional + +_DB_PERSISTED_MARKER = "_db_persisted" + + +def _fresh_compaction_message_copy(msg: Dict[str, Any]) -> Dict[str, Any]: + """Copy a message for compaction assembly without persistence markers. + + Live cached-gateway transcripts stamp ``_db_persisted`` during incremental + flushes. Shallow ``.copy()`` propagates that marker into the post-rotation + compressed list, so ``_flush_messages_to_session_db`` skips every row when + writing to the new child session (#57491). + + This strips at the copy site (clearest intent, and cheap), but the + authoritative guarantee is the single terminal sweep in ``compress()`` + (``_strip_persistence_markers``): no message may leave ``compress()`` + carrying ``_db_persisted`` regardless of how many intermediate copy sites + a future refactor adds. + """ + fresh = msg.copy() + fresh.pop(_DB_PERSISTED_MARKER, None) + return fresh + + +def _template_visible_role(message: Any) -> Optional[str]: + """Role as counted by strict chat-template alternation checks. + + Mistral-family templates (Devstral, Mistral Small 3.x, Magistral) + enforce user/assistant alternation at render time but EXEMPT the tool + flow from the check: ``tool`` results and assistant messages carrying + ``tool_calls`` are skipped. A summary role chosen against the *literal* + neighbouring roles can therefore still violate alternation as the + template sees it. The canonical failure: the protected head ends + ``[user, assistant(tool_calls), tool]``, so the literal last role is + ``tool`` and the summary is pinned to ``role="user"`` -- but the last + role the template counts is ``user``, the template sees user -> user, + and llama.cpp / Mistral-hosted backends reject the ENTIRE request with + a Jinja alternation error (HTTP 500). Because the summary persists in + the stored conversation, every retry replays the same poisoned history + and the session is unrecoverable. + + Returns ``None`` for messages the alternation check skips. + """ + if not isinstance(message, dict): + return None + role = message.get("role") + if role == "tool": + return None + if role == "assistant" and message.get("tool_calls"): + return None + return role + + +def _strip_persistence_markers(messages: List[Dict[str, Any]]) -> None: + """Enforce the compaction invariant: no assembled message carries a + session-store persistence marker. + + ``compress()`` copies protected head/tail messages out of the live + cached-gateway transcript, which stamps ``_db_persisted`` on every message + over the life of the session. If any copied dict keeps that marker, the + rotation flush to the child session skips it and the compacted transcript is + lost from ``state.db`` (#57491). Stripping at each copy site is necessary + but *positional* — a copy site added after the assembly loops would re-leak. + This single terminal sweep makes the guarantee structural instead: run it + once on the fully-assembled list so the invariant holds no matter where the + copies happened. Mutates in place (the dicts are compaction-local copies). + """ + for msg in messages: + if isinstance(msg, dict): + msg.pop(_DB_PERSISTED_MARKER, None) diff --git a/tests/agent/test_context_compressor_message_helpers_seam.py b/tests/agent/test_context_compressor_message_helpers_seam.py new file mode 100644 index 000000000000..37adbc39d184 --- /dev/null +++ b/tests/agent/test_context_compressor_message_helpers_seam.py @@ -0,0 +1,207 @@ +"""Seam contracts for the message-marker helper extraction.""" + +from __future__ import annotations + +import ast +import hashlib +import inspect +import json +import subprocess +import sys +from pathlib import Path +from unittest.mock import Mock, patch + +import agent.context_compressor as owner +import pytest + + +ROOT = Path(__file__).parents[2] +EXPECTED_BYTES = 3204 +EXPECTED_SHA256 = "8c6860876e9a511927f9eecffce5bb878179d590218df284edd48c3cf3adc31d" +HELPER_NAMES = { + "_fresh_compaction_message_copy", + "_template_visible_role", + "_strip_persistence_markers", +} + + +@pytest.fixture() +def extracted(): + import agent.context_compressor_message_helpers as helpers + + return helpers + + +def test_leaf_source_is_cycle_free_exact_move_and_only_definition_owner(extracted): + source = inspect.getsource(extracted) + helper_tree = ast.parse(source) + owner_tree = ast.parse( + (ROOT / "agent" / "context_compressor.py").read_text(encoding="utf-8") + ) + + imported_modules = { + alias.name + for node in ast.walk(helper_tree) + if isinstance(node, (ast.Import, ast.ImportFrom)) + for alias in node.names + } + assert "agent.context_compressor" not in imported_modules + + owner_definitions = { + node.name + for node in owner_tree.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + } + helper_definitions = { + node.name + for node in helper_tree.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + } + assert HELPER_NAMES.isdisjoint(owner_definitions) + assert HELPER_NAMES <= helper_definitions + + helper_bytes = ( + ROOT / "agent" / "context_compressor_message_helpers.py" + ).read_bytes() + approved = helper_bytes[helper_bytes.index(b"def _fresh_compaction_message_copy") :] + assert len(approved) == EXPECTED_BYTES + assert hashlib.sha256(approved).hexdigest() == EXPECTED_SHA256 + + +def test_original_module_reexports_extracted_function_objects(extracted): + for name in ( + "_fresh_compaction_message_copy", + "_template_visible_role", + "_strip_persistence_markers", + ): + owner_value = getattr(owner, name) + extracted_value = getattr(extracted, name) + assert owner_value is extracted_value + assert owner_value.__module__ == "agent.context_compressor_message_helpers" + + +def test_original_module_reexports_marker_constant(extracted): + assert owner._DB_PERSISTED_MARKER == extracted._DB_PERSISTED_MARKER + assert owner._DB_PERSISTED_MARKER == "_db_persisted" + + +def test_fresh_copy_is_distinct_and_preserves_nested_identity(extracted): + nested = {"tool_calls": [{"id": "call-1"}]} + source = { + "role": "assistant", + "content": "done", + "metadata": nested, + "_db_persisted": True, + } + + copied = extracted._fresh_compaction_message_copy(source) + + assert copied is not source + assert copied["metadata"] is nested + assert "_db_persisted" not in copied + assert source["_db_persisted"] is True + + +@pytest.mark.parametrize( + ("message", "expected"), + [ + (None, None), + ("not a message", None), + ({"role": "tool", "content": "result"}, None), + ({"role": "assistant", "tool_calls": [{"id": "c"}]}, None), + ({"role": "assistant", "tool_calls": []}, "assistant"), + ({"role": "user"}, "user"), + ({"role": "system"}, "system"), + ], +) +def test_template_visible_role_contract(extracted, message, expected): + assert extracted._template_visible_role(message) == expected + + +def test_strip_markers_mutates_dicts_only_and_preserves_nested_data(extracted): + nested = {"keep": [1, 2, 3]} + first = {"role": "user", "metadata": nested, "_db_persisted": True} + second = {"role": "assistant", "content": "ok"} + messages = [first, "foreign-row", second] + + result = extracted._strip_persistence_markers(messages) + + assert result is None + assert "_db_persisted" not in first + assert first["metadata"] is nested + assert messages[1] == "foreign-row" + assert second == {"role": "assistant", "content": "ok"} + + +def test_original_module_patch_authority_survives_extraction(monkeypatch): + sentinel = {"role": "user", "content": "patched"} + patched = lambda _message: sentinel # noqa: E731 + + monkeypatch.setattr(owner, "_fresh_compaction_message_copy", patched) + + assert owner._fresh_compaction_message_copy({"role": "user"}) is sentinel + + +def test_import_orders_are_cycle_free(): + snippets = ( + "import agent.context_compressor_message_helpers as h; " + "import agent.context_compressor as o", + "import agent.context_compressor as o; " + "import agent.context_compressor_message_helpers as h", + ) + check = ( + "; import json; print(json.dumps({" + "'fresh': o._fresh_compaction_message_copy is h._fresh_compaction_message_copy, " + "'role': o._template_visible_role is h._template_visible_role, " + "'strip': o._strip_persistence_markers is h._strip_persistence_markers, " + "'marker': o._DB_PERSISTED_MARKER == h._DB_PERSISTED_MARKER}))" + ) + + for snippet in snippets: + completed = subprocess.run( + [sys.executable, "-c", snippet + check], + check=True, + capture_output=True, + text=True, + ) + assert json.loads(completed.stdout) == { + "fresh": True, + "role": True, + "strip": True, + "marker": True, + } + + +def test_compress_resolves_fresh_copy_through_original_patch_surface(): + with patch( + "agent.context_compressor.get_model_context_length", return_value=100000 + ): + compressor = owner.ContextCompressor( + model="test/model", + threshold_percent=0.85, + protect_first_n=2, + protect_last_n=2, + quiet_mode=True, + ) + _ = compressor.context_length + + messages = [ + { + "role": "user" if index % 2 == 0 else "assistant", + "content": f"m{index}", + "_db_persisted": True, + } + for index in range(10) + ] + leaking_copy = Mock(side_effect=lambda message: message.copy()) + + with ( + patch.object(owner, "_fresh_compaction_message_copy", leaking_copy), + patch( + "agent.context_compressor.call_llm", side_effect=RuntimeError("no provider") + ), + ): + result = compressor.compress(messages) + + assert leaking_copy.call_count > 0 + assert all("_db_persisted" not in message for message in result)