diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 89794fc2480..6bb4c91f6d5 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -1188,22 +1188,18 @@ async def make_agent_public( if agent is None: raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found") - 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): + stored_public_agent_groups: Final = await proxy_config.append_to_list_setting( + section_name="litellm_settings", + key="public_agent_groups", + item=agent.agent_id, + listed_as=AGENT_REGISTRY.ids_for_agent(agent.agent_id), + ) + if stored_public_agent_groups is None: raise HTTPException( status_code=400, detail=f"Agent with name {agent.agent_name} already in public agent groups", ) - updated_public_agent_groups: Final = [*current_public_agent_groups, agent.agent_id] - - if "litellm_settings" not in config or config["litellm_settings"] is None: - config["litellm_settings"] = {} - - config["litellm_settings"]["public_agent_groups"] = updated_public_agent_groups - - await proxy_config.save_config(new_config=config) + updated_public_agent_groups: Final = list(stored_public_agent_groups) litellm.public_agent_groups = updated_public_agent_groups diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a767f5423ed..1389ea46990 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5485,6 +5485,17 @@ _CONFIG_PERSISTED_SECTIONS: Final = ("general_settings", "router_settings", "lit _CONFIG_UNMANAGED_EXCLUSIONS: Final = frozenset(("environment_variables", "model_list")) _CONFIG_SECTION_VALUES: Final = TypeAdapter(Mapping[str, JsonValue]) _CONFIG_SECTION_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock(hashtext($1))" +_STRING_LIST: Final = TypeAdapter(list[str]) +_CONFIG_MAPPING: Final = TypeAdapter(Mapping[str, object]) + + +def _with_list_item( + section: Mapping[str, JsonValue], *, key: str, item: str, listed_as: frozenset[str] +) -> Mapping[str, JsonValue]: + current: Final = _STRING_LIST.validate_python(section.get(key) or []) + if not listed_as.isdisjoint(current): + return section + return MappingProxyType({**section, key: [*current, item]}) class _ConfigParamWhere(TypedDict): @@ -6016,6 +6027,63 @@ class ProxyConfig: changed_keys: Mapping[str, JsonValue], removed_keys: frozenset[str], prisma_client: PrismaClient, + ) -> Mapping[str, JsonValue] | None: + return await self._rewrite_config_section( + section_name=section_name, + prisma_client=prisma_client, + rewrite=lambda existing_section: MappingProxyType( + { + key: value + for key, value in chain( + ((key, value) for key, value in existing_section.items() if key not in removed_keys), + changed_keys.items(), + ) + } + ), + ) + + async def append_to_list_setting( + self, *, section_name: Section, key: str, item: str, listed_as: frozenset[str] + ) -> tuple[str, ...] | None: + if prisma_client is None or not ( + general_settings_view().get("store_model_in_db", False) is True or store_model_in_db + ): + return await self._append_to_file_list_setting( + section_name=section_name, key=key, item=item, listed_as=listed_as + ) + store: Final = self._settings_stores[section_name] + if store.owned_by_config(key): + if not listed_as.isdisjoint(_STRING_LIST.validate_python(store.get(key) or [])): + return None + self._raise_config_owned(section_name=section_name, rejected=(key,), store=store) + written_section: Final = await self._rewrite_config_section( + section_name=section_name, + prisma_client=prisma_client, + rewrite=lambda existing_section: _with_list_item(existing_section, key=key, item=item, listed_as=listed_as), + ) + if written_section is None: + return None + store.apply_db_row(section_name, written_section) + await invalidate_config_param(section_name) + return tuple(_STRING_LIST.validate_python(written_section[key])) + + async def _append_to_file_list_setting( + self, *, section_name: Section, key: str, item: str, listed_as: frozenset[str] + ) -> tuple[str, ...] | None: + config: Final = _CONFIG_MAPPING.validate_python(await self.get_config()) + section: Final = _CONFIG_SECTION_VALUES.validate_python(config.get(section_name) or {}) + updated_section: Final = _with_list_item(section, key=key, item=item, listed_as=listed_as) + if updated_section == section: + return None + await self.save_config(new_config={**config, section_name: dict(updated_section)}) + return tuple(_STRING_LIST.validate_python(updated_section[key])) + + async def _rewrite_config_section( + self, + *, + section_name: str, + prisma_client: PrismaClient, + rewrite: Callable[[Mapping[str, JsonValue]], Mapping[str, JsonValue]], ) -> Mapping[str, JsonValue] | None: async with prisma_client.tx() as tx: await tx.query_raw(_CONFIG_SECTION_LOCK_SQL, section_name) @@ -6030,15 +6098,7 @@ class ProxyConfig: if isinstance(existing_value, Mapping) else MappingProxyType({}) ) - merged_section: Final[Mapping[str, JsonValue]] = MappingProxyType( - { - key: value - for key, value in chain( - ((key, value) for key, value in existing_section.items() if key not in removed_keys), - changed_keys.items(), - ) - } - ) + merged_section: Final = rewrite(existing_section) if merged_section == existing_section: return None serialized_section: Final = json.dumps(dict(merged_section)) diff --git a/tests/unit/proxy/agent_endpoints/test_endpoints.py b/tests/unit/proxy/agent_endpoints/test_endpoints.py index 81b6c09ca12..8eecdf02fb8 100644 --- a/tests/unit/proxy/agent_endpoints/test_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_endpoints.py @@ -1,6 +1,9 @@ +import asyncio import json from collections.abc import Mapping +from dataclasses import dataclass from datetime import datetime, timezone +from pathlib import Path from types import SimpleNamespace from typing import Final @@ -8,6 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import yaml from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient from prisma.models import LiteLLM_AgentsTable @@ -15,6 +19,7 @@ from prisma.models import LiteLLM_AgentsTable from litellm.constants import REDACTED_BY_LITELM_STRING from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints import endpoints as agent_endpoints +from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( RestrictedAgentAccess, ) @@ -1143,77 +1148,164 @@ def test_merged_agent_card_url_has_no_double_slash_without_proxy_base_url( 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 - proxy_config: Final = ProxyConfig() - db_values: Final = proxy_config._prepared_db_settings_values("litellm_settings", db_param_value) - proxy_config._apply_litellm_settings_db_values(db_values) - return {"litellm_settings": dict(proxy_config.litellm_settings.resolved())} - - 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 {}) +@dataclass(frozen=True, slots=True) +class _StoredConfigRow: + param_name: str + param_value: dict[str, object] -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.""" +class _ConfigRows: + """`LiteLLM_Config` as a section-locked transaction reads it, shared by every replica.""" + + def __init__(self, rows: Mapping[str, Mapping[str, object]]) -> None: + self.serialized_rows: dict[str, str] = {name: json.dumps(dict(value)) for name, value in rows.items()} + self.section_lock: Final = asyncio.Lock() + + async def find_first(self, *, where: Mapping[str, str]) -> _StoredConfigRow | None: + name: Final = where["param_name"] + serialized: Final = self.serialized_rows.get(name) + await asyncio.sleep(0) + return None if serialized is None else _StoredConfigRow(param_name=name, param_value=json.loads(serialized)) + + async def upsert(self, *, where: Mapping[str, str], data: Mapping[str, Mapping[str, str]]) -> None: + self.serialized_rows[where["param_name"]] = data["update"]["param_value"] + + def stored(self, name: str) -> dict[str, object]: + return json.loads(self.serialized_rows[name]) + + +class _ConfigTransaction: + def __init__(self, rows: _ConfigRows) -> None: + self.litellm_config: Final = rows + self.holds_section_lock = False + + async def __aenter__(self) -> "_ConfigTransaction": + return self + + async def __aexit__(self, *_: object) -> None: + if self.holds_section_lock: + self.litellm_config.section_lock.release() + + async def query_raw(self, query: str, section_name: str) -> None: + await self.litellm_config.section_lock.acquire() + self.holds_section_lock = True + + +class _StaleReplicaPrisma: + """A replica whose config reads still return what it cached before another replica wrote.""" + + def __init__(self, rows: _ConfigRows, cached_rows: Mapping[str, Mapping[str, object]]) -> None: + self.rows: Final = rows + self._cached_rows: Final = cached_rows + + async def get_generic_data(self, key: str, value: str, table_name: str) -> SimpleNamespace | None: + cached: Final = self._cached_rows.get(value) + return None if cached is None else SimpleNamespace(param_name=value, param_value=dict(cached)) + + def tx(self) -> _ConfigTransaction: + return _ConfigTransaction(self.rows) + + +_PROXY_ADMIN: Final = UserAPIKeyAuth(user_id="test-user", user_role=LitellmUserRoles.PROXY_ADMIN) + + +def _replica_that_missed_a_publish( + monkeypatch: pytest.MonkeyPatch, *, stored_public_agent_groups: list[str], store_model_in_db: bool = True +) -> tuple[_ConfigRows, AgentRegistry]: import litellm from litellm.proxy.agent_endpoints import agent_registry as agent_registry_module - from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry + from litellm.proxy.proxy_server import ProxyConfig 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")) + rows: Final = _ConfigRows({"litellm_settings": {"public_agent_groups": stored_public_agent_groups}}) + cached_rows: Final = {"litellm_settings": {"public_agent_groups": []}} 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()) + monkeypatch.setattr(litellm, "public_agent_groups", []) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _StaleReplicaPrisma(rows, cached_rows)) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", store_model_in_db) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"store_model_in_db": store_model_in_db}) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", ProxyConfig()) + return rows, registry - 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"] +@pytest.mark.asyncio +async def test_make_agent_public_keeps_an_agent_another_replica_published(monkeypatch: pytest.MonkeyPatch) -> None: + rows, registry = _replica_that_missed_a_publish(monkeypatch, stored_public_agent_groups=["agent-1"]) + + response: Final = await agent_endpoints.make_agent_public(agent_id="agent-2", user_api_key_dict=_PROXY_ADMIN) + + assert response["public_agent_groups"] == ["agent-1", "agent-2"] + assert rows.stored("litellm_settings")["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 +@pytest.mark.asyncio +async def test_make_agent_public_concurrent_publishes_all_stay_public(monkeypatch: pytest.MonkeyPatch) -> None: + rows, _ = _replica_that_missed_a_publish(monkeypatch, stored_public_agent_groups=[]) - 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"]}), + await asyncio.gather( + agent_endpoints.make_agent_public(agent_id="agent-1", user_api_key_dict=_PROXY_ADMIN), + agent_endpoints.make_agent_public(agent_id="agent-2", user_api_key_dict=_PROXY_ADMIN), ) - duplicate: Final = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"}) + assert sorted(rows.stored("litellm_settings")["public_agent_groups"]) == ["agent-1", "agent-2"] - assert duplicate.status_code == 400 - assert "already in public agent groups" in duplicate.json()["detail"] + +@pytest.mark.asyncio +async def test_make_agent_public_rejects_an_agent_another_replica_published(monkeypatch: pytest.MonkeyPatch) -> None: + rows, _ = _replica_that_missed_a_publish(monkeypatch, stored_public_agent_groups=["agent-1"]) + + with pytest.raises(HTTPException) as duplicate: + await agent_endpoints.make_agent_public(agent_id="agent-1", user_api_key_dict=_PROXY_ADMIN) + + assert duplicate.value.status_code == 400 + assert "already in public agent groups" in duplicate.value.detail + assert rows.stored("litellm_settings")["public_agent_groups"] == ["agent-1"] + + +@pytest.mark.asyncio +async def test_make_agent_public_refuses_a_list_the_config_file_owns( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + from litellm.proxy import proxy_server + + rows, _ = _replica_that_missed_a_publish(monkeypatch, stored_public_agent_groups=[]) + config_file: Final = tmp_path / "config.yaml" + config_file.write_text("litellm_settings:\n public_agent_groups: [agent-1]\n") + monkeypatch.setattr("litellm.proxy.proxy_server.user_config_file_path", str(config_file)) + await proxy_server.proxy_config.get_config() + + with pytest.raises(HTTPException) as already_public: + await agent_endpoints.make_agent_public(agent_id="agent-1", user_api_key_dict=_PROXY_ADMIN) + with pytest.raises(HTTPException) as refused: + await agent_endpoints.make_agent_public(agent_id="agent-2", user_api_key_dict=_PROXY_ADMIN) + + assert already_public.value.status_code == 400 + assert "already in public agent groups" in already_public.value.detail + assert refused.value.status_code == 400 + assert refused.value.detail["keys"] == ["public_agent_groups"] + assert rows.stored("litellm_settings")["public_agent_groups"] == [] + + +@pytest.mark.asyncio +async def test_make_agent_public_appends_to_the_config_file_without_db_settings( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _replica_that_missed_a_publish(monkeypatch, stored_public_agent_groups=[], store_model_in_db=False) + config_file: Final = tmp_path / "config.yaml" + config_file.write_text("litellm_settings:\n public_agent_groups: [agent-1]\n") + monkeypatch.setattr("litellm.proxy.proxy_server.user_config_file_path", str(config_file)) + + response: Final = await agent_endpoints.make_agent_public(agent_id="agent-2", user_api_key_dict=_PROXY_ADMIN) + with pytest.raises(HTTPException) as duplicate: + await agent_endpoints.make_agent_public(agent_id="agent-1", user_api_key_dict=_PROXY_ADMIN) + + assert response["public_agent_groups"] == ["agent-1", "agent-2"] + assert yaml.safe_load(config_file.read_text())["litellm_settings"]["public_agent_groups"] == ["agent-1", "agent-2"] + assert duplicate.value.status_code == 400 @pytest.mark.parametrize("enabled, claim_field, expected", [(True, "azp", True), (False, "azp", False), (True, None, False)])