diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 144de52d0d2..c7b6bca72cf 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -6,8 +6,12 @@ from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Final, NamedTuple, Protocol, TypedDict +from pydantic import TypeAdapter, ValidationError + import litellm +from litellm.constants import REDACTED_BY_LITELM_STRING from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, ) @@ -52,6 +56,9 @@ class AgentRecord(Protocol): @property def agent_name(self) -> str: ... + @property + def litellm_params(self) -> Mapping[str, object] | None: ... + @property def object_permission_id(self) -> str | None: ... @@ -121,6 +128,188 @@ def _dump_agent_params(raw: Mapping[str, object]) -> dict[str, object]: return dict(raw) if raw else {} +_AGENT_PARAMS_MASKER: Final = SensitiveDataMasker() +_REDACT_AGENT_PARAMS_MAX_DEPTH: Final = 10 +_AGENT_PARAMS_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter( + dict[str, object] +) # mutable-ok: safe_dumps() and AgentResponse.litellm_params both require a real dict, not a Mapping +_AGENT_PARAMS_SEQUENCE_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...]) +_EMPTY_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def redact_sensitive_agent_litellm_params(litellm_params: object, _depth: int = 0) -> object: + """ + Replace credential-bearing values in an agent's litellm_params with + ``REDACTED_BY_LITELM_STRING`` while preserving non-secret keys (``model``, + ``is_public``, rate-limit config). Used so list/get/create/update + responses never echo a stored provider credential back to the caller. + + Handles a plain dict, a JSON-serialized string (some callers hold the + in-memory registry's params that way), and ``None`` at the top level; + anything else is passed through. Recursion depth is bounded to match the + convention documented in ``tests/code_coverage_tests/recursive_detector.py``. + """ + if litellm_params is None: + return None + if isinstance(litellm_params, str): + if _depth >= _REDACT_AGENT_PARAMS_MAX_DEPTH: + return REDACTED_BY_LITELM_STRING + try: + parsed_params: Final = _AGENT_PARAMS_ADAPTER.validate_json(litellm_params) + except ValidationError: + return REDACTED_BY_LITELM_STRING + return json.dumps(_redact_agent_params_tree(parsed_params, _depth + 1)) + return _redact_agent_params_tree(litellm_params, _depth) + + +def _redact_agent_params_tree(value: object, _depth: int) -> object: + """Structural recursion over an already-parsed litellm_params value: a + dict redacts sensitive keys and recurses into the rest, a list redacts + each element (so a secret nested inside a list of provider configs is + still caught), and anything else -- including a plain string leaf, which + must never be re-interpreted as a JSON blob -- passes through unchanged. + """ + if _depth >= _REDACT_AGENT_PARAMS_MAX_DEPTH: + return REDACTED_BY_LITELM_STRING + if isinstance(value, list): + typed_items: Final = _AGENT_PARAMS_SEQUENCE_ADAPTER.validate_python(value) + return tuple(_redact_agent_params_tree(item, _depth + 1) for item in typed_items) + if not isinstance(value, dict): + return value + typed_params: Final = _AGENT_PARAMS_ADAPTER.validate_python(value) + return { + key: ( + REDACTED_BY_LITELM_STRING + if _AGENT_PARAMS_MASKER.is_sensitive_key(key) + else _redact_agent_params_tree(nested_value, _depth + 1) + ) + for key, nested_value in typed_params.items() + } # mutable-ok: consumed by json.dumps()/AgentResponse.litellm_params, both of which require a real dict + + +def parse_agent_litellm_params(value: object) -> Mapping[str, object]: + """Normalize a stored litellm_params column to a read-only mapping. + + The prisma Json column comes back as either an already-parsed dict or a + JSON string depending on the read path, so handle both rather than + assuming one. Only ever read from (merge-source lookups), never mutated + or re-serialized directly, so a read-only view is enough here. + """ + if isinstance(value, str): + try: + return _AGENT_PARAMS_ADAPTER.validate_json(value) + except ValidationError: + return _EMPTY_LITELLM_PARAMS + if isinstance(value, Mapping): + try: + return _AGENT_PARAMS_ADAPTER.validate_python(value) + except ValidationError: + return _EMPTY_LITELLM_PARAMS + return _EMPTY_LITELLM_PARAMS + + +_MISSING_AGENT_PARAM: Final = object() +_RESTORE_AGENT_PARAMS_MAX_DEPTH: Final = 10 + + +def _restore_redacted_nested_value(incoming_value: object, existing_value: object, _depth: int) -> object: + """Recurse into a non-sensitively-named dict/list value so a secret + nested underneath it (e.g. inside a list of per-provider configs) is + still restored, not just top-level keys. Mirrors the shapes + ``redact_sensitive_agent_litellm_params`` recurses into on read, so + restore and redact stay symmetric. + + List elements are paired with the existing list by position: with no + stable per-element identity in an arbitrary ``dict[str, object]`` schema, + index is the same correspondence every other part of this restore (and + the endpoints' existing full-replace-on-PUT semantics) already assumes. + This correctly preserves a masked secret across an ordinary edit of that + same entry's other fields; it does not protect against a caller who both + reorders/resizes the list AND echoes back a masked marker in the same + request, which is a known, narrow limitation (see LIT-6736 PR discussion) + rather than a cross-entry credential leak in the common case. + + A value collapsed to the flat marker by the read side's depth cap is + recovered wholesale from ``existing_value`` (rather than the marker + string itself getting persisted) whenever ``existing_value`` isn't + already that same flat marker. Depth-bounded like its read-side + counterpart; a value at the cap is returned unchanged rather than + corrupted. + """ + if incoming_value == REDACTED_BY_LITELM_STRING and existing_value != REDACTED_BY_LITELM_STRING: + return existing_value + if _depth >= _RESTORE_AGENT_PARAMS_MAX_DEPTH: + return incoming_value + if isinstance(incoming_value, Mapping): + typed_incoming_map: Final = _AGENT_PARAMS_ADAPTER.validate_python(incoming_value) + existing_map: Final = ( + _AGENT_PARAMS_ADAPTER.validate_python(existing_value) + if isinstance(existing_value, Mapping) + else _EMPTY_LITELLM_PARAMS + ) + return _restore_redacted_litellm_params(typed_incoming_map, existing_map, _depth + 1) + if isinstance(incoming_value, (list, tuple)): + typed_incoming_seq: Final = _AGENT_PARAMS_SEQUENCE_ADAPTER.validate_python(incoming_value) + existing_seq: Final = ( + _AGENT_PARAMS_SEQUENCE_ADAPTER.validate_python(existing_value) + if isinstance(existing_value, (list, tuple)) + else () + ) + return tuple( + _restore_redacted_nested_value( + item, + existing_seq[index] if index < len(existing_seq) else None, + _depth + 1, + ) + for index, item in enumerate(typed_incoming_seq) + ) + return incoming_value + + +def _resolved_agent_param_value( + key: str, + incoming: Mapping[str, object], + existing: Mapping[str, object], + _depth: int, +) -> object: + """The value ``key`` should end up with in a restored litellm_params, or + ``_MISSING_AGENT_PARAM`` when it should be dropped entirely.""" + if key in incoming: + value: Final = incoming[key] + if _AGENT_PARAMS_MASKER.is_sensitive_key(key): + return existing.get(key, _MISSING_AGENT_PARAM) if value == REDACTED_BY_LITELM_STRING else value + return _restore_redacted_nested_value(value, existing.get(key), _depth) + if _AGENT_PARAMS_MASKER.is_sensitive_key(key): + return existing.get(key, _MISSING_AGENT_PARAM) + return _MISSING_AGENT_PARAM + + +def _restore_redacted_litellm_params( + incoming: Mapping[str, object], + existing: Mapping[str, object], + _depth: int = 0, +) -> dict[str, object]: + """Restore the real credential behind any litellm_params value the caller + echoed back as ``REDACTED_BY_LITELM_STRING``, and behind any sensitive key + omitted entirely, so an edit to an unrelated field never overwrites (or + silently drops) a stored provider credential -- the UI never has to + read-and-resend a secret to keep it. Recurses into nested dicts and lists + so a secret nested under a non-sensitively-named key is restored too. + + A sensitive key given a real (non-marker) value, including an explicit + empty string, is treated as a deliberate update -- that's how a caller + clears a credential. Non-sensitive keys always take the incoming value + (recursed into), matching the endpoints' existing full-replace-on-PUT / + merge-on-PATCH semantics for everything that isn't a secret. + """ + all_keys: Final = frozenset(incoming) | frozenset(existing) + return { + key: value + for key in all_keys + if (value := _resolved_agent_param_value(key, incoming, existing, _depth)) is not _MISSING_AGENT_PARAM + } # mutable-ok: fed to safe_dumps() for JSON-column storage, which requires a real dict + + class GrantMigrationResult(NamedTuple): rewritten: int missed: int @@ -301,9 +490,14 @@ class AgentRegistry: try: agent_name: Final = agent.get("agent_name") - # Serialize litellm_params + # Serialize litellm_params. A create has no stored row to restore a + # secret behind, so a sensitive key submitted as the redaction + # marker (e.g. a stray client re-post) is dropped rather than + # persisted as the literal placeholder string. litellm_params_obj: Final = agent.get("litellm_params", {}) - litellm_params_dict: Final[dict[str, object]] = _dump_agent_params(litellm_params_obj) + litellm_params_dict: Final = _restore_redacted_litellm_params( + _dump_agent_params(litellm_params_obj), _EMPTY_LITELLM_PARAMS + ) litellm_params: Final[str] = safe_dumps(litellm_params_dict) # Serialize agent_card_params @@ -410,8 +604,14 @@ class AgentRegistry: update_data: Final[dict[str, object]] = {} if augment_agent.get("agent_name"): update_data["agent_name"] = augment_agent.get("agent_name") - if augment_agent.get("litellm_params"): - update_data["litellm_params"] = safe_dumps(augment_agent.get("litellm_params")) + if "litellm_params" in agent: + existing_litellm_params: Final = parse_agent_litellm_params(existing_agent.get("litellm_params")) + update_data["litellm_params"] = safe_dumps( + _restore_redacted_litellm_params( + _dump_agent_params(agent.get("litellm_params") or _EMPTY_LITELLM_PARAMS), + existing_litellm_params, + ) + ) if augment_agent.get("agent_card_params"): update_data["agent_card_params"] = safe_dumps(augment_agent.get("agent_card_params")) @@ -474,9 +674,22 @@ class AgentRegistry: try: agent_name: Final = agent.get("agent_name") + # A PUT fully replaces litellm_params from the request body, so the + # existing row is read up front to restore any sensitive key the + # caller echoed back redacted (or omitted) rather than persisting + # the marker -- or nothing -- over the real stored credential. + existing_row: Final = await agents_table(prisma_client).find_unique( + where={"agent_id": agent_id} # mutable-ok: prisma's query builder rejects a Mapping/MappingProxyType + ) + existing_litellm_params: Final = parse_agent_litellm_params( + existing_row.litellm_params if existing_row is not None else None + ) + # Serialize litellm_params litellm_params_obj: Final = agent.get("litellm_params", {}) - litellm_params_dict: Final[dict[str, object]] = _dump_agent_params(litellm_params_obj) + litellm_params_dict: Final = _restore_redacted_litellm_params( + _dump_agent_params(litellm_params_obj), existing_litellm_params + ) litellm_params: Final[str] = safe_dumps(litellm_params_dict) # Serialize agent_card_params @@ -512,9 +725,8 @@ class AgentRegistry: update_data[rate_field] = _val if agent.get("object_permission") is not None: - existing_agent: Final = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id}) existing_object_permission_id: Final = ( - existing_agent.object_permission_id if existing_agent is not None else None + existing_row.object_permission_id if existing_row is not None else None ) agent_copy: Final = dict(agent) object_permission_id: Final = await handle_update_object_permission_common( diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index b6c41a17503..3e4dc07a521 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -20,7 +20,6 @@ from typing_extensions import ReadOnly, Required import litellm from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.litellm_logging import _get_masked_values from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._types import ( CommonProxyErrors, @@ -33,6 +32,10 @@ from litellm.proxy.a2a.agent_card import ( merge_agent_card, normalize_protocol_version, ) +from litellm.proxy.agent_endpoints.agent_registry import ( + parse_agent_litellm_params, + redact_sensitive_agent_litellm_params, +) from litellm.proxy.agent_endpoints.agent_search import ( DEFAULT_AGENT_SEARCH_TOP_K, AgentSearchEmbeddingFailed, @@ -139,25 +142,37 @@ async def _attach_keys_to_agents(agents: Sequence[AgentResponse], prisma_client) agent.keys = matched_keys or None +def _redact_agent_litellm_params_dict( + litellm_params: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: AgentResponse.litellm_params is declared as a plain dict, not Mapping + """Type-narrowing wrapper: a dict in always yields a dict back from + ``redact_sensitive_agent_litellm_params``, which the function's general + (possible-JSON-string, possibly-None) signature can't express.""" + return dict( # mutable-ok: AgentResponse.litellm_params is declared as a plain dict, not Mapping + parse_agent_litellm_params(redact_sensitive_agent_litellm_params(litellm_params)) + ) + + def _redact_sensitive_agent_fields( agents: Sequence[AgentResponse], + *, + is_admin: bool, ) -> list[AgentResponse]: """ - Return copies of the given agents with sensitive configuration fields - redacted. The original objects are not modified. + Return copies of the given agents with credential-bearing litellm_params + values replaced by a fixed marker (never returned to ANY caller, + admin included) and, for non-admin callers, virtual-key and header + fields stripped entirely. The original objects are not modified. """ redacted: Final[list[AgentResponse]] = [] for agent in agents: copy = agent.model_copy(deep=True) - copy.static_headers = None - copy.extra_headers = None - copy.keys = None + if not is_admin: + copy.static_headers = None + copy.extra_headers = None + copy.keys = None if copy.litellm_params: - copy.litellm_params = _get_masked_values( - copy.litellm_params, - unmasked_length=4, - number_of_asterisks=4, - ) + copy.litellm_params = _redact_agent_litellm_params_dict(copy.litellm_params) redacted.append(copy) return redacted @@ -345,13 +360,13 @@ async def get_agents( global_agent_registry.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups) ) - # Redact sensitive fields for non-admin users + # litellm_params secrets are always redacted; keys/headers stay + # admin-only. is_admin: Final = ( user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value ) - if not is_admin: - returned_agents = _redact_sensitive_agent_fields(returned_agents) + returned_agents = _redact_sensitive_agent_fields(returned_agents, is_admin=is_admin) if health_check: agents_with_url: Final = [agent for agent in returned_agents if (agent.agent_card_params or {}).get("url")] @@ -505,7 +520,9 @@ async def create_agent( "Failed to register agent '%s' (ID: %s) in memory: %s", agent_name, agent_id, reg_error ) - return result + # The caller is a proxy admin (enforced above); litellm_params + # secrets are still never echoed back in the response. + return _redact_sensitive_agent_fields((result,), is_admin=True)[0] except HTTPException: raise @@ -578,13 +595,13 @@ async def get_agent_by_id( await _attach_keys_to_agents([agent], prisma_client) - # Redact sensitive fields for non-admin users + # litellm_params secrets are always redacted; keys/headers stay + # admin-only. is_admin = ( user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value ) - if not is_admin: - agent = _redact_sensitive_agent_fields([agent])[0] + agent = _redact_sensitive_agent_fields((agent,), is_admin=is_admin)[0] return agent except HTTPException: @@ -688,7 +705,7 @@ async def update_agent( "Successfully updated agent '%s' (ID: %s) in memory", existing_agent.get("agent_name"), agent_id ) - return result + return _redact_sensitive_agent_fields((result,), is_admin=True)[0] except HTTPException: raise except Exception as e: @@ -791,7 +808,7 @@ async def patch_agent( "Successfully updated agent '%s' (ID: %s) in memory", existing_agent.get("agent_name"), agent_id ) - return result + return _redact_sensitive_agent_fields((result,), is_admin=True)[0] except HTTPException: raise except Exception as e: diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 790956156b0..0578dc60119 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -64,6 +64,8 @@ IGNORE_FUNCTIONS = [ "_flatten_form_field", # bounded by the nesting depth of the already-parsed request body (a finite JSON tree, no cycles possible). "_flatten_form_data_field", # bounded by the nesting depth of the already-parsed request body (a finite JSON tree, no cycles possible). "_json_safe", # max depth set (_MAX_DEPTH) plus a seen-ids cycle guard for self-referential input. + "_redact_agent_params_tree", # max depth set (default 10), same shape as _redact_sensitive_litellm_params. + "_restore_redacted_nested_value", # max depth set (default 10), mirrors _redact_agent_params_tree on the write side. ] diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py index 7ce62fdf648..231626c7eb5 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py @@ -8,7 +8,18 @@ from unittest.mock import AsyncMock, MagicMock import pytest -from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry, GrantMigrationResult +from litellm.constants import REDACTED_BY_LITELM_STRING +from litellm.proxy.agent_endpoints.agent_registry import ( + AgentRegistry, + GrantMigrationResult, + _restore_redacted_litellm_params, + redact_sensitive_agent_litellm_params, +) + +# Obviously-fake stand-ins for a real AWS credential pair (LIT-6736 regression +# fixtures) -- never a real key shape, and must never appear in any response. +SENTINEL_AWS_ACCESS_KEY_ID: Final = "AKIATESTSENTINEL0000" +SENTINEL_AWS_SECRET_ACCESS_KEY: Final = "test-sentinel-do-not-use-secret-value" def _sample_agent_card_params() -> dict: @@ -49,6 +60,7 @@ async def test_update_agent_in_db_clears_static_headers_and_extra_headers_when_o mock_update = AsyncMock(return_value=updated_agent) mock_prisma.db.litellm_agentstable.update = mock_update + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=None) # Agent config WITHOUT static_headers or extra_headers (omitted) agent_config = { @@ -95,6 +107,7 @@ async def test_update_agent_in_db_preserves_explicit_static_headers_and_extra_he mock_update = AsyncMock(return_value=updated_agent) mock_prisma.db.litellm_agentstable.update = mock_update + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=None) agent_config = { "agent_name": "Updated Agent", @@ -436,6 +449,9 @@ async def test_update_agent_in_db_raises_when_row_deleted_mid_update(): guard the code dereferences None and reports an opaque AttributeError instead of the id.""" registry: Final = AgentRegistry() mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace(litellm_params={}, object_permission_id=None) + ) mock_prisma.db.litellm_agentstable.update = AsyncMock(return_value=None) with pytest.raises(Exception, match="Error updating agent in DB") as exc_info: @@ -485,3 +501,492 @@ async def test_delete_agent_from_db_raises_when_row_already_gone(): await registry.delete_agent_from_db(agent_id="agent-123", prisma_client=mock_prisma) assert str(exc_info.value) == "Error deleting agent from DB: Agent not found, passed agent_id=agent-123" + + +# ---------- LIT-6736: agent litellm_params secret redaction ---------- + + +def test_redact_sensitive_agent_litellm_params_masks_secrets_keeps_the_rest(): + """The sentinel secret must never appear in the redacted output; non-secret + keys (model reference, is_public) must survive untouched.""" + redacted = redact_sensitive_agent_litellm_params( + { + "aws_access_key_id": SENTINEL_AWS_ACCESS_KEY_ID, + "aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, + "model": "bedrock/agentcore/my-agent", + "is_public": True, + } + ) + + assert SENTINEL_AWS_ACCESS_KEY_ID not in json.dumps(redacted) + assert SENTINEL_AWS_SECRET_ACCESS_KEY not in json.dumps(redacted) + assert redacted["aws_access_key_id"] == REDACTED_BY_LITELM_STRING + assert redacted["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING + assert redacted["model"] == "bedrock/agentcore/my-agent" + assert redacted["is_public"] is True + + +def test_redact_sensitive_agent_litellm_params_recurses_into_nested_dicts(): + """A secret nested one level down (e.g. a per-provider sub-config) must + also be redacted, not just top-level keys.""" + redacted = redact_sensitive_agent_litellm_params( + {"provider_config": {"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "region": "us-east-1"}} + ) + + assert SENTINEL_AWS_SECRET_ACCESS_KEY not in json.dumps(redacted) + assert redacted["provider_config"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING + assert redacted["provider_config"]["region"] == "us-east-1" + + +def test_redact_sensitive_agent_litellm_params_handles_none_and_json_string(): + assert redact_sensitive_agent_litellm_params(None) is None + + serialized = json.dumps({"api_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "model": "gpt-4"}) + redacted = redact_sensitive_agent_litellm_params(serialized) + + assert SENTINEL_AWS_SECRET_ACCESS_KEY not in redacted + assert json.loads(redacted)["api_key"] == REDACTED_BY_LITELM_STRING + assert json.loads(redacted)["model"] == "gpt-4" + + +def test_redact_sensitive_agent_litellm_params_recurses_into_lists_of_dicts(): + """A secret nested inside a list of provider sub-configs (a shape a + non-sensitively-named key can legitimately hold) must also be redacted, + not silently returned as-is.""" + redacted = redact_sensitive_agent_litellm_params( + { + "provider_configs": [ + {"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "region": "us-east-1"}, + {"aws_secret_access_key": "other-" + SENTINEL_AWS_SECRET_ACCESS_KEY, "region": "us-west-2"}, + ] + } + ) + + assert SENTINEL_AWS_SECRET_ACCESS_KEY not in json.dumps(redacted) + assert redacted["provider_configs"][0]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING + assert redacted["provider_configs"][0]["region"] == "us-east-1" + assert redacted["provider_configs"][1]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING + assert redacted["provider_configs"][1]["region"] == "us-west-2" + + +def test_redact_sensitive_agent_litellm_params_redacts_secrets_inside_model_list(): + """The exact shape flagged in review: litellm_params.model_list, where each + entry carries its own nested litellm_params with a provider credential.""" + redacted = redact_sensitive_agent_litellm_params( + { + "model_list": [ + { + "model_name": "gpt-4", + "litellm_params": {"api_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "model": "gpt-4"}, + }, + { + "model_name": "claude", + "litellm_params": { + "aws_secret_access_key": "other-" + SENTINEL_AWS_SECRET_ACCESS_KEY, + "model": "bedrock/claude", + }, + }, + ] + } + ) + + assert SENTINEL_AWS_SECRET_ACCESS_KEY not in json.dumps(redacted) + assert redacted["model_list"][0]["litellm_params"]["api_key"] == REDACTED_BY_LITELM_STRING + assert redacted["model_list"][0]["litellm_params"]["model"] == "gpt-4" + assert redacted["model_list"][1]["litellm_params"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING + assert redacted["model_list"][1]["litellm_params"]["model"] == "bedrock/claude" + + +def test_restore_redacted_litellm_params_preserves_secret_inside_model_list(): + """The write-side counterpart: a caller editing a model_list entry's own + non-secret field (renaming it) while leaving that same entry's nested + secret masked must not corrupt the stored per-deployment credential. + List entries correspond by position (see the module docstring on + ``_restore_redacted_nested_value``), so this -- the common "edit this + entry, keep its secret" pattern -- must keep working.""" + existing = { + "agent_name": "my-agent", + "model_list": [ + { + "model_name": "gpt-4", + "litellm_params": {"api_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "model": "gpt-4"}, + }, + ], + } + incoming = { + "agent_name": "my-agent-renamed", + "model_list": [ + { + "model_name": "gpt-4-renamed", + "litellm_params": {"api_key": REDACTED_BY_LITELM_STRING, "model": "gpt-4"}, + }, + ], + } + + restored = _restore_redacted_litellm_params(incoming, existing) + + assert SENTINEL_AWS_SECRET_ACCESS_KEY == restored["model_list"][0]["litellm_params"]["api_key"] + assert restored["model_list"][0]["model_name"] == "gpt-4-renamed" + assert restored["agent_name"] == "my-agent-renamed" + + +def test_restore_redacted_litellm_params_matches_list_entries_by_position(): + """Documents the accepted trade-off: a list has no stable per-element + identity in a plain ``dict[str, object]`` schema, so restoration matches + entries by index, the same correspondence every other part of this merge + (and the endpoints' full-replace-on-PUT semantics) already assumes. If a + caller both reorders the list AND echoes back a masked marker in the same + request, a credential can end up attached to a different logical entry. + That is a known, narrow limitation -- not a leak between different + agents or tenants, since it only reshuffles one agent's own stored + values -- and this test pins the current, deliberate behavior rather + than asserting it away.""" + existing = { + "model_list": [ + {"model_name": "gpt-4", "litellm_params": {"api_key": SENTINEL_AWS_SECRET_ACCESS_KEY}}, + {"model_name": "claude", "litellm_params": {"api_key": "other-" + SENTINEL_AWS_SECRET_ACCESS_KEY}}, + ], + } + incoming = { + "model_list": [ + # Same index (0) now holds what used to be at index 1's entry. + {"model_name": "claude", "litellm_params": {"api_key": REDACTED_BY_LITELM_STRING}}, + ], + } + + restored = _restore_redacted_litellm_params(incoming, existing) + + assert restored["model_list"][0]["litellm_params"]["api_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY + + +def test_restore_redacted_litellm_params_recovers_a_whole_subtree_collapsed_by_the_depth_cap(): + """Past the read-side recursion depth cap, a whole nested subtree is + collapsed to the flat REDACTED_BY_LITELM marker rather than a dict/list. + If the caller echoes that flat marker back unchanged, the whole + subtree -- not just the literal marker string -- must be restored.""" + existing_subtree = {"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "region": "us-east-1"} + incoming = {"provider_config": REDACTED_BY_LITELM_STRING} + existing = {"provider_config": existing_subtree} + + restored = _restore_redacted_litellm_params(incoming, existing) + + assert restored["provider_config"] == existing_subtree + + +def test_redact_sensitive_agent_litellm_params_does_not_reinterpret_plain_string_values_as_json(): + """A plain non-JSON string value (most string leaves) must pass through + unchanged rather than failing to parse and getting redacted.""" + redacted = redact_sensitive_agent_litellm_params({"model": "bedrock/agentcore/my-agent", "is_public": True}) + + assert redacted["model"] == "bedrock/agentcore/my-agent" + assert redacted["is_public"] is True + + +@pytest.mark.asyncio +async def test_add_agent_to_db_drops_a_sentinel_value_instead_of_storing_the_placeholder(): + """A create has nothing stored to restore behind a redaction marker, so a + sensitive key submitted as the literal marker is dropped rather than + persisted as the placeholder string itself.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + created_agent = MagicMock() + created_agent.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {}, + "object_permission": None, + } + created_agent.object_permission = None + mock_create = AsyncMock(return_value=created_agent) + mock_prisma.db.litellm_agentstable.create = mock_create + + await registry.add_agent_to_db( + agent={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": { + "aws_secret_access_key": REDACTED_BY_LITELM_STRING, + "model": "bedrock/agentcore/my-agent", + }, + }, + prisma_client=mock_prisma, + created_by="test-user", + ) + + stored_params: Final = json.loads(mock_create.call_args.kwargs["data"]["litellm_params"]) + assert "aws_secret_access_key" not in stored_params + assert stored_params["model"] == "bedrock/agentcore/my-agent" + + +@pytest.mark.asyncio +async def test_update_agent_in_db_preserves_secret_when_echoed_back_redacted(): + """PUT round-trips the GET response, which shows the secret redacted. Saving + an unrelated field change must not overwrite the real stored credential + with the redaction marker.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace( + litellm_params={ + "aws_access_key_id": SENTINEL_AWS_ACCESS_KEY_ID, + "aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, + "model": "bedrock/agentcore/my-agent", + }, + object_permission_id=None, + ) + ) + updated_agent = MagicMock() + updated_agent.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "Renamed Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {}, + "object_permission": None, + } + updated_agent.object_permission = None + mock_update = AsyncMock(return_value=updated_agent) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.update_agent_in_db( + agent_id="agent-123", + agent={ + "agent_name": "Renamed Agent", + "agent_card_params": _sample_agent_card_params(), + # The UI round-tripped the redacted secret and the untouched + # access key id verbatim; only agent_name actually changed. + "litellm_params": { + "aws_access_key_id": SENTINEL_AWS_ACCESS_KEY_ID, + "aws_secret_access_key": REDACTED_BY_LITELM_STRING, + "model": "bedrock/agentcore/my-agent", + }, + }, + prisma_client=mock_prisma, + updated_by="test-user", + ) + + stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"]) + assert stored_params["aws_secret_access_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY + assert stored_params["aws_access_key_id"] == SENTINEL_AWS_ACCESS_KEY_ID + assert stored_params["model"] == "bedrock/agentcore/my-agent" + + +@pytest.mark.asyncio +async def test_update_agent_in_db_preserves_secret_when_key_omitted_entirely(): + """Omitting the sensitive key altogether must fall back to the stored + value too, not just an explicit redaction-marker round-trip.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace( + litellm_params={"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY}, + object_permission_id=None, + ) + ) + updated_agent = MagicMock() + updated_agent.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {}, + "object_permission": None, + } + updated_agent.object_permission = None + mock_update = AsyncMock(return_value=updated_agent) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.update_agent_in_db( + agent_id="agent-123", + agent={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {"model": "bedrock/agentcore/my-agent"}, + }, + prisma_client=mock_prisma, + updated_by="test-user", + ) + + stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"]) + assert stored_params["aws_secret_access_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY + + +@pytest.mark.asyncio +async def test_update_agent_in_db_preserves_secret_nested_under_a_non_sensitive_key(): + """A secret nested inside a dict held by a non-sensitively-named key + (e.g. a per-provider sub-config) must also survive an echoed-back + redaction marker, not just top-level secret keys.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace( + litellm_params={ + "provider_config": { + "aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, + "region": "us-east-1", + } + }, + object_permission_id=None, + ) + ) + updated_agent = MagicMock() + updated_agent.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {}, + "object_permission": None, + } + updated_agent.object_permission = None + mock_update = AsyncMock(return_value=updated_agent) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.update_agent_in_db( + agent_id="agent-123", + agent={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": { + # The GET response redacted the nested secret; the caller + # round-trips it verbatim while changing nothing. + "provider_config": { + "aws_secret_access_key": REDACTED_BY_LITELM_STRING, + "region": "us-west-2", + } + }, + }, + prisma_client=mock_prisma, + updated_by="test-user", + ) + + stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"]) + assert stored_params["provider_config"]["aws_secret_access_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY + assert stored_params["provider_config"]["region"] == "us-west-2" + + +@pytest.mark.asyncio +async def test_update_agent_in_db_clears_secret_on_explicit_empty_value(): + """An explicit empty string is a deliberate clear, distinct from an omitted + key or the redaction marker, and must actually clear the stored secret.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace( + litellm_params={"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY}, + object_permission_id=None, + ) + ) + updated_agent = MagicMock() + updated_agent.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {}, + "object_permission": None, + } + updated_agent.object_permission = None + mock_update = AsyncMock(return_value=updated_agent) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.update_agent_in_db( + agent_id="agent-123", + agent={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {"aws_secret_access_key": ""}, + }, + prisma_client=mock_prisma, + updated_by="test-user", + ) + + stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"]) + assert stored_params["aws_secret_access_key"] == "" + + +@pytest.mark.asyncio +async def test_patch_agent_in_db_preserves_secret_when_litellm_params_omitted(): + """A PATCH that only renames the agent must not touch (let alone drop) the + stored litellm_params secret.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "Old Name", + "litellm_params": {"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY}, + "object_permission_id": None, + } + ) + patched_agent = MagicMock() + patched_agent.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "New Name", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY}, + "object_permission": None, + } + patched_agent.object_permission = None + mock_update = AsyncMock(return_value=patched_agent) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.patch_agent_in_db( + agent_id="agent-123", + agent={"agent_name": "New Name"}, + prisma_client=mock_prisma, + updated_by="test-user", + ) + + update_data: Final = mock_update.call_args.kwargs["data"] + assert "litellm_params" not in update_data + + +@pytest.mark.asyncio +async def test_patch_agent_in_db_preserves_secret_when_echoed_back_redacted(): + """A PATCH that includes litellm_params (e.g. to flip an unrelated flag) + with the secret round-tripped as the redaction marker must not clobber + the stored credential.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "Test Agent", + "litellm_params": { + "aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, + "is_public": False, + }, + "object_permission_id": None, + } + ) + patched_agent = MagicMock() + patched_agent.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {}, + "object_permission": None, + } + patched_agent.object_permission = None + mock_update = AsyncMock(return_value=patched_agent) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.patch_agent_in_db( + agent_id="agent-123", + agent={ + "litellm_params": { + "aws_secret_access_key": REDACTED_BY_LITELM_STRING, + "is_public": True, + } + }, + prisma_client=mock_prisma, + updated_by="test-user", + ) + + stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"]) + assert stored_params["aws_secret_access_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY + assert stored_params["is_public"] is True diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index ea196bda529..a78b3238a9a 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -4,6 +4,7 @@ import pytest from fastapi import FastAPI from fastapi.testclient import TestClient +from litellm.constants import REDACTED_BY_LITELM_STRING from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints import endpoints as agent_endpoints from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( @@ -484,11 +485,15 @@ class TestAgentRBACInternalUserViewOnly: assert resp.status_code == 403 +SENTINEL_AGENT_API_KEY = "sk-test-sentinel-do-not-use" + + class TestAgentRBACProxyAdminViewOnly: """Read-only proxy admins go through the object-permission scoped branch on GET /v1/agents (the admin fast path stays full PROXY_ADMIN only, so viewers - cannot fan out health checks beyond their allowlist), and secret unredaction - also stays gated on full PROXY_ADMIN.""" + cannot fan out health checks beyond their allowlist). litellm_params + secrets are redacted for every caller, admin included (LIT-6736); only the + virtual-key/header visibility stays gated on full PROXY_ADMIN.""" @pytest.fixture(autouse=True) def _setup(self, monkeypatch): @@ -501,7 +506,7 @@ class TestAgentRBACProxyAdminViewOnly: agent_id=f"agent-{index}", agent_name=f"Agent {index}", agent_card_params=_sample_agent_card_params(), - litellm_params={"api_key": "sk-super-secret-agent-key"}, + litellm_params={"api_key": SENTINEL_AGENT_API_KEY}, ) for index in (1, 2) ] @@ -544,7 +549,7 @@ class TestAgentRBACProxyAdminViewOnly: def test_should_still_redact_secrets_for_view_only_admin(self): """An unrestricted viewer sees the same agents as an admin but with keys - stripped and litellm_params masked.""" + stripped; litellm_params secrets never appear in either response.""" self.allowed_agents_spy.return_value = UnrestrictedAgentAccess() viewer_resp = self._list_agents(self.viewer_client) admin_resp = self._list_agents(self.admin_client) @@ -553,14 +558,12 @@ class TestAgentRBACProxyAdminViewOnly: viewer_by_id = {agent["agent_id"]: agent for agent in viewer_resp.json()} assert set(viewer_by_id) == {"agent-1", "agent-2"} assert viewer_by_id["agent-1"]["keys"] is None - assert "sk-super-secret-agent-key" not in viewer_resp.text + assert SENTINEL_AGENT_API_KEY not in viewer_resp.text admin_by_id = {agent["agent_id"]: agent for agent in admin_resp.json()} assert admin_by_id["agent-1"]["keys"][0]["token"] == "hash-aaa" - assert ( - admin_by_id["agent-1"]["litellm_params"]["api_key"] - == "sk-super-secret-agent-key" - ) + assert SENTINEL_AGENT_API_KEY not in admin_resp.text + assert admin_by_id["agent-1"]["litellm_params"]["api_key"] == REDACTED_BY_LITELM_STRING class TestAgentRBACProxyAdmin: @@ -616,6 +619,109 @@ class TestAgentRBACProxyAdmin: # Security scheme is the LiteLLM scheme. assert "LiteLLMKey" in stored_card["securitySchemes"] + def test_create_agent_response_never_echoes_secret(self): + """LIT-6736: POST /v1/agents must not echo the stored secret back, even + though it's the caller's own value and even for a proxy admin.""" + with patch("litellm.proxy.proxy_server.prisma_client"): # test-quality-ok: proxy_server module global is the endpoint's only injection point + self.mock_registry.get_agent_by_name = MagicMock(return_value=None) + self.mock_registry.add_agent_to_db = AsyncMock( + return_value=AgentResponse( + agent_id="agent-123", + agent_name="Test Agent", + agent_card_params=_sample_agent_card_params(), + litellm_params={ + "aws_secret_access_key": SENTINEL_AGENT_API_KEY, + "model": "bedrock/agentcore/my-agent", + }, + ) + ) + self.mock_registry.register_agent = MagicMock() + + resp = self.admin_client.post( + "/v1/agents", + json={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": { + "aws_secret_access_key": SENTINEL_AGENT_API_KEY, + "model": "bedrock/agentcore/my-agent", + }, + }, + headers={"Authorization": "Bearer k"}, + ) + + assert resp.status_code == 200 + assert SENTINEL_AGENT_API_KEY not in resp.text + body = resp.json() + assert body["litellm_params"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING + assert body["litellm_params"]["model"] == "bedrock/agentcore/my-agent" + + def test_update_agent_response_never_echoes_secret(self): + """LIT-6736: PUT /v1/agents/{id} must not echo the stored secret back.""" + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: # test-quality-ok: proxy_server module global is the endpoint's only injection point + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "Existing Agent", + "agent_card_params": _sample_agent_card_params(), + } + ) + self.mock_registry.update_agent_in_db = AsyncMock( + return_value=AgentResponse( + agent_id="agent-123", + agent_name="Test Agent", + agent_card_params=_sample_agent_card_params(), + litellm_params={"aws_secret_access_key": SENTINEL_AGENT_API_KEY}, + ) + ) + self.mock_registry.deregister_agent = MagicMock() + self.mock_registry.register_agent = MagicMock() + + resp = self.admin_client.put( + "/v1/agents/agent-123", + json={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {"aws_secret_access_key": REDACTED_BY_LITELM_STRING}, + }, + headers={"Authorization": "Bearer k"}, + ) + + assert resp.status_code == 200 + assert SENTINEL_AGENT_API_KEY not in resp.text + assert resp.json()["litellm_params"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING + + def test_patch_agent_response_never_echoes_secret(self): + """LIT-6736: PATCH /v1/agents/{id} must not echo the stored secret back.""" + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: # test-quality-ok: proxy_server module global is the endpoint's only injection point + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "Existing Agent", + "agent_card_params": _sample_agent_card_params(), + } + ) + self.mock_registry.patch_agent_in_db = AsyncMock( + return_value=AgentResponse( + agent_id="agent-123", + agent_name="Renamed Agent", + agent_card_params=_sample_agent_card_params(), + litellm_params={"aws_secret_access_key": SENTINEL_AGENT_API_KEY}, + ) + ) + self.mock_registry.deregister_agent = MagicMock() + self.mock_registry.register_agent = MagicMock() + + resp = self.admin_client.patch( + "/v1/agents/agent-123", + json={"agent_name": "Renamed Agent"}, + headers={"Authorization": "Bearer k"}, + ) + + assert resp.status_code == 200 + assert SENTINEL_AGENT_API_KEY not in resp.text + assert resp.json()["litellm_params"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING + def test_should_allow_admin_to_delete_agent(self): existing = { "agent_id": "agent-123",