Merge pull request #39554 from BerriAI/litellm_fix_flaky_model_hub_e2e

fix(agents): keep the published agent in public_agent_groups

litellm/proxy/agent_endpoints/endpoints.py

@@ -939,35 +939,32 @@ async def make_agent_public(
             if agent is None:
                 raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found")
 
-        if litellm.public_agent_groups is None:
-            litellm.public_agent_groups = []
-        # handle duplicates
-        if not AGENT_REGISTRY.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups):
+        config: Final = await proxy_config.get_config()
+
+        current_public_agent_groups: Final = list(litellm.public_agent_groups or [])
+        if not AGENT_REGISTRY.ids_for_agent(agent.agent_id).isdisjoint(current_public_agent_groups):
             raise HTTPException(
                 status_code=400,
                 detail=f"Agent with name {agent.agent_name} already in public agent groups",
             )
-        litellm.public_agent_groups.append(agent.agent_id)
+        updated_public_agent_groups: Final = [*current_public_agent_groups, agent.agent_id]
 
-        # Load existing config
-        config: Final = await proxy_config.get_config()
-
-        # Update config with new settings
         if "litellm_settings" not in config or config["litellm_settings"] is None:
             config["litellm_settings"] = {}
 
-        config["litellm_settings"]["public_agent_groups"] = litellm.public_agent_groups
+        config["litellm_settings"]["public_agent_groups"] = updated_public_agent_groups
 
-        # Save the updated config
         await proxy_config.save_config(new_config=config)
 
+        litellm.public_agent_groups = updated_public_agent_groups
+
         verbose_proxy_logger.debug(
-            "Updated public agent groups to: %s by user: %s", litellm.public_agent_groups, user_api_key_dict.user_id
+            "Updated public agent groups to: %s by user: %s", updated_public_agent_groups, user_api_key_dict.user_id
         )
 
         return {
             "message": "Successfully updated public agent groups",
-            "public_agent_groups": litellm.public_agent_groups,
+            "public_agent_groups": updated_public_agent_groups,
             "updated_by": user_api_key_dict.user_id,
         }
     except HTTPException:

tests/e2e/ui/tests/modelHub/modelHub.spec.ts

@@ -23,10 +23,12 @@ test.describe("AI Hub (internal admin view)", () => {
     await expect(modal.getByText(/Select All \(\d+\)/)).toBeVisible({ timeout: 5_000 });
 
     // Step 1: pick the seeded models via "Select All"
-    await modal.getByText(/Select All/i).click();
+    await modal.getByRole("checkbox", { name: /Select All/ }).check();
 
     // Move to confirm step
-    await modal.getByRole("button", { name: "Next" }).click();
+    const next = modal.getByRole("button", { name: "Next" });
+    await expect(next).toBeEnabled();
+    await next.click();
     await expect(modal.getByText("Confirm Making Models Public")).toBeVisible({ timeout: 5_000 });
 
     // Submit

tests/test_litellm/proxy/agent_endpoints/test_endpoints.py

@@ -1,3 +1,5 @@
+import json
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -1062,3 +1064,77 @@ def test_merged_agent_card_url_has_no_double_slash_without_proxy_base_url(
     interface_url = merged["supportedInterfaces"][0]["url"]
     assert interface_url == f"{base_url.rstrip('/')}/a2a/agent-xyz"
     assert "//a2a" not in interface_url
+
+
+class _DbBackedProxyConfig:
+    """Round-trips `litellm_settings` through the DB overlay the proxy applies on every
+    `get_config()`, which is what re-assigns the `litellm.public_*` globals in production.
+
+    Storage goes through JSON the way the `litellm_config` row does, so every read hands back
+    freshly built values instead of the objects the endpoint still holds a reference to."""
+
+    def __init__(self, stored_litellm_settings: dict[str, object] | None = None) -> None:
+        self.stored_litellm_settings_json: str = json.dumps(stored_litellm_settings or {})
+
+    async def get_config(self) -> dict[str, dict[str, object]]:
+        from litellm.proxy.proxy_server import ProxyConfig
+
+        config: Final[dict[str, dict[str, object]]] = {"litellm_settings": {}}
+        db_param_value: Final[dict[str, object]] = json.loads(self.stored_litellm_settings_json)
+        if not db_param_value:
+            return config
+        return ProxyConfig()._update_config_fields(
+            current_config=config,
+            param_name="litellm_settings",
+            db_param_value=db_param_value,
+        )
+
+    async def save_config(self, new_config: dict[str, dict[str, object]]) -> None:
+        self.stored_litellm_settings_json = json.dumps(new_config.get("litellm_settings") or {})
+
+
+def test_make_agent_public_twice_keeps_both_agents_public(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A second /make_public call must not drop the agent published by the first one."""
+    import litellm
+    from litellm.proxy.agent_endpoints import agent_registry as agent_registry_module
+    from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
+
+    registry: Final = AgentRegistry()
+    registry.register_agent(_sample_agent_response(agent_id="agent-1", agent_name="Agent One"))
+    registry.register_agent(_sample_agent_response(agent_id="agent-2", agent_name="Agent Two"))
+
+    monkeypatch.setattr(agent_registry_module, "global_agent_registry", registry)
+    monkeypatch.setattr(litellm, "public_agent_groups", None)
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
+    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", _DbBackedProxyConfig())
+
+    first: Final = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"})
+    second: Final = client.post("/v1/agents/agent-2/make_public", headers={"Authorization": "Bearer test-key"})
+
+    assert first.status_code == 200
+    assert second.status_code == 200
+    assert second.json()["public_agent_groups"] == ["agent-1", "agent-2"]
+    assert [agent.agent_id for agent in registry.get_public_agent_list()] == ["agent-1", "agent-2"]
+
+
+def test_make_agent_public_rejects_an_agent_published_only_in_the_db(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The duplicate guard must fire off the stored list, not just what this process published."""
+    import litellm
+    from litellm.proxy.agent_endpoints import agent_registry as agent_registry_module
+    from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
+
+    registry: Final = AgentRegistry()
+    registry.register_agent(_sample_agent_response(agent_id="agent-1", agent_name="Agent One"))
+
+    monkeypatch.setattr(agent_registry_module, "global_agent_registry", registry)
+    monkeypatch.setattr(litellm, "public_agent_groups", None)
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.proxy_config",
+        _DbBackedProxyConfig({"public_agent_groups": ["agent-1"]}),
+    )
+
+    duplicate: Final = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"})
+
+    assert duplicate.status_code == 400
+    assert "already in public agent groups" in duplicate.json()["detail"]