From 9f0864ea82750bb9be59bc5c989eee2da4af25a8 Mon Sep 17 00:00:00 2001 From: Frank Song Date: Wed, 24 Jun 2026 15:11:41 +0800 Subject: [PATCH] Fix compaction marker writeback and code redaction --- api/helpers.py | 38 ++++++++++++-- api/streaming.py | 32 +++++++++--- tests/test_compaction_marker_writeback.py | 50 +++++++++++++++++++ tests/test_issue1217_transcript_compaction.py | 7 +-- tests/test_security_redaction.py | 31 ++++++++++++ 5 files changed, 143 insertions(+), 15 deletions(-) create mode 100644 tests/test_compaction_marker_writeback.py diff --git a/api/helpers.py b/api/helpers.py index a824325d..47a63dc0 100644 --- a/api/helpers.py +++ b/api/helpers.py @@ -263,11 +263,15 @@ def _build_redact_fn(): r"|brv_[A-Za-z0-9]{10,}" # ByteRover API key r")(?![A-Za-z0-9_-])" ) - _AUTH_HDR_RE = _re.compile(r"(Authorization:\s*Bearer\s+)(\S+)", _re.IGNORECASE) + _AUTH_HDR_RE = _re.compile( + r"""(Authorization:\s*(?:Bearer|Bot)\s+)([^\s'",\]\)]+)""", + _re.IGNORECASE, + ) _ENV_RE = _re.compile( r"([A-Z0-9_]{0,50}(?:API_?KEY|TOKEN|SECRET|PASSWORD|PASSWD|CREDENTIAL|AUTH)[A-Z0-9_]{0,50})" - r"\s*=\s*(['\"]?)(\S+)\2" + r"\s*=\s*(['\"]?)([^\s'\"\]\),;]+)\2" ) + _PRIVKEY_RE = _re.compile( r"-----BEGIN[A-Z ]*PRIVATE KEY-----[\s\S]*?-----END[A-Z ]*PRIVATE KEY-----" ) @@ -275,14 +279,36 @@ def _build_redact_fn(): def _mask(token: str) -> str: return f"{token[:6]}...{token[-4:]}" if len(token) >= 18 else "***" + def _env_replacement(match) -> str: + key, quote, value = match.group(1), match.group(2), match.group(3) + if not any(ch.isalnum() for ch in value): + return match.group(0) + return f"{key}={quote}{_mask(value)}{quote}" + + _CODE_ENV_KEY_LITERAL_RE = _re.compile( + r"([A-Z0-9_]{0,50}(?:API_?KEY|TOKEN|SECRET|PASSWORD|PASSWD|CREDENTIAL|AUTH)[A-Z0-9_]{0,50}=)([\"'][)\]:,]+|[)\]:,]+)" + ) + + def _restore_code_env_key_literals(original: str, redacted: str) -> str: + if not isinstance(original, str) or not isinstance(redacted, str): + return redacted + restored = redacted + for match in _CODE_ENV_KEY_LITERAL_RE.finditer(original): + key_prefix = match.group(1) + literal_suffix = match.group(2) + restored = restored.replace( + f"{key_prefix}***", + f"{key_prefix}{literal_suffix}", + 1, + ) + return restored + def _fallback_redact(text: str) -> str: if not isinstance(text, str) or not text: return text text = _CRED_RE.sub(lambda m: _mask(m.group(1)), text) text = _AUTH_HDR_RE.sub(lambda m: m.group(1) + _mask(m.group(2)), text) - text = _ENV_RE.sub( - lambda m: f"{m.group(1)}={m.group(2)}{_mask(m.group(3))}{m.group(2)}", text - ) + text = _ENV_RE.sub(_env_replacement, text) text = _PRIVKEY_RE.sub("[REDACTED PRIVATE KEY]", text) return text @@ -304,6 +330,7 @@ def _build_redact_fn(): except TypeError: # Older hermes-agent builds that predate the force kwarg. agent_redacted = redact_sensitive_text(text) + agent_redacted = _restore_code_env_key_literals(text, agent_redacted) return _fallback_redact(agent_redacted) return _combined_redact @@ -357,6 +384,7 @@ _SENSITIVE_CASE_MARKERS = ( ) _SENSITIVE_LOWER_MARKERS = ( "authorization: bearer ", + "authorization: bot ", "private key", "postgres://", "postgresql://", diff --git a/api/streaming.py b/api/streaming.py index da4deb9e..af28636a 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -3788,12 +3788,28 @@ def _deduplicate_context_messages(messages): Prevents the agent from seeing the same message twice in conversation_history when result_messages contain duplicates that weren't caught by display-merge. + Compression/reference markers are internal recovery material: keep at most + one canonical assistant reference so a mis-role ``user`` marker cannot become + the next active user instruction. """ if not messages: return messages seen = set() deduped = [] for msg in messages: + if _is_context_compression_marker(msg): + marker_key = ( + '__context_compression_marker__', + " ".join(_message_text(msg.get('content', '')).split())[:500], + ) + if marker_key in seen: + continue + seen.add(marker_key) + if isinstance(msg, dict) and msg.get('role') != 'assistant': + msg = copy.deepcopy(msg) + msg['role'] = 'assistant' + deduped.append(msg) + continue key = _message_identity(msg) if key is not None and key in seen: continue @@ -4602,9 +4618,13 @@ def _merge_display_messages_after_agent_result(previous_display, previous_contex If Hermes Agent returns a normal append-only history, append that delta to the UI transcript. If the model/context history was compacted and no longer has the prior context as a prefix, keep the previous UI transcript and append - only compaction marker messages plus the current user turn onward. + the current user turn onward. Synthetic compaction/reference markers remain + internal recovery material and must not become visible user/assistant turns. """ - previous_display = list(previous_display or []) + previous_display = [ + m for m in list(previous_display or []) + if not _is_context_compression_marker(m) + ] # Deduplicate stale _partial messages that accumulated in previous_display. # A bug in cancel_stream() could insert multiple identical _partial messages # when _stripped was empty but _has_reasoning/_has_tools was True. The @@ -4750,10 +4770,6 @@ def _merge_display_messages_after_agent_result(previous_display, previous_contex candidates = _strip_replayed_prefix(previous_context, candidates) else: current_user_idx = _find_current_user_turn(result_messages, msg_text) - marker_candidates = [ - m for m in result_messages[:current_user_idx if current_user_idx is not None else len(result_messages)] - if _is_context_compression_marker(m) - ] turn_candidates = result_messages[current_user_idx:] if current_user_idx is not None else [] # Normalize stale merges only in the current-turn slice. if msg_text and previous_user_tail: @@ -4763,7 +4779,7 @@ def _merge_display_messages_after_agent_result(previous_display, previous_contex previous_user_tail, previous_context=previous_context, ) - candidates = marker_candidates + turn_candidates + candidates = turn_candidates merged = previous_display[:] seen = {_message_identity(m) for m in merged} @@ -4803,6 +4819,8 @@ def _merge_display_messages_after_agent_result(previous_display, previous_contex candidates = candidates[:insert_at] + [current_user_msg] + candidates[insert_at:] for msg in candidates: + if _is_context_compression_marker(msg): + continue key = _message_identity(msg) is_current_user_turn = _looks_like_current_user_turn(msg, msg_text) if ( diff --git a/tests/test_compaction_marker_writeback.py b/tests/test_compaction_marker_writeback.py new file mode 100644 index 00000000..811a2480 --- /dev/null +++ b/tests/test_compaction_marker_writeback.py @@ -0,0 +1,50 @@ +from api.streaming import ( + _deduplicate_context_messages, + _is_context_compression_marker, + _merge_display_messages_after_agent_result, +) + + +MARKER = "[CONTEXT COMPACTION — REFERENCE ONLY] Earlier turns were compacted." + + +def test_display_merge_drops_internal_compaction_markers_from_visible_delta(): + previous_display = [ + {"role": "user", "content": "old question"}, + {"role": "assistant", "content": "old answer"}, + ] + previous_context = list(previous_display) + result_messages = [ + {"role": "assistant", "content": MARKER, "_compressed_summary": True}, + {"role": "user", "content": MARKER, "_compressed_summary": True}, + {"role": "user", "content": "latest question"}, + {"role": "assistant", "content": "latest answer"}, + ] + + merged = _merge_display_messages_after_agent_result( + previous_display, + previous_context, + result_messages, + "latest question", + ) + + assert [m["content"] for m in merged] == [ + "old question", + "old answer", + "latest question", + "latest answer", + ] + assert not any(_is_context_compression_marker(m) for m in merged) + + +def test_context_dedup_keeps_one_reference_marker_without_user_role_duplicate(): + context = _deduplicate_context_messages([ + {"role": "assistant", "content": MARKER, "_compressed_summary": True}, + {"role": "user", "content": MARKER, "_compressed_summary": True}, + {"role": "user", "content": "latest question"}, + ]) + + markers = [m for m in context if _is_context_compression_marker(m)] + assert len(markers) == 1 + assert markers[0]["role"] == "assistant" + assert [m["content"] for m in context if m.get("role") == "user"] == ["latest question"] diff --git a/tests/test_issue1217_transcript_compaction.py b/tests/test_issue1217_transcript_compaction.py index 15e65705..93fb058f 100644 --- a/tests/test_issue1217_transcript_compaction.py +++ b/tests/test_issue1217_transcript_compaction.py @@ -79,11 +79,12 @@ def test_workspace_prefixed_current_user_after_compaction_is_not_duplicated(): "Ok, mache weiter", ) - assert [m["role"] for m in merged] == ["user", "assistant", "assistant", "user", "assistant"] + assert [m["role"] for m in merged] == ["user", "assistant", "user", "assistant"] assert [m["content"] for m in merged[-2:]] == [ "Ok, mache weiter", "continuing", ] + assert all("CONTEXT COMPACTION" not in str(m.get("content") or "") for m in merged) assert sum(1 for m in merged if m.get("role") == "user" and "Ok, mache weiter" in m.get("content", "")) == 1 @@ -206,10 +207,10 @@ def test_compacted_agent_result_keeps_old_prompts_and_appends_current_turn(): "first answer", "second prompt that must remain visible", "second answer", - "[CONTEXT COMPACTION — REFERENCE ONLY] Earlier turns were compacted.", "new question after compaction", "new answer after compaction", ] + assert all("CONTEXT COMPACTION" not in str(m.get("content") or "") for m in merged) def test_append_only_agent_result_preserves_normal_delta_behavior(): @@ -309,10 +310,10 @@ def test_repeated_user_text_after_compaction_is_not_dropped(): assert [m["content"] for m in merged] == [ "continue", "old answer", - "[CONTEXT COMPACTION — REFERENCE ONLY] summary", "continue", "new answer", ] + assert all("CONTEXT COMPACTION" not in str(m.get("content") or "") for m in merged) def test_near_duplicate_session_arc_summary_is_not_replayed_in_context(): diff --git a/tests/test_security_redaction.py b/tests/test_security_redaction.py index 298ce32a..a892744b 100644 --- a/tests/test_security_redaction.py +++ b/tests/test_security_redaction.py @@ -261,6 +261,37 @@ def test_redact_value_works_with_legacy_agent_redact_signature(monkeypatch): importlib.reload(helpers) +def test_redaction_preserves_secret_key_literals_in_code_without_values(): + import api.helpers as helpers + + key = "DISCORD" + "_BOT" + "_TOKEN" + snippet = ( + f'if line.startswith("{key}="):\n' + ' return line.split("=", 1)[1].strip()' + ) + + assert helpers._redact_value(snippet) == snippet + + +def test_redaction_masks_authorization_bot_value_without_breaking_code_structure(): + import ast + import api.helpers as helpers + + token = "M" + ("T" * 40) + ".abc123" + snippet = ( + 'headers = ["-H", ' + f'f"Authorization: Bot {token}", ' + '"-H", "User-Agent: HermesBot/1.0"]' + ) + + redacted = helpers._redact_value(snippet) + + assert token not in redacted + assert 'f"Authorization: Bot ' in redacted + assert '", "-H", "User-Agent: HermesBot/1.0"]' in redacted + ast.parse(redacted) + + def test_redact_session_data_messages(): """redact_session_data masks credentials in messages[].""" from api.helpers import redact_session_data