From 2f4c1e24df4d178e241e4c6e0fe7ee01c3596b51 Mon Sep 17 00:00:00 2001 From: Konstantin Sivakov Date: Tue, 7 Jul 2026 17:47:44 +0200 Subject: [PATCH] scrub user input messages from secrets --- anton/core/session.py | 22 ++++++++++++++++++++++ anton/tools.py | 8 ++++++-- tests/test_scrubbing.py | 34 ++++++++++++++++++++++++++++++++++ tests/test_tools.py | 28 ++++++++++++++++++++++++++++ 4 files changed, 90 insertions(+), 2 deletions(-) diff --git a/anton/core/session.py b/anton/core/session.py index f2b718ad..043c654c 100644 --- a/anton/core/session.py +++ b/anton/core/session.py @@ -99,6 +99,26 @@ def _extract_datasources(tool_call: ToolCall) -> List[str]: seen.add(m.group(1).lower()) return list(seen) + +def _scrub_user_input(user_input: str | list[dict]) -> str | list[dict]: + """Scrub credential values from an inbound user message. + + Applied at the `turn`/`turn_stream` entry, before the first + `_append_history`, so a secret pasted into chat never reaches model + context, episodic memory, or the trace sinks downstream of the LLM + gateway (Langfuse). Only text blocks are scrubbed; image and file + blocks carry no scrubbable text. + """ + if isinstance(user_input, str): + return scrub_credentials(user_input) + return [ + {**b, "text": scrub_credentials(b.get("text", ""))} + if b.get("type") == "text" + else b + for b in user_input + ] + + @dataclass class ChatSessionConfig: """All construction parameters for a ChatSession. @@ -1379,6 +1399,7 @@ def _schedule_cerebellum_flush(self) -> None: cb.reset() async def turn(self, user_input: str | list[dict]) -> str: + user_input = _scrub_user_input(user_input) self._append_history({"role": "user", "content": user_input}) user_msg_str = ( @@ -1541,6 +1562,7 @@ async def turn_stream( (e.g. an eval-run id) without any change to Anton. """ self._current_turn_id = turn_id + user_input = _scrub_user_input(user_input) self._append_history({"role": "user", "content": user_input}) # Log user input to episodic memory diff --git a/anton/tools.py b/anton/tools.py index 45fec2f8..1a6ae53e 100644 --- a/anton/tools.py +++ b/anton/tools.py @@ -18,7 +18,11 @@ "password", "secret", "token", "api_key", "key", "auth", "credential", "private", ) -SCRUBBED_VALUE_RE = re.compile(r"^\[DS_\w+\]$") +# All marker forms `scrub_credentials` emits: [DS_*] for vault secrets, +# [] labels for provider keys, [REDACTED_API_KEY] for shape matches. +# User messages are scrubbed too, so the model may see these and echo them +# back as known_variables — they must never be saved as real credentials. +SCRUBBED_VALUE_RE = re.compile(r"^\[(?:DS_\w+|[A-Z][A-Z0-9_]*)\]$") def looks_secret(field_name: str) -> bool: @@ -77,7 +81,7 @@ async def handle_connect_datasource(session: ChatSession, tc_input: dict) -> str console.print() console.print( f"[anton.warning](anton)[/] Ignoring scrubbed-placeholder values " - f"for {', '.join(dropped_scrubbed)} — those [DS_...] strings are " + f"for {', '.join(dropped_scrubbed)} — those bracketed strings are " f"scrub-markers, not real credentials. Pass the actual secret " f"values instead." ) diff --git a/tests/test_scrubbing.py b/tests/test_scrubbing.py index 2fbfd315..a7cc4772 100644 --- a/tests/test_scrubbing.py +++ b/tests/test_scrubbing.py @@ -5,6 +5,7 @@ import pytest +from anton.core.session import _scrub_user_input from anton.utils.datasources import ( _DS_KNOWN_VARS, _DS_SECRET_VARS, @@ -105,3 +106,36 @@ def test_short_sk_and_base_url_left_readable(self, monkeypatch): result = scrub_credentials("sk-abc connecting to https://api.openai.com/v1") assert "sk-abc" in result assert "https://api.openai.com/v1" in result + + +class TestScrubUserInput: + """User messages are scrubbed before entering session history (ENG-583).""" + + def test_string_input_key_redacted(self): + result = _scrub_user_input( + "use this key: sk-ant-api03-abcdefghij1234567890XYZ" + ) + assert "sk-ant-api03" not in result + assert "[REDACTED_API_KEY]" in result + + def test_plain_string_unchanged(self): + text = "please connect me to my staging database" + assert _scrub_user_input(text) == text + + def test_text_blocks_scrubbed_other_blocks_untouched(self): + blocks = [ + {"type": "text", "text": "key is mdb_AAAAAAAAAA.BBBBBBBBBBBBCCCC"}, + {"type": "image", "source": {"type": "base64", "data": "aGk="}}, + ] + result = _scrub_user_input(blocks) + assert "mdb_AAAAAAAAAA" not in result[0]["text"] + assert "[REDACTED_API_KEY]" in result[0]["text"] + assert result[1] is blocks[1] + + def test_known_secret_env_value_redacted_with_label(self, monkeypatch): + """A pasted value matching a stored provider secret gets its var label.""" + key = "sk-proj-abcDEF1234567890abcDEF1234567890" + monkeypatch.setenv("OPENAI_API_KEY", key) + result = _scrub_user_input(f"my key is {key}") + assert key not in result + assert "[OPENAI_API_KEY]" in result diff --git a/tests/test_tools.py b/tests/test_tools.py index 92e3c9bd..54842470 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -307,6 +307,34 @@ async def test_scrubbed_placeholder_values_are_dropped(self, vault_dir): assert "scrubbed-placeholder" in printed assert "api_key" in printed + @pytest.mark.asyncio + async def test_redacted_and_env_label_markers_are_dropped(self, vault_dir): + """User-message scrub markers ([REDACTED_*], []) must not be vaulted.""" + session = _make_session(vault_dir) + interactive = _fake_interactive(session) + with patch( + "anton.commands.datasource.handle_connect_datasource", new=interactive + ): + await handle_connect_datasource( + session, + { + "engine": "posthog", + "known_variables": { + "api_key": "[REDACTED_API_KEY]", + "host": "[OPENAI_API_KEY]", + }, + }, + ) + + interactive.assert_called_once() + assert session._data_vault.list_connections() == [] + printed = " ".join( + str(c.args[0]) + for c in session._console.print.call_args_list + if c.args + ) + assert "scrubbed-placeholder" in printed + @pytest.mark.asyncio async def test_mixed_real_and_scrubbed_values(self, vault_dir, tmp_path): """Real values are saved; scrubbed placeholders are dropped with a warning."""