From 6e499eb1dd7f9260afe74b07fb2c86a79a5e35d3 Mon Sep 17 00:00:00 2001 From: b3nw Date: Sun, 28 Jun 2026 17:37:32 +0000 Subject: [PATCH] fix(session): semantic context truncate on rewind + evict agent Refs nesquena/hermes-webui#5096 Co-authored-by: b3nw --- api/routes.py | 23 ++----- api/session_ops.py | 39 ++++++++++++ tests/test_issue2914_truncation_watermark.py | 66 ++++++++++++++++++++ tests/test_truncate_session_at_keep.py | 25 ++++++++ 4 files changed, 135 insertions(+), 18 deletions(-) create mode 100644 tests/test_truncate_session_at_keep.py diff --git a/api/routes.py b/api/routes.py index 279b3ca0..9b78c21f 100644 --- a/api/routes.py +++ b/api/routes.py @@ -12227,24 +12227,9 @@ def handle_post(handler, parsed) -> bool: if keep < 0: return bad(handler, "keep_count must be non-negative") with _get_session_agent_lock(body["session_id"]): - old_msg_count = len(s.messages or []) - old_ctx_count = len(getattr(s, 'context_messages', None) or []) - s.messages = s.messages[:keep] - # Truncate context_messages in sync with messages so the agent's - # model-facing context doesn't retain rows the user removed via - # Edit / Regenerate. Without this, context_messages still contains - # the full pre-truncation history and the agent sees "deleted" - # turns on the next turn (#2914). - if isinstance(getattr(s, 'context_messages', None), list): - s.context_messages = s.context_messages[:keep] - try: - from api.session_ops import _truncation_watermark_for - s.truncation_watermark = _truncation_watermark_for(s.messages) - # Persist the original truncate cutoff. - s.truncation_boundary = s.truncation_watermark - except Exception: - s.truncation_watermark = 0.0 - s.truncation_boundary = 0.0 + from api.session_ops import truncate_session_at_keep + + old_msg_count, old_ctx_count = truncate_session_at_keep(s, keep) s.save() logger.info( "truncate %s: messages %d→%d, context_messages %d→%d, watermark=%.2f", @@ -12252,6 +12237,8 @@ def handle_post(handler, parsed) -> bool: old_ctx_count, len(getattr(s, 'context_messages', None) or []), s.truncation_watermark or 0, ) + from api.config import _evict_session_agent + _evict_session_agent(body["session_id"]) return j( handler, {"ok": True, "session": s.compact() | {"messages": s.messages}} ) diff --git a/api/session_ops.py b/api/session_ops.py index 9f33075f..755f24e3 100644 --- a/api/session_ops.py +++ b/api/session_ops.py @@ -67,6 +67,45 @@ def _truncation_watermark_for(messages): return 0.0 +def truncate_context_for_display_keep( + context_messages: list | None, + full_messages: list | None, + keep: int, +) -> list: + """Align model context with display prefix ``full_messages[:keep]``.""" + if keep <= 0: + return [] + ctx = context_messages if isinstance(context_messages, list) else [] + msgs = full_messages if isinstance(full_messages, list) else [] + if not ctx: + return [] + if len(ctx) == len(msgs): + return ctx[:keep] + if len(msgs) == 0: + return [] + prefix_len = max(0, len(ctx) - len(msgs)) + prefix = ctx[:prefix_len] + suffix = ctx[prefix_len:] + return prefix + suffix[:keep] + + +def truncate_session_at_keep(session, keep: int) -> tuple[int, int]: + """Truncate display + context; set watermark/boundary. Returns old counts.""" + full_messages = list(session.messages or []) + old_msg_count = len(full_messages) + old_ctx_count = len(getattr(session, 'context_messages', None) or []) + session.messages = full_messages[:keep] + if isinstance(getattr(session, 'context_messages', None), list): + session.context_messages = truncate_context_for_display_keep( + session.context_messages, + full_messages, + keep, + ) + session.truncation_watermark = _truncation_watermark_for(session.messages) + session.truncation_boundary = session.truncation_watermark + return old_msg_count, old_ctx_count + + def retry_last(session_id: str) -> dict[str, Any]: """Truncate the session to before the last user message, return its text. diff --git a/tests/test_issue2914_truncation_watermark.py b/tests/test_issue2914_truncation_watermark.py index 51a53ccf..462a938b 100644 --- a/tests/test_issue2914_truncation_watermark.py +++ b/tests/test_issue2914_truncation_watermark.py @@ -139,6 +139,72 @@ def test_truncate_endpoint_also_truncates_context_messages(monkeypatch, tmp_path assert loaded.truncation_watermark == 2.0 +def test_truncate_endpoint_compaction_leading_context_row(monkeypatch, tmp_path): + """Context longer than display (leading compaction row): naive [:keep] would + leave a stale tail; truncate must align suffix to display prefix (#5096 / C). + """ + import json + from io import BytesIO + from types import SimpleNamespace + + import api.models as models + import api.routes as routes + from api.models import Session + + session_dir = tmp_path / "sessions" + session_dir.mkdir(parents=True) + monkeypatch.setattr(models, "SESSION_DIR", session_dir) + monkeypatch.setattr(models, "SESSION_INDEX_FILE", session_dir / "_index.json") + models.SESSIONS.clear() + monkeypatch.setattr( + "api.config._evict_session_agent", + lambda _sid: None, + ) + + display = [ + _msg("user", "u1", 1.0, "u1"), + _msg("assistant", "a1", 2.0, "a1"), + _msg("user", "REMOVE", 3.0, "u2"), + ] + context = [ + _msg("user", "compaction-only", 0.5, "cref"), + _msg("user", "u1", 1.0, "cu1"), + _msg("assistant", "a1", 2.0, "ca1"), + _msg("user", "REMOVE", 3.0, "cu2"), + ] + session = Session( + session_id="issue5096truncatectx", + messages=display, + context_messages=context, + ) + session.save() + + body = {"session_id": "issue5096truncatectx", "keep_count": 2} + body_bytes = json.dumps(body).encode() + monkeypatch.setattr(routes, "_check_csrf", lambda handler: True) + + captured_response = {} + + def fake_j(handler, payload, status=200, extra_headers=None): + captured_response["payload"] = payload + + monkeypatch.setattr(routes, "j", fake_j) + + handler = SimpleNamespace( + headers={"Content-Length": str(len(body_bytes))}, + rfile=BytesIO(body_bytes), + ) + routes.handle_post(handler, SimpleNamespace(path="/api/session/truncate")) + + assert captured_response["payload"].get("ok") is True + loaded = Session.load("issue5096truncatectx") + assert loaded is not None + assert [m["content"] for m in loaded.messages] == ["u1", "a1"] + assert len(loaded.context_messages) == 3 + assert loaded.context_messages[0]["content"] == "compaction-only" + assert "REMOVE" not in [m["content"] for m in loaded.context_messages] + + def test_truncate_without_context_messages_truncation_leaks_to_agent(monkeypatch, tmp_path): """Prove the bug: if context_messages is NOT truncated, agent sees old rows.""" import api.models as models diff --git a/tests/test_truncate_session_at_keep.py b/tests/test_truncate_session_at_keep.py new file mode 100644 index 00000000..ade1939f --- /dev/null +++ b/tests/test_truncate_session_at_keep.py @@ -0,0 +1,25 @@ +"""truncate_session_at_keep aligns context when lengths differ (#5096 C).""" + +from api.models import Session +from api.session_ops import truncate_session_at_keep + + +def test_truncate_session_at_keep_compaction_prefix(): + s = Session( + session_id="t1", + messages=[ + {"role": "user", "content": "u1", "timestamp": 1}, + {"role": "assistant", "content": "a1", "timestamp": 2}, + {"role": "user", "content": "gone", "timestamp": 3}, + ], + context_messages=[ + {"role": "user", "content": "cref", "timestamp": 0.5}, + {"role": "user", "content": "u1", "timestamp": 1}, + {"role": "assistant", "content": "a1", "timestamp": 2}, + {"role": "user", "content": "gone", "timestamp": 3}, + ], + ) + truncate_session_at_keep(s, 2) + assert len(s.messages) == 2 + assert len(s.context_messages) == 3 + assert "gone" not in [m["content"] for m in s.context_messages] \ No newline at end of file