diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 3b8151d1064..cc17672553b 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/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: diff --git a/tests/e2e/ui/tests/modelHub/modelHub.spec.ts b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts index 6877fc9c48d..1fa3e4c530e 100644 --- a/tests/e2e/ui/tests/modelHub/modelHub.spec.ts +++ b/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 diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index a78b3238a9a..067f5a9f64c 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/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"]