fix(proxy): tell users with an already exported master key to replace it in place, because it wins over .env

litellm/proxy/auth/master_key_boot_check.py

@@ -1,5 +1,3 @@
-"""Decides at boot whether the proxy may start with the master key it resolved."""
-
 import atexit
 import sys
 from collections.abc import Callable, Mapping
@@ -57,6 +55,7 @@ class UnsafeMasterKeyAllowed:
 class UnsafeMasterKeyRefused:
     reason: UnsafeMasterKeyReason
     source: MasterKeySource
+    environment_variable_is_set: bool
     stored_credentials_need_rotation: bool
 
 
@@ -90,6 +89,7 @@ def master_key_boot_verdict(
             if MASTER_KEY_SETTING in general_settings and not config_file_only_relays_the_environment
             else EnvironmentSource()
         ),
+        environment_variable_is_set=environment_master_key is not None,
         stored_credentials_need_rotation=(
             reason is UnsafeMasterKeyReason.PUBLICLY_KNOWN and not salt_key_is_set and database_is_configured
         ),
@@ -116,11 +116,7 @@ def enforce_master_key_boot_verdict(verdict: MasterKeyBootVerdict, announce: Cal
 
 
 def announce_on_stderr_at_exit(message: str) -> None:
-    """Keeps the fix as the last thing on screen and away from the log handlers.
-
-    A failed lifespan prints a traceback hundreds of lines long (a frame pair per included router) that buries
-    anything written before it, and the log redactor strips the key-shaped command from anything sent to a logger.
-    """
+    """A logger would redact the key-shaped command and the lifespan traceback would bury it, so print at exit."""
     atexit.register(_flush_stdout_then_write_stderr, message)
 
 
@@ -133,7 +129,7 @@ def render_refusal(refusal: UnsafeMasterKeyRefused) -> str:
     return "\n\n".join(
         (
             f"LiteLLM proxy refused to start: {_REFUSAL_HEADLINE[refusal.reason]}\n{_source_line(refusal)}",
-            _ROTATE_INSTEAD_OF_REPLACING if refusal.stored_credentials_need_rotation else _fix_steps(refusal.source),
+            _ROTATE_INSTEAD_OF_REPLACING if refusal.stored_credentials_need_rotation else _fix_steps(refusal),
             _OVERRIDE_HINT,
         )
     )
@@ -164,6 +160,14 @@ _SAVE_KEY_STEP: Final = (
     f"   {MASTER_KEY_ENV_VAR} environment variable instead."
 )
 
+_REPLACE_EXPORTED_KEY_STEP: Final = (
+    "Generate a key:\n"
+    f"     {PRINT_NEW_MASTER_KEY_COMMAND}\n"
+    f"   Put it in place of the current {MASTER_KEY_ENV_VAR} value wherever that is set: a shell export, your\n"
+    "   container or deployment environment, or its line in .env. Do not just add it to .env, because a value\n"
+    "   already exported in the environment wins over .env."
+)
+
 _ROTATE_INSTEAD_OF_REPLACING: Final = (
     f"Credentials stored in your database are encrypted with this master key because {SALT_KEY_ENV_VAR} is not\n"
     "set, so replacing the key makes them undecryptable. Rotate it by following this guide, which re-encrypts them:\n"
@@ -213,15 +217,16 @@ def _source_line(refusal: UnsafeMasterKeyRefused) -> str:
             assert_never(refusal.source)
 
 
-def _fix_steps(source: MasterKeySource) -> str:
-    match source:
-        case ConfigFileSource():
+def _fix_steps(refusal: UnsafeMasterKeyRefused) -> str:
+    set_key_step: Final = _REPLACE_EXPORTED_KEY_STEP if refusal.environment_variable_is_set else _SAVE_KEY_STEP
+    match refusal.source:
+        case ConfigFileSource() as source:
             return (
                 f"1. Make sure {_config_label(source)} reads the key from the environment:\n"
                 f"     general_settings:\n       {MASTER_KEY_SETTING}: os.environ/{MASTER_KEY_ENV_VAR}\n"
-                f"2. {_SAVE_KEY_STEP}"
+                f"2. {set_key_step}"
             )
         case EnvironmentSource():
-            return f"1. {_SAVE_KEY_STEP}"
+            return f"1. {set_key_step}"
         case _:
-            assert_never(source)
+            assert_never(refusal.source)

tests/test_litellm/proxy/auth/test_master_key_boot_check.py

@@ -1,6 +1,7 @@
 import re
 import shutil
 import subprocess
+import sys
 from collections.abc import Mapping
 from pathlib import Path
 
@@ -21,6 +22,7 @@ from litellm.proxy.auth.master_key_boot_check import (
     UnsafeMasterKeyError,
     UnsafeMasterKeyReason,
     UnsafeMasterKeyRefused,
+    announce_on_stderr_at_exit,
     enforce_master_key_boot_verdict,
     master_key_boot_verdict,
     render_refusal,
@@ -144,10 +146,14 @@ def test_rotation_is_only_needed_when_the_known_key_encrypts_a_database(
 def _refusal(
     reason: UnsafeMasterKeyReason = UnsafeMasterKeyReason.PUBLICLY_KNOWN,
     source: ConfigFileSource | EnvironmentSource = EnvironmentSource(),
+    environment_variable_is_set: bool = False,
     stored_credentials_need_rotation: bool = False,
 ) -> UnsafeMasterKeyRefused:
     return UnsafeMasterKeyRefused(
-        reason=reason, source=source, stored_credentials_need_rotation=stored_credentials_need_rotation
+        reason=reason,
+        source=source,
+        environment_variable_is_set=environment_variable_is_set,
+        stored_credentials_need_rotation=stored_credentials_need_rotation,
     )
 
 
@@ -160,13 +166,42 @@ def test_config_refusal_names_the_file_and_tells_it_to_read_the_environment():
 
 
 def test_environment_refusal_gives_the_command_without_a_config_step():
-    text = render_refusal(_refusal(source=EnvironmentSource()))
+    text = render_refusal(_refusal(reason=UnsafeMasterKeyReason.NOT_SET, source=EnvironmentSource()))
 
     assert f"the {MASTER_KEY_ENV_VAR} environment variable" in text
     assert GENERATE_MASTER_KEY_COMMAND in text
     assert "os.environ/" not in text
 
 
+@pytest.mark.parametrize(
+    ("master_key", "general_settings", "environment_master_key", "is_set"),
+    [
+        (None, {}, None, False),
+        ("sk-1234", {"master_key": "sk-1234"}, None, False),
+        ("sk-1234", {}, "sk-1234", True),
+        ("sk-1234", {"master_key": "sk-1234"}, "", True),
+    ],
+)
+def test_refusal_records_whether_the_environment_variable_is_already_set(
+    master_key: str | None, general_settings: Mapping[str, object], environment_master_key: str | None, is_set: bool
+):
+    refusal = _verdict(master_key, general_settings, environment_master_key=environment_master_key)
+
+    assert isinstance(refusal, UnsafeMasterKeyRefused)
+    assert refusal.environment_variable_is_set is is_set
+
+
+@pytest.mark.parametrize("source", [EnvironmentSource(), ConfigFileSource(config_file_path="/app/config.yaml")])
+def test_refusal_never_tells_a_user_with_an_exported_key_to_append_to_the_env_file(
+    source: ConfigFileSource | EnvironmentSource,
+):
+    text = render_refusal(_refusal(source=source, environment_variable_is_set=True))
+
+    assert PRINT_NEW_MASTER_KEY_COMMAND in text
+    assert "tee" not in text
+    assert "wins over .env" in text
+
+
 def test_unset_key_refusal_says_nothing_supplied_one():
     text = render_refusal(_refusal(reason=UnsafeMasterKeyReason.NOT_SET, source=EnvironmentSource()))
 
@@ -242,3 +277,19 @@ def test_safe_and_overridden_keys_boot_without_announcing(verdict: MasterKeyBoot
     enforce_master_key_boot_verdict(verdict, announce=announced.append)
 
     assert announced == []
+
+
+def test_announced_fix_is_the_last_thing_a_crashing_process_prints():
+    crash_after_announcing = (
+        "from litellm.proxy.auth.master_key_boot_check import announce_on_stderr_at_exit\n"
+        "announce_on_stderr_at_exit('THE FIX')\n"
+        "print('buffered stdout')\n"
+        "raise RuntimeError('lifespan failed')\n"
+    )
+
+    completed = subprocess.run([sys.executable, "-c", crash_after_announcing], capture_output=True, text=True)
+
+    assert completed.returncode != 0
+    assert "RuntimeError: lifespan failed" in completed.stderr
+    assert completed.stderr.endswith("THE FIX")
+    assert completed.stdout == "buffered stdout\n"