diff --git a/api/compression_anchor.py b/api/compression_anchor.py index 6da968ad..57118749 100644 --- a/api/compression_anchor.py +++ b/api/compression_anchor.py @@ -64,9 +64,10 @@ def is_context_compression_marker(message): message.get("content", ""), part_types={"text", "input_text", "output_text"}, ).lower().lstrip() + synthetic_unbracketed_marker = bool(message.get("_compressed_summary")) return ( text.startswith("[context compaction") - or text.startswith("context compaction") + or (synthetic_unbracketed_marker and text.startswith("context compaction")) or text.startswith("[your active task list was preserved across context compression]") or text.startswith("[session arc summary") ) diff --git a/api/helpers.py b/api/helpers.py index 47a63dc0..ed5b7755 100644 --- a/api/helpers.py +++ b/api/helpers.py @@ -269,7 +269,7 @@ def _build_redact_fn(): ) _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( @@ -288,20 +288,47 @@ def _build_redact_fn(): _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}=)([\"'][)\]:,]+|[)\]:,]+)" ) + _ENV_KEY_PREFIX_RE = _re.compile( + r"([A-Z0-9_]{0,50}(?:API_?KEY|TOKEN|SECRET|PASSWORD|PASSWD|CREDENTIAL|AUTH)[A-Z0-9_]{0,50}=)" + ) + _REDACTED_ENV_VALUE_RE = _re.compile( + r"(?:\*{3,}|[A-Za-z0-9][A-Za-z0-9_.:/+-]{0,32}\.\.\.[A-Za-z0-9_.:/+-]{1,16})" + ) 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): + literal_occurrences: dict[tuple[str, int], str] = {} + original_counts: dict[str, int] = {} + for match in _ENV_KEY_PREFIX_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 + occurrence = original_counts.get(key_prefix, 0) + original_counts[key_prefix] = occurrence + 1 + literal_match = _CODE_ENV_KEY_LITERAL_RE.match(original, match.start()) + if literal_match: + literal_occurrences[(key_prefix, occurrence)] = literal_match.group(2) + if not literal_occurrences: + return redacted + redacted_counts: dict[str, int] = {} + pieces = [] + last = 0 + for match in _ENV_KEY_PREFIX_RE.finditer(redacted): + key_prefix = match.group(1) + occurrence = redacted_counts.get(key_prefix, 0) + redacted_counts[key_prefix] = occurrence + 1 + literal_suffix = literal_occurrences.get((key_prefix, occurrence)) + if literal_suffix is None: + continue + value_match = _REDACTED_ENV_VALUE_RE.match(redacted, match.end()) + if not value_match: + continue + pieces.append(redacted[last:value_match.start()]) + pieces.append(literal_suffix) + last = value_match.end() + if not pieces: + return redacted + pieces.append(redacted[last:]) + return "".join(pieces) def _fallback_redact(text: str) -> str: if not isinstance(text, str) or not text: diff --git a/tests/test_compaction_marker_writeback.py b/tests/test_compaction_marker_writeback.py index 811a2480..641ee81a 100644 --- a/tests/test_compaction_marker_writeback.py +++ b/tests/test_compaction_marker_writeback.py @@ -48,3 +48,48 @@ def test_context_dedup_keeps_one_reference_marker_without_user_role_duplicate(): assert len(markers) == 1 assert markers[0]["role"] == "assistant" assert [m["content"] for m in context if m.get("role") == "user"] == ["latest question"] + + +def test_real_user_marker_prefix_text_is_preserved(): + prompt = "context compaction is broken; explain what happened" + previous_display = [ + {"role": "user", "content": "old question"}, + {"role": "assistant", "content": "old answer"}, + ] + previous_context = list(previous_display) + result_messages = [ + {"role": "user", "content": prompt}, + {"role": "assistant", "content": "It was treated as an internal marker."}, + ] + + merged = _merge_display_messages_after_agent_result( + previous_display, + previous_context, + result_messages, + prompt, + ) + context = _deduplicate_context_messages(result_messages) + + assert not _is_context_compression_marker(result_messages[0]) + assert [m["content"] for m in merged] == [ + "old question", + "old answer", + prompt, + "It was treated as an internal marker.", + ] + assert context[0]["role"] == "user" + assert context[0]["content"] == prompt + + +def test_unbracketed_marker_requires_internal_metadata(): + marker_text = "context compaction summary for internal recovery" + + assert not _is_context_compression_marker({ + "role": "user", + "content": marker_text, + }) + assert _is_context_compression_marker({ + "role": "assistant", + "content": marker_text, + "_compressed_summary": True, + }) diff --git a/tests/test_security_redaction.py b/tests/test_security_redaction.py index a892744b..20d86241 100644 --- a/tests/test_security_redaction.py +++ b/tests/test_security_redaction.py @@ -292,6 +292,67 @@ def test_redaction_masks_authorization_bot_value_without_breaking_code_structure ast.parse(redacted) +def _reload_helpers_with_fallback_redactor(monkeypatch): + fake_agent = types.ModuleType("agent") + fake_redact = types.ModuleType("agent.redact") + monkeypatch.setitem(sys.modules, "agent", fake_agent) + monkeypatch.setitem(sys.modules, "agent.redact", fake_redact) + + import api.helpers as helpers + return importlib.reload(helpers) + + +@pytest.mark.parametrize("text,leaked", [ + ("API_KEY=val,with,commas", "val,with,commas"), + ('TOKEN="quoted"', "quoted"), + ("PASSWORD=ends)", "ends)"), +]) +def test_fallback_env_redaction_masks_values_with_delimiters(monkeypatch, text, leaked): + helpers = _reload_helpers_with_fallback_redactor(monkeypatch) + try: + redacted = helpers._redact_value(text) + assert leaked not in redacted + assert "..." in redacted or "***" in redacted + finally: + importlib.reload(helpers) + + +def test_agent_redaction_restore_is_occurrence_aware_for_code_key_literals(monkeypatch): + key = "DISCORD" + "_BOT" + "_TOKEN" + secret = "real-secret-value-123456789" + mask = "real-s...6789" + original = ( + f"env line: {key}={secret}\n" + f'if line.startswith("{key}="):\n' + " return True" + ) + + fake_agent = types.ModuleType("agent") + fake_redact = types.ModuleType("agent.redact") + + def _fake_redact_sensitive_text(text, force=True): + assert force is True + return ( + text + .replace(f"{key}={secret}", f"{key}={mask}") + .replace(f'{key}="):', f"{key}={mask}") + ) + + fake_redact.redact_sensitive_text = _fake_redact_sensitive_text + monkeypatch.setitem(sys.modules, "agent", fake_agent) + monkeypatch.setitem(sys.modules, "agent.redact", fake_redact) + + import api.helpers as helpers + helpers = importlib.reload(helpers) + try: + redacted = helpers._redact_value(original) + assert secret not in redacted + assert f'if line.startswith("{key}="):' in redacted + assert f"env line: {key}=\"):" not in redacted + finally: + importlib.reload(helpers) + + def test_redact_session_data_messages(): """redact_session_data masks credentials in messages[].""" from api.helpers import redact_session_data