diff --git a/tests/code_coverage_tests/test_e2e_metadata.py b/tests/code_coverage_tests/test_e2e_metadata.py index 042a89610ae..bb0eff38304 100644 --- a/tests/code_coverage_tests/test_e2e_metadata.py +++ b/tests/code_coverage_tests/test_e2e_metadata.py @@ -23,7 +23,7 @@ from types import UnionType from typing import Final, cast, get_args, get_type_hints import pytest -from e2e_metadata import MAX_STEPS, STEP_FRAMES, STEPS, step +from e2e_metadata import MASK, MAX_STEPS, STEP_FRAMES, STEPS, StepRecorder, environment_secrets, step from proxy_client import ProxyClient from pydantic import BaseModel, Field @@ -284,6 +284,55 @@ class TestLabelTemplates: assert STEPS.taken() == ("GET /v1/batches/{id}",) +class TestSecretMasking: + """Steps are published with the results, so a credential the run holds is + masked wherever it shows up in a label: a nested model field nobody marked + `repr=False`, a dict value, or a prompt.""" + + def test_a_secret_anywhere_in_a_label_is_masked(self) -> None: + recorder: Final = StepRecorder(secrets=lambda: ("sk-live-abcdef123", "wandb-9f8e7d6c")) + recorder.record("Generate a virtual key with callback vars: wandb api key: wandb-9f8e7d6c") + recorder.record('Send "use sk-live-abcdef123 please" to claude-haiku-4-5') + assert recorder.taken() == ( + f"Generate a virtual key with callback vars: wandb api key: {MASK}", + f'Send "use {MASK} please" to claude-haiku-4-5', + ) + + def test_a_secret_is_masked_before_the_label_is_cut(self) -> None: + secret: Final = "s3cr3t-" + "x" * 40 + recorder: Final = StepRecorder(secrets=lambda: (secret,)) + recorder.record("a" * 170 + " " + secret) + assert recorder.taken() == ("a" * 170 + f" {MASK}",) + + def test_a_longer_secret_containing_a_shorter_one_is_masked_whole(self) -> None: + recorder: Final = StepRecorder(secrets=lambda: ("abcdefgh", "abcdefgh-ijklmnop")) + recorder.record("key abcdefgh-ijklmnop") + assert recorder.taken() == (f"key {MASK}",) + + def test_only_secret_named_variables_long_enough_to_be_credentials_count(self) -> None: + environ: Final = { + "OPENAI_API_KEY": "sk-proj-0123456789", + "AWS_SECRET_ACCESS_KEY": "wJalrXUtnFEMI/K7MDENG", + "LITELLM_MASTER_KEY": "sk-1234", + "GOOGLE_APPLICATION_CREDENTIALS": "/secrets/vertex.json", + "KEYCLOAK_URL": "http://localhost:8080", + "E2E_MODEL": "claude-haiku-4-5", + } + assert environment_secrets(environ) == frozenset( + {"sk-proj-0123456789", "wJalrXUtnFEMI/K7MDENG", "/secrets/vertex.json"} + ) + + def test_the_shared_log_masks_the_live_environment(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("WANDB_API_KEY", "wandb-live-5a4b3c2d") + + @step("Generate a virtual key with {body}") + def generate_key(body: _KeyBody) -> None: + return None + + generate_key(_KeyBody(team_id="wandb-live-5a4b3c2d")) + assert STEPS.taken() == (f"Generate a virtual key with team id: {MASK}",) + + class TestNestedSteps: """Harness layers call each other, so a step's helper routinely calls other decorated helpers. Only the outermost records.""" diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index c4a44447fe1..cbecc1adee7 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -156,7 +156,7 @@ Generate a virtual key with models: claude-haiku-4-5 and rpm limit: 3 Send a /chat/completions request to claude-haiku-4-5 with the prompt "reply with one word d3940a1c4288" ``` -A request model prints only the fields the test set, and a dotted placeholder like `{body.litellm_params.model}` prints just one field. A field marked `Field(repr=False)` never prints, so mark every secret field that way, and never put a key, token or credential in a label. A placeholder that isn't one of the helper's parameters fails at import, and a literal brace is written `{{id}}`. A filled-in label is squashed onto one line and cut at 200 characters +A request model prints only the fields the test set, and a dotted placeholder like `{body.litellm_params.model}` prints just one field. A field marked `Field(repr=False)` never prints, so mark every secret field that way, and never put a key, token or credential in a label. As a backstop, the recorder replaces the value of every secret-named environment variable (`*_KEY`, `*_SECRET`, `*_TOKEN`, `*_PASSWORD`, `*_CREDENTIALS`) with `***` wherever it shows up in a label. That only covers secrets the environment holds, so a key the proxy hands back during the test is still never named in a label. A placeholder that isn't one of the helper's parameters fails at import, and a literal brace is written `{{id}}`. A filled-in label is squashed onto one line and cut at 200 characters ### Nesting and the step log diff --git a/tests/e2e/e2e_metadata.py b/tests/e2e/e2e_metadata.py index bbd7eef0d16..e5cd016e9d2 100644 --- a/tests/e2e/e2e_metadata.py +++ b/tests/e2e/e2e_metadata.py @@ -13,6 +13,7 @@ this one, so it imports only the stdlib and pydantic. from __future__ import annotations import inspect +import os import re import string import threading @@ -33,13 +34,31 @@ _Y = TypeVar("_Y") MAX_STEPS: Final = 50 MAX_STEP_CHARS: Final = 200 +SECRET_ENV_NAME: Final = re.compile(r"(^|_)(KEY|SECRET|TOKEN|PASSWORD|CREDENTIALS?)(_|$)", re.IGNORECASE) +MIN_SECRET_CHARS: Final = 8 +MASK: Final = "***" + + +def environment_secrets(environ: Mapping[str, str] = os.environ) -> frozenset[str]: + """The credentials a live run holds: every secret-named environment variable's + value, long enough that masking it can't blank out ordinary words.""" + return frozenset( + value for name, value in environ.items() if SECRET_ENV_NAME.search(name) and len(value) >= MIN_SECRET_CHARS + ) + + +def _masked(label: str, secrets: Iterable[str]) -> str: + longest_first: Final = sorted(secrets, key=len, reverse=True) + return reduce(lambda text, secret: text.replace(secret, MASK), longest_first, label) + + STEP_FRAMES: Final = 1 """Frames a `@step` wrapper puts between a helper and its caller. A decorated helper that warns about its caller adds this to `stacklevel` (`stacklevel=2 + STEP_FRAMES`), or the warning is reported at the wrapper.""" -class _StepRecorder: +class StepRecorder: """The ordered step log for the running test. A plain lock-guarded list rather than a ContextVar: ContextVars do not @@ -48,7 +67,8 @@ class _StepRecorder: cross-test bleed beyond what the per-test reset already handles. """ - def __init__(self) -> None: + def __init__(self, secrets: Callable[[], Iterable[str]] = environment_secrets) -> None: + self._secrets = secrets self._lock = threading.Lock() self._steps: deque[str] = deque(maxlen=MAX_STEPS) self._dropped = 0 @@ -69,8 +89,11 @@ class _StepRecorder: the story rather than fifty, and past MAX_STEPS the oldest step makes way. It is the oldest that goes because the last step is the one that has to survive: it is where a failing test died. + + Any credential the run holds is masked before the label is kept, however it + got into the label, since the steps are published with the results. """ - cleaned = " ".join(label.split())[:MAX_STEP_CHARS] + cleaned = " ".join(_masked(label, self._secrets()).split())[:MAX_STEP_CHARS] if not cleaned: return with self._lock: @@ -88,7 +111,7 @@ class _StepRecorder: return dropped + tuple(self._steps) -STEPS: Final = _StepRecorder() +STEPS: Final = StepRecorder() def _joined(phrases: tuple[str, ...]) -> str: