Merge 44c25aafeb into e9a6a25fe3

litellm/proxy/agent_endpoints/agent_registry.py

@@ -88,6 +88,9 @@ class AgentRecord(Protocol):
     @property
     def kill_switch(self) -> Mapping[str, object] | None: ...
 
+    @property
+    def static_headers(self) -> Mapping[str, object] | None: ...
+
     @property
     def spend(self) -> float: ...
 
@@ -302,6 +305,39 @@ def serialize_agent_kill_switch(incoming: object, existing: object) -> str:
     return safe_dumps(restored.model_dump() if restored is not None else None)
 
 
+_STATIC_HEADERS_ADAPTER: Final[TypeAdapter[Mapping[str, str]]] = TypeAdapter(Mapping[str, str])
+_EMPTY_STATIC_HEADERS: Final[Mapping[str, str]] = MappingProxyType({})
+
+
+def redact_agent_static_headers(
+    headers: Mapping[str, str] | None,
+) -> dict[str, str] | None:  # mutable-ok: AgentResponse.static_headers is declared as a plain dict, not Mapping
+    if headers is None:
+        return None
+    return {name: REDACTED_BY_LITELM_STRING for name in headers}
+
+
+def _parse_agent_static_headers(value: object) -> Mapping[str, str]:
+    if value is None:
+        return _EMPTY_STATIC_HEADERS
+    try:
+        if isinstance(value, str):
+            return _STATIC_HEADERS_ADAPTER.validate_json(value)
+        return _STATIC_HEADERS_ADAPTER.validate_python(value)
+    except ValidationError:
+        return _EMPTY_STATIC_HEADERS
+
+
+def serialize_agent_static_headers(incoming: Mapping[str, str] | None, existing: object) -> str:
+    """A header echoed back as the redaction marker keeps its stored value, or is dropped if none is stored"""
+    stored: Final = _parse_agent_static_headers(existing)
+    restored: Final = (
+        (name, stored.get(name) if value == REDACTED_BY_LITELM_STRING else value)
+        for name, value in (incoming or _EMPTY_STATIC_HEADERS).items()
+    )
+    return safe_dumps({name: value for name, value in restored if value is not None})
+
+
 _MISSING_AGENT_PARAM: Final = object()
 _RESTORE_AGENT_PARAMS_MAX_DEPTH: Final = 10
 
@@ -615,7 +651,9 @@ class AgentRegistry:
 
             # Serialize static_headers
             static_headers_obj: Final = agent.get("static_headers")
-            static_headers_val: Final[str | None] = safe_dumps(dict(static_headers_obj)) if static_headers_obj else None
+            static_headers_val: Final[str | None] = (
+                serialize_agent_static_headers(static_headers_obj, None) if static_headers_obj else None
+            )
 
             extra_headers_val: Final = agent.get("extra_headers")
             access_group_ids_val: Final = agent.get("access_group_ids")
@@ -750,8 +788,9 @@ class AgentRegistry:
                 if rate_field in agent:
                     update_data[rate_field] = agent.get(rate_field)
             if "static_headers" in agent:
-                headers_value: Final = agent.get("static_headers")
-                update_data["static_headers"] = safe_dumps(dict(headers_value) if headers_value is not None else {})
+                update_data["static_headers"] = serialize_agent_static_headers(
+                    agent.get("static_headers"), existing_agent.get("static_headers")
+                )
             if "extra_headers" in agent:
                 extra_headers_value: Final = agent.get("extra_headers")
                 update_data["extra_headers"] = extra_headers_value if extra_headers_value is not None else []
@@ -827,9 +866,8 @@ class AgentRegistry:
             agent_card_params: Final[str] = safe_dumps(agent_card_params_dict)
 
             # Serialize static_headers for update
-            static_headers_obj_u: Final = agent.get("static_headers")
-            static_headers_val_u: Final[str] = (
-                safe_dumps(dict(static_headers_obj_u)) if static_headers_obj_u is not None else safe_dumps({})
+            static_headers_val_u: Final[str] = serialize_agent_static_headers(
+                agent.get("static_headers"), existing_row.static_headers if existing_row is not None else None
             )
             extra_headers_val_u: Final = agent.get("extra_headers") or []
             access_group_ids_val_u: Final = tuple(dict.fromkeys(agent.get("access_group_ids") or ()))

litellm/proxy/agent_endpoints/endpoints.py

@@ -40,6 +40,7 @@ from litellm.proxy.agent_endpoints.agent_registry import (
     AgentIdWhere,
     parse_agent_kill_switch,
     parse_agent_litellm_params,
+    redact_agent_static_headers,
     redact_sensitive_agent_litellm_params,
 )
 from litellm.proxy.agent_endpoints.agent_search import (
@@ -215,10 +216,10 @@ def _redact_sensitive_agent_fields(
 ) -> list[AgentResponse]:
     """
     Return copies of the given agents with credential-bearing litellm_params
-    values and kill-switch auth secrets replaced by a fixed marker (never
-    returned to ANY caller, admin included) and, for non-admin callers,
-    virtual-key, header and kill-switch fields stripped entirely. The original
-    objects are not modified.
+    values, static header values and kill-switch auth secrets replaced by a
+    fixed marker (never returned to ANY caller, admin included) and, for
+    non-admin callers, virtual-key, header and kill-switch fields stripped
+    entirely. The original objects are not modified.
     """
     from litellm.proxy.proxy_server import general_settings, jwt_handler
 
@@ -242,6 +243,7 @@ def _redact_sensitive_agent_fields(
                 )
         if copy.litellm_params:
             copy.litellm_params = _redact_agent_litellm_params_dict(copy.litellm_params)
+        copy.static_headers = redact_agent_static_headers(copy.static_headers)
         copy.kill_switch = redact_kill_switch(copy.kill_switch)
         redacted.append(copy)
     return redacted

tests/unit/proxy/agent_endpoints/test_agent_registry.py

@@ -738,6 +738,7 @@ async def test_add_agent_to_db_drops_a_sentinel_value_instead_of_storing_the_pla
                 "aws_secret_access_key": REDACTED_BY_LITELM_STRING,
                 "model": "bedrock/agentcore/my-agent",
             },
+            "static_headers": {"Authorization": REDACTED_BY_LITELM_STRING, "X-Region": "eu"},
         },
         prisma_client=mock_prisma,
         created_by="test-user",
@@ -746,6 +747,7 @@ async def test_add_agent_to_db_drops_a_sentinel_value_instead_of_storing_the_pla
     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"
+    assert json.loads(mock_create.call_args.kwargs["data"]["static_headers"]) == {"X-Region": "eu"}
 
 
 @pytest.mark.asyncio

tests/unit/proxy/agent_endpoints/test_endpoints.py

@@ -221,7 +221,46 @@ def test_identity_settings_edit_preserves_runtime_configuration_on_readback(
     )
     preserved: Final = {*runtime, "agent_name", "enabled", "agent_id"}
     assert stored.model_dump(include=preserved) == expected.model_dump(include=preserved)
-    assert {key: readback.json()[key] for key in preserved} == expected.model_dump(mode="json", include=preserved)
+    shown: Final = expected.model_copy(update={"static_headers": {"X-Runtime": REDACTED_BY_LITELM_STRING}})
+    assert {key: readback.json()[key] for key in preserved} == shown.model_dump(mode="json", include=preserved)
+
+
+@pytest.mark.parametrize("method", ["PUT", "PATCH"])
+def test_edit_echoing_masked_static_headers_keeps_the_stored_values(
+    monkeypatch: pytest.MonkeyPatch, method: str
+) -> None:
+    from litellm.proxy import proxy_server
+    from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
+    from tests.unit.proxy.agent_endpoints.test_agent_registry import _stored_agent_row
+
+    secret_headers: Final = {"Authorization": f"Bearer {SENTINEL_AGENT_API_KEY}"}
+    table: Final = _AgentPersistence(
+        _stored_agent_row({"agent_card_params": _sample_agent_card_params(), "static_headers": secret_headers})
+    )
+    database: Final = SimpleNamespace(
+        litellm_agentstable=table,
+        litellm_verificationtoken=SimpleNamespace(find_many=AsyncMock(return_value=[])),
+    )
+    monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=database, writer_db=database))
+    monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", AgentRegistry())
+
+    shown: Final = client.get("/v1/agents/agent-123").json()["static_headers"]
+    assert shown == {"Authorization": REDACTED_BY_LITELM_STRING}
+
+    response: Final = client.request(
+        method,
+        "/v1/agents/agent-123",
+        json={
+            "agent_name": "Renamed agent",
+            "agent_card_params": _sample_agent_card_params(),
+            "static_headers": {**shown, "X-Region": "eu", "X-Unknown": REDACTED_BY_LITELM_STRING},
+        },
+    )
+
+    assert response.status_code == 200, response.text
+    assert SENTINEL_AGENT_API_KEY not in response.text
+    stored: Final = AgentResponse.model_validate(table.row.model_dump())
+    assert stored.static_headers == {**secret_headers, "X-Region": "eu"}
 
 
 def test_get_agent_by_id_not_found(
@@ -599,6 +638,7 @@ class TestAgentRBACProxyAdminViewOnly:
                 agent_name=f"Agent {index}",
                 agent_card_params=_sample_agent_card_params(),
                 litellm_params={"api_key": SENTINEL_AGENT_API_KEY},
+                static_headers={"Authorization": f"Bearer {SENTINEL_AGENT_API_KEY}"},
             )
             for index in (1, 2)
         ]
@@ -656,6 +696,7 @@ class TestAgentRBACProxyAdminViewOnly:
         assert admin_by_id["agent-1"]["keys"][0]["token"] == "hash-aaa"
         assert SENTINEL_AGENT_API_KEY not in admin_resp.text
         assert admin_by_id["agent-1"]["litellm_params"]["api_key"] == REDACTED_BY_LITELM_STRING
+        assert admin_by_id["agent-1"]["static_headers"] == {"Authorization": REDACTED_BY_LITELM_STRING}
 
 
 class TestAgentRBACProxyAdmin: