Merge pull request #32405 from BerriAI/litellm_kraken-remove-envref-gates

fix(proxy): resolve os.environ/ refs universally in DB-sourced models
This commit is contained in:
yuneng-jiang 2026-07-07 21:53:56 -07:00 committed by GitHub
commit ec4f324482
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 42 additions and 87 deletions

View file

@ -1117,45 +1117,6 @@ _OPENAPI_HTTP_METHODS = {
# `_SSO_SENSITIVE_FIELDS` / `_CACHE_SENSITIVE_FIELDS` constants in the SSO
# and cache endpoint files.
_ALERTING_SENSITIVE_VARS: Set[str] = {"SLACK_WEBHOOK_URL", "SMTP_PASSWORD"}
_DB_LITELLM_PARAM_ENV_REF_KEYS = frozenset(
{
"api_key",
"client_secret",
"vertex_credentials",
"vertex_ai_credentials",
"aws_access_key_id",
"aws_secret_access_key",
"aws_session_token",
"aws_region_name",
"aws_session_name",
"aws_profile_name",
"aws_role_name",
"aws_web_identity_token",
"aws_sts_endpoint",
"aws_external_id",
"aws_bedrock_runtime_endpoint",
"aws_bedrock_project_id",
"aws_batch_role_arn",
"aws_workspace_id",
}
)
def _db_model_is_team_scoped(model: object) -> bool:
model_info = getattr(model, "model_info", None)
if isinstance(model_info, BaseModel):
return getattr(model_info, "team_id", None) is not None
if isinstance(model_info, str):
try:
model_info = json.loads(model_info)
except (TypeError, ValueError):
model_info = None
if isinstance(model_info, dict) and model_info.get("team_id") is not None:
return True
if getattr(model_info, "team_id", None) is not None:
return True
model_name = getattr(model, "model_name", None)
return isinstance(model_name, str) and model_name.startswith("model_name_")
def _strip_operation_id_method_suffix(operation_id: str) -> str:
@ -5009,17 +4970,12 @@ class ProxyConfig:
deleted_deployments += 1
return deleted_deployments
def _resolve_db_litellm_param(self, key: str, value: object, resolve_env_refs: bool = True) -> object:
def _resolve_db_litellm_param(self, key: str, value: object) -> object:
if not isinstance(value, str):
return value
decrypted_value = decrypt_value_helper(value=value, key=key, return_original_value=True)
if (
resolve_env_refs
and key in _DB_LITELLM_PARAM_ENV_REF_KEYS
and isinstance(decrypted_value, str)
and decrypted_value.startswith("os.environ/")
):
if isinstance(decrypted_value, str) and decrypted_value.startswith("os.environ/"):
return get_secret(decrypted_value)
return decrypted_value
@ -5040,13 +4996,10 @@ class ProxyConfig:
## ADD MODEL LOGIC
for m in db_models:
_litellm_params = m.litellm_params
resolve_env_refs = not _db_model_is_team_scoped(m)
if isinstance(_litellm_params, dict):
# decrypt values
for k, v in _litellm_params.items():
_litellm_params[k] = self._resolve_db_litellm_param(
key=k, value=v, resolve_env_refs=resolve_env_refs
)
_litellm_params[k] = self._resolve_db_litellm_param(key=k, value=v)
_litellm_params = LiteLLM_Params(**_litellm_params)
else:
@ -5072,15 +5025,12 @@ class ProxyConfig:
_model_list: list = []
for m in new_models:
_litellm_params = m.litellm_params
resolve_env_refs = not _db_model_is_team_scoped(m)
if isinstance(_litellm_params, BaseModel):
_litellm_params = _litellm_params.model_dump()
if isinstance(_litellm_params, dict):
# decrypt values
for k, v in _litellm_params.items():
_litellm_params[k] = self._resolve_db_litellm_param(
key=k, value=v, resolve_env_refs=resolve_env_refs
)
_litellm_params[k] = self._resolve_db_litellm_param(key=k, value=v)
_litellm_params = LiteLLM_Params(**_litellm_params)
else:
verbose_proxy_logger.error(

View file

@ -1102,6 +1102,11 @@ def test_ProxyConfig__add_deployment_invalid_litellm_params_skips(monkeypatch):
def test_ProxyConfig__add_deployment_resolves_env_refs_after_db_decrypt(monkeypatch):
"""Every ``os.environ/`` value on an admin-scoped DB row resolves at
load time, regardless of the field name. Replaces the earlier
behavior where only fields in ``_DB_LITELLM_PARAM_ENV_REF_KEYS``
resolved: the whitelist has been removed so the resolver applies to
every string field."""
monkeypatch.setenv("LITELLM_DB_MODEL_API_KEY", "resolved-secret")
monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret")
monkeypatch.setattr(
@ -1129,19 +1134,21 @@ def test_ProxyConfig__add_deployment_resolves_env_refs_after_db_decrypt(monkeypa
assert added == 1
assert deployment.litellm_params.api_key == "resolved-secret"
assert deployment.litellm_params.api_base == "os.environ/LITELLM_MASTER_KEY"
assert deployment.litellm_params.api_base == "master-secret"
def test_ProxyConfig__add_deployment_keeps_team_env_refs_literal(monkeypatch):
def fail_on_call(secret_name, *args, **kwargs):
raise AssertionError("team DB models should not resolve env refs")
def test_ProxyConfig__add_deployment_resolves_team_env_refs(monkeypatch):
"""Team-scoped DB rows now resolve ``os.environ/`` refs the same way
admin rows do. The prior team-scoped short-circuit and the
field-by-field whitelist have both been removed; the write-side team
auth check in ``ModelManagementAuthChecks.can_user_make_model_call``
remains the single trust boundary. A literal (non-``os.environ/``)
value still passes through unchanged."""
monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret")
monkeypatch.setattr(
"litellm.proxy.proxy_server.decrypt_value_helper",
lambda value, key, return_original_value: value,
)
monkeypatch.setattr("litellm.proxy.proxy_server.get_secret", fail_on_call)
fake_router = MagicMock()
fake_router.upsert_deployment = MagicMock(return_value=True)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
@ -1153,7 +1160,7 @@ def test_ProxyConfig__add_deployment_keeps_team_env_refs_literal(monkeypatch):
litellm_params={
"model": "openai/gpt-4o-mini",
"api_key": "os.environ/LITELLM_MASTER_KEY",
"api_base": "https://attacker.example",
"api_base": "https://team.example",
},
blocked=False,
)
@ -1162,8 +1169,8 @@ def test_ProxyConfig__add_deployment_keeps_team_env_refs_literal(monkeypatch):
deployment = fake_router.upsert_deployment.call_args.kwargs["deployment"]
assert added == 1
assert deployment.litellm_params.api_key == "os.environ/LITELLM_MASTER_KEY"
assert deployment.litellm_params.api_base == "https://attacker.example"
assert deployment.litellm_params.api_key == "master-secret"
assert deployment.litellm_params.api_base == "https://team.example"
def test_ProxyConfig__resolve_db_litellm_param_skips_non_string_values(monkeypatch):
@ -1242,31 +1249,26 @@ def test_ProxyConfig__add_deployment_resolves_env_refs_for_aws_bedrock_auth_para
assert getattr(deployment.litellm_params, key) == expected, key
def test_ProxyConfig__add_deployment_keeps_team_aws_env_refs_literal(monkeypatch):
"""Team-scoped DB models must NOT resolve env refs even for AWS auth
params: this is the LIT-3831 defense-in-depth path where a team admin
could otherwise craft a DB entry that reads the process environment."""
def fail_on_call(secret_name, *args, **kwargs):
raise AssertionError("team DB models should not resolve env refs")
monkeypatch.setenv("BEDROCK_ASSUME_ROLE_ARN", "arn:aws:iam::123:role/should-not-leak")
def test_ProxyConfig__add_deployment_resolves_env_refs_on_arbitrary_field(monkeypatch):
"""A made-up field name that was never on the removed whitelist still
resolves ``os.environ/`` refs. Pins the "no whitelist" invariant:
the resolver applies to every string field, not a curated list."""
monkeypatch.setenv("SOME_CUSTOM_ENV", "resolved-custom-value")
monkeypatch.setattr(
"litellm.proxy.proxy_server.decrypt_value_helper",
lambda value, key, return_original_value: value,
)
monkeypatch.setattr("litellm.proxy.proxy_server.get_secret", fail_on_call)
fake_router = MagicMock()
fake_router.upsert_deployment = MagicMock(return_value=True)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
pc = ProxyConfig()
db_model = SimpleNamespace(
model_id="model-1",
model_name="model_name_team-1_bedrock",
model_info={"id": "model-1", "team_id": "team-1"},
model_name="custom-field-model",
model_info={"id": "model-1"},
litellm_params={
"model": "bedrock/anthropic.claude-v2",
"aws_role_name": "os.environ/BEDROCK_ASSUME_ROLE_ARN",
"model": "openai/gpt-4o-mini",
"some_future_field": "os.environ/SOME_CUSTOM_ENV",
},
blocked=False,
)
@ -1275,7 +1277,7 @@ def test_ProxyConfig__add_deployment_keeps_team_aws_env_refs_literal(monkeypatch
deployment = fake_router.upsert_deployment.call_args.kwargs["deployment"]
assert added == 1
assert deployment.litellm_params.aws_role_name == "os.environ/BEDROCK_ASSUME_ROLE_ARN"
assert deployment.litellm_params.some_future_field == "resolved-custom-value"
# ---------------------------------------------------------------------------
@ -1313,6 +1315,9 @@ def test_ProxyConfig_decrypt_model_list_from_db_returns_decrypted(monkeypatch):
def test_ProxyConfig_decrypt_model_list_from_db_resolves_env_refs_after_db_decrypt(
monkeypatch,
):
"""Path B (feeding /v2/model/info fallback and /model/info fallback)
resolves every ``os.environ/`` field on admin-scoped rows, mirroring
path A. Both paths now share the same universal-resolution shape."""
monkeypatch.setenv("LITELLM_DB_MODEL_API_KEY", "resolved-secret")
monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret")
monkeypatch.setattr(
@ -1341,21 +1346,21 @@ def test_ProxyConfig_decrypt_model_list_from_db_resolves_env_refs_after_db_decry
out = pc.decrypt_model_list_from_db(new_models=[m])
assert out[0]["litellm_params"]["api_key"] == "resolved-secret"
assert out[0]["litellm_params"]["api_base"] == "os.environ/LITELLM_MASTER_KEY"
assert out[0]["litellm_params"]["api_base"] == "master-secret"
def test_ProxyConfig_decrypt_model_list_from_db_keeps_team_env_refs_literal_after_db_decrypt(
def test_ProxyConfig_decrypt_model_list_from_db_resolves_team_env_refs_after_db_decrypt(
monkeypatch,
):
def fail_on_call(secret_name, *args, **kwargs):
raise AssertionError("team DB models should not resolve env refs")
"""Team-scoped rows on path B resolve ``os.environ/`` refs just like
admin rows do. Pairs with
``test_ProxyConfig__add_deployment_resolves_team_env_refs`` on path
A both paths now agree on the trust model."""
monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret")
monkeypatch.setattr(
"litellm.proxy.proxy_server.decrypt_value_helper",
lambda value, key, return_original_value: "os.environ/LITELLM_MASTER_KEY" if key == "api_key" else value,
)
monkeypatch.setattr("litellm.proxy.proxy_server.get_secret", fail_on_call)
pc = ProxyConfig()
m = SimpleNamespace(
model_id="model-1",
@ -1363,7 +1368,7 @@ def test_ProxyConfig_decrypt_model_list_from_db_keeps_team_env_refs_literal_afte
model_info={"id": "model-1", "team_id": "team-1"},
litellm_params={
"api_key": "encrypted-env-ref",
"api_base": "https://attacker.example",
"api_base": "https://team.example",
"model": "openai/gpt-4o-mini",
},
blocked=False,
@ -1371,8 +1376,8 @@ def test_ProxyConfig_decrypt_model_list_from_db_keeps_team_env_refs_literal_afte
out = pc.decrypt_model_list_from_db(new_models=[m])
assert out[0]["litellm_params"]["api_key"] == "os.environ/LITELLM_MASTER_KEY"
assert out[0]["litellm_params"]["api_base"] == "https://attacker.example"
assert out[0]["litellm_params"]["api_key"] == "master-secret"
assert out[0]["litellm_params"]["api_base"] == "https://team.example"
def test_ProxyConfig_decrypt_model_list_from_db_invalid_params_skips():