mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(agents): redact secret litellm_params fields from all /v1/agents responses (#39389)
* fix(agents): redact secret litellm_params fields from all /v1/agents responses Secret-bearing litellm_params fields (aws_secret_access_key, api_key, and similar) are now write-only: list, get, create, update, and patch responses always replace them with a fixed marker, regardless of caller role. Editing an agent no longer requires resending a real credential -- an update that omits a sensitive field, or echoes the marker back, preserves the stored value; a real value still rotates it. * fix(agents): redact secrets nested inside dicts/lists in litellm_params too Greptile found that a secret nested one level down under a non-sensitively-named key, or inside a list of per-provider configs, was neither redacted on read nor restored symmetrically on write (the marker string could get persisted as the real value). Recurse into lists on the read side, and mirror that recursion on the write side so restoration isn't limited to top-level keys. Also fixes a regression the redact rewrite introduced (a plain string leaf like a model name was being misinterpreted as a JSON blob and redacted), and suppresses 3 new test-quality-gate findings on an established repo-wide mocking pattern this PR's new tests also use. * fix(agents): guard list-position credential restore against misassignment Two more real gaps Greptile/veria found in the recursive redact/restore mechanism, verified directly against the exact reported shape (litellm_params.model_list, each entry carrying its own nested litellm_params.api_key/aws_secret_access_key) before fixing: - Positional restoration inside a list could attach one entry's stored credential to a different entry if the list were reordered or resized between GET and PUT/PATCH. Restoration by index now only fires when the incoming and existing entries match on every non-secret field; otherwise the caller's own value is used (never a guessed cross-entry secret). - A subtree collapsed to the flat REDACTED_BY_LITELM marker by the read-side recursion depth cap couldn't be recovered on write (the marker string itself would get persisted). Restore now recognizes that shape and recovers the whole existing subtree. Both covered by regression tests mirroring the exact model_list shape reported, mutation-verified. * fix(agents): simplify list-entry credential restore to positional matching The content-match guard from the previous commit fixed one Greptile finding (cross-entry misassignment on reorder) but introduced a worse one: it also rejected restoration whenever an entry's own non-secret fields changed, which is the common case (rename a model_list entry while leaving its own secret masked) -- silently dropping the stored credential on an ordinary edit. There is no stable per-element identity in a plain dict[str, object] schema, so no rule can satisfy both 'restore whenever the entry itself only had its secret masked' and 'never restore across a reorder' at once. Positional correspondence is what every other part of this restore (and the endpoints' full-replace-on-PUT semantics) already assumes, so drop the content-match gate and rely on it here too: this fixes the common case correctly and accepts cross-entry misassignment on a simultaneous reorder-plus-masked-echo as a known, narrow, documented limitation (not a leak between different agents or tenants, since it only reshuffles one agent's own stored values). Tests updated to pin the accepted trade-off explicitly rather than asserting it away, and to cover the previously broken ordinary-edit case.
This commit is contained in:
Yassin Kortam
GitHub
parent
25991fe78a
commit
8e65265bb4
5 changed files with 879 additions and 37 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue