scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 114

fix(proxy): sync search tools into the router on management writes

Browse source 
Creating a search tool through the UI only wrote the row; the router was updated
solely by the add_deployment job, so the tool was unusable for up to
PROXY_CONFIG_RELOAD_INTERVAL_SECONDS (30s by default) even on the worker that
served the write. Tools declared in config.yaml load straight into the router at
startup, which is why they never showed the delay.

The create, update and delete endpoints now refresh the router inline, matching
what the MCP server endpoints already do. The refresh is best-effort: the row is
already committed, so a failure must not surface as a 500 and push the caller
into a retry that creates duplicates.

Two related gaps go with it. _init_search_tools_in_db skipped the router update
whenever the merged list came back empty, so deleting the last search tool left
it live in memory forever. And in store_model_in_db-off deployments the
add_deployment job is never scheduled, so DB-backed search tools never reached
the router at all; that branch now loads them at startup and keeps them fresh on
its own interval, the same way MCP servers already do.
This commit is contained in:
Yuneng Jiang 2026-08-26 11:46:54 -07:00
parent e52f05566d
commit 91e7eb115d
No known key found for this signature in database
4 changed files with 253 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
litellm/proxy/proxy_server.py
View file

@ -7470,11 +7470,9 @@ class ProxyConfig:
len(db_search_tools),
)
if llm_router is not None and search_tools:
if llm_router is not None:
await SearchAPIRouter.update_router_search_tools(router_instance=llm_router, search_tools=search_tools)
verbose_proxy_logger.info("Successfully loaded %s search tool(s) into router", len(search_tools))
elif llm_router is not None:
verbose_proxy_logger.debug("No search tools found in config or database, skipping router update")
else:
verbose_proxy_logger.debug(
"Router not initialized yet, search tools will be added when router is created"
@ -7485,6 +7483,19 @@ class ProxyConfig:
"litellm.proxy.proxy_server.py::ProxyConfig:_init_search_tools_in_db - %s", e
)
async def reload_search_tools_from_db(self) -> None:
"""Refresh this worker's router from the search tools table.
Driven by the management endpoints so the worker that served the write is correct
immediately, and by the periodic job in store_model_in_db-off deployments. Gated the same
way as startup, so an admin who excluded search_tools from supported_db_objects opts out.
"""
if not self._should_load_db_object(object_type="search_tools"):
return
if prisma_client is None:
return
await self._init_search_tools_in_db(prisma_client=prisma_client)
@staticmethod
def _merge_config_and_db_search_tools(
config_search_tools: list[SearchToolTypedDict],
@ -9104,7 +9115,18 @@ class ProxyStartupEvent:
if store_model_in_db is not True:
await proxy_config.init_mcp_servers_from_db()
# Without this branch's own refresh, a UI-created search tool never reaches the router:
# the add_deployment job that carries it in store_model_in_db=True mode is not scheduled.
await proxy_config.reload_search_tools_from_db()
if prisma_client is not None:
scheduler.add_job(
proxy_config.reload_search_tools_from_db,
"interval",
seconds=config_reload_interval_seconds,
id="reload_search_tools_job",
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
# DB-backed MCP servers are live objects in every mode, so the registry refresh that
# store_model_in_db=True deployments get via the add_deployment job must run here
# too; without it, a server whose OAuth discovery failed at startup is rebuilt only

28
litellm/proxy/search_endpoints/search_tool_management.py
View file

@ -51,6 +51,20 @@ def _convert_datetime_to_str(value: datetime | str | None) -> str | None:
TeamObjectLookup: TypeAlias = Callable[[str, UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable]]
async def _refresh_router_search_tools() -> None:
"""Push the search tools table into this worker's router.
Best-effort: the row is already committed, so a refresh failure must not surface as a 500 and
push the caller into a retry that creates duplicates.
"""
from litellm.proxy.proxy_server import proxy_config
try:
await proxy_config.reload_search_tools_from_db()
except Exception as e: # noqa: BLE001 # the row is committed; no refresh failure may reach the caller
verbose_proxy_logger.exception("Search tool router refresh failed after a management write: %s", e)
async def _team_object_from_db(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable:
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.proxy_server import (
@ -305,8 +319,10 @@ async def create_search_tool(request: CreateSearchToolRequest):
search_tool=request.search_tool, prisma_client=prisma_client
)
await _refresh_router_search_tools()
verbose_proxy_logger.debug(
"Successfully added search tool '%s' to database. Router will be updated by the cron job.",
"Successfully added search tool '%s' to database.",
result.get("search_tool_name"),
)
@ -388,8 +404,10 @@ async def update_search_tool(search_tool_id: str, request: UpdateSearchToolReque
prisma_client=prisma_client,
)
await _refresh_router_search_tools()
verbose_proxy_logger.debug(
"Successfully updated search tool '%s' in database. Router will be updated by the cron job.",
"Successfully updated search tool '%s' in database.",
result.get("search_tool_name"),
)
@ -445,9 +463,9 @@ async def delete_search_tool(search_tool_id: str):
search_tool_id=search_tool_id, prisma_client=prisma_client
)
verbose_proxy_logger.debug(
"Successfully deleted search tool from database. Router will be updated by the cron job."
)
await _refresh_router_search_tools()
verbose_proxy_logger.debug("Successfully deleted search tool from database.")
return result
except HTTPException as e:

156
tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py
View file

@ -992,3 +992,159 @@ async def test_list_search_tools_reports_a_missing_real_team_as_404():
assert response.status_code == 404
assert "search_tools" not in response.json()
# ---------------------------------------------------------------------------
# Router sync on management writes (LIT-3379)
#
# The proxy resolves prisma_client / proxy_config / llm_router from
# litellm.proxy.proxy_server module globals at call time and reaches its DB layer through a
# module-level registry singleton, so there is no constructor or parameter to inject through.
# Patching those globals is the only seam that exercises the endpoint end to end.
# ---------------------------------------------------------------------------
def _search_tool_row(name: str, provider: str = "tavily") -> dict:
return {
"search_tool_id": f"{name}-id",
"search_tool_name": name,
"litellm_params": {"search_provider": provider, "api_key": "sk-test"},
"search_tool_info": {"description": name},
}
def _fake_registry(db_rows: list) -> MagicMock:
"""A registry singleton whose writes land in db_rows, so the refresh reads back real state."""
async def _add(search_tool, **_):
row = _search_tool_row(
search_tool["search_tool_name"],
provider=search_tool.get("litellm_params", {}).get("search_provider", "tavily"),
)
db_rows.append(row)
return row
async def _update(search_tool_id, search_tool, **_):
row = _search_tool_row(
search_tool["search_tool_name"],
provider=search_tool.get("litellm_params", {}).get("search_provider", "tavily"),
)
db_rows[:] = [row if existing["search_tool_id"] == search_tool_id else existing for existing in db_rows]
return row
async def _delete(search_tool_id, **_):
db_rows[:] = [existing for existing in db_rows if existing["search_tool_id"] != search_tool_id]
return {"message": "deleted", "search_tool_name": search_tool_id}
async def _get_by_id(search_tool_id, **_):
return next((row for row in db_rows if row["search_tool_id"] == search_tool_id), None)
registry = MagicMock()
registry.add_search_tool_to_db = AsyncMock(side_effect=_add)
registry.update_search_tool_in_db = AsyncMock(side_effect=_update)
registry.delete_search_tool_from_db = AsyncMock(side_effect=_delete)
registry.get_search_tool_by_id_from_db = AsyncMock(side_effect=_get_by_id)
return registry
@contextlib.contextmanager
def _live_router_and_db(db_rows: list):
"""Drive the endpoints against a real ProxyConfig so the router refresh actually runs."""
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
proxy_config.update_config_state({})
fake_router = MagicMock()
fake_router.search_tools = list(db_rows)
with contextlib.ExitStack() as stack:
stack.enter_context(patch("litellm.proxy.proxy_server.prisma_client", MagicMock())) # test-quality-ok: proxy globals are the only seam; see the module note above
stack.enter_context(patch("litellm.proxy.proxy_server.proxy_config", proxy_config)) # test-quality-ok: proxy globals are the only seam; see the module note above
stack.enter_context(patch("litellm.proxy.proxy_server.llm_router", fake_router)) # test-quality-ok: proxy globals are the only seam; see the module note above
stack.enter_context(
patch( # test-quality-ok: proxy globals are the only seam; see the module note above
"litellm.proxy.search_endpoints.search_tool_management.SEARCH_TOOL_REGISTRY",
_fake_registry(db_rows),
)
)
stack.enter_context(
patch( # test-quality-ok: proxy globals are the only seam; see the module note above
"litellm.proxy.search_endpoints.search_tool_registry.SearchToolRegistry.get_all_search_tools_from_db",
AsyncMock(side_effect=lambda **_: list(db_rows)),
)
)
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
)
try:
yield fake_router
finally:
app.dependency_overrides.pop(user_api_key_auth, None)
@pytest.mark.asyncio
async def test_create_search_tool_reaches_the_router_before_the_response():
"""A UI-created tool must be usable immediately, not only after the next config reload tick."""
with _live_router_and_db([]) as fake_router:
response = TestClient(app).post(
"/search_tools",
json={
"search_tool": {
"search_tool_name": "tavily-search",
"litellm_params": {"search_provider": "tavily"},
}
},
)
assert response.status_code == 200
assert [tool["search_tool_name"] for tool in fake_router.search_tools] == ["tavily-search"]
@pytest.mark.asyncio
async def test_update_search_tool_reaches_the_router_before_the_response():
with _live_router_and_db([_search_tool_row("tavily-search", provider="tavily")]) as fake_router:
response = TestClient(app).put(
"/search_tools/tavily-search-id",
json={
"search_tool": {
"search_tool_name": "tavily-search",
"litellm_params": {"search_provider": "exa_ai"},
}
},
)
assert response.status_code == 200
assert fake_router.search_tools[0]["litellm_params"]["search_provider"] == "exa_ai"
@pytest.mark.asyncio
async def test_delete_search_tool_removes_it_from_the_router():
"""Deleting the last tool must clear the router; the old empty-list guard left it live."""
with _live_router_and_db([_search_tool_row("tavily-search")]) as fake_router:
response = TestClient(app).delete("/search_tools/tavily-search-id")
assert response.status_code == 200
assert fake_router.search_tools == []
@pytest.mark.asyncio
async def test_create_search_tool_survives_a_failing_router_refresh():
"""The row is already committed, so a refresh failure must not turn into a 500."""
with _live_router_and_db([]):
with patch( # test-quality-ok: forcing the refresh to fail needs the refresh itself replaced
"litellm.proxy.proxy_server.ProxyConfig.reload_search_tools_from_db",
AsyncMock(side_effect=RuntimeError("registry boom")),
):
response = TestClient(app).post(
"/search_tools",
json={
"search_tool": {
"search_tool_name": "tavily-search",
"litellm_params": {"search_provider": "tavily"},
}
},
)
assert response.status_code == 200
assert response.json()["search_tool_name"] == "tavily-search"

55
tests/test_litellm/proxy/proxy_server/test_proxy_config.py
View file

@ -1253,26 +1253,69 @@ async def test_ProxyConfig__init_search_tools_in_db_loads_merged_tools(monkeypat
@pytest.mark.asyncio
async def test_ProxyConfig__init_search_tools_in_db_skips_empty_router_update(monkeypatch):
async def test_ProxyConfig__init_search_tools_in_db_clears_router_when_last_tool_is_deleted(monkeypatch):
"""Deleting the last search tool must clear the router, not leave the tool live in memory."""
from litellm.proxy import proxy_server
from litellm.router_utils.search_api_router import SearchAPIRouter
pc = ProxyConfig()
pc.update_config_state({})
fake_router = MagicMock()
fake_router.search_tools = [{"search_tool_name": "deleted-search", "litellm_params": {}}]
mock_get_db_tools = AsyncMock(return_value=[])
mock_update_router = AsyncMock()
monkeypatch.setattr(proxy_server, "llm_router", MagicMock())
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
monkeypatch.setattr(
"litellm.proxy.search_endpoints.search_tool_registry.SearchToolRegistry.get_all_search_tools_from_db",
mock_get_db_tools,
)
monkeypatch.setattr(SearchAPIRouter, "update_router_search_tools", mock_update_router)
await pc._init_search_tools_in_db(prisma_client=MagicMock())
mock_get_db_tools.assert_awaited_once()
mock_update_router.assert_not_awaited()
assert fake_router.search_tools == []
@pytest.mark.asyncio
async def test_ProxyConfig_reload_search_tools_from_db_refreshes_router(monkeypatch):
from litellm.proxy import proxy_server
pc = ProxyConfig()
mock_init = AsyncMock()
monkeypatch.setattr(pc, "_init_search_tools_in_db", mock_init)
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
await pc.reload_search_tools_from_db()
mock_init.assert_awaited_once()
@pytest.mark.asyncio
async def test_ProxyConfig_reload_search_tools_from_db_honors_supported_db_objects(monkeypatch):
from litellm.proxy import proxy_server
pc = ProxyConfig()
mock_init = AsyncMock()
monkeypatch.setattr(pc, "_init_search_tools_in_db", mock_init)
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["models"]})
await pc.reload_search_tools_from_db()
mock_init.assert_not_awaited()
@pytest.mark.asyncio
async def test_ProxyConfig_reload_search_tools_from_db_noops_without_prisma(monkeypatch):
from litellm.proxy import proxy_server
pc = ProxyConfig()
mock_init = AsyncMock()
monkeypatch.setattr(pc, "_init_search_tools_in_db", mock_init)
monkeypatch.setattr(proxy_server, "prisma_client", None)
await pc.reload_search_tools_from_db()
mock_init.assert_not_awaited()
# ---------------------------------------------------------------------------