From 91b2013644a95fb06ea58bdadf5a47b74a1a43b9 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 4 Jun 2026 20:40:40 -0700 Subject: [PATCH] fix(proxy): stop team BYOK model name corruption on model edit (#29731) * fix(proxy): stop team model name corruption on edit (#28382) (#29001) Team-scoped ("Team-BYOK") models store an internal routing key model_name_{team_id}_{uuid} in the model_name column and the user-facing name in model_info.team_public_model_name. The internal name leaked into /v1, /v2, and /model/info responses; the dashboard bound its edit form to it, so any non-rename save (e.g. a TPM tweak) PATCHed the internal name back. The update path then treated it as a rename, overwriting team_public_model_name and rewriting the team's models[] ACL with the mangled string -- breaking team key calls with team_model_access_denied. Two-layer fix: - Read path (root cause): add _translate_model_name_for_response and apply it in model_info_v2 and _get_proxy_model_info so /v1, /v2, and /model/info surface the public name for team-scoped rows. The DB column and router index keep the internal name as the routing key; this is a presentation-layer swap on a shallow copy (never mutates input). - Write path (defense in depth): harden _get_public_model_name so a value matching the internal shape, or a no-op against the current DB column, is never treated as a rename -- for both the top-level model_name and an explicit model_info.team_public_model_name. Tests: regression for the reported scenario, full branch coverage of _get_public_model_name, two internal-shape guard cases, an end-to-end PATCH through _update_team_model_in_db (asserts the team ACL is untouched), and four response-translation cases. 60 passed (model management), 181 passed (proxy server). * fix(ui): key Agent Builder agent selection on model_info.id (#29729) * fix(ui): key Agent Builder agent selection on model_info.id Once team-scoped BYOK models can share a public name (the backend now returns the public name on /model/info instead of the internal routing key), selecting agents by model_name collides. Key selection, create, update and delete on the stable model_info.id instead, falling back to model_name only for config-defined agents that have no id. * fix(ui): add name-match fallback to post-create agent selection If the just-created agent's id is not yet present in the re-fetched list, try matching by name before falling back to the first agent. Addresses greptile review on #29729. --------- Co-authored-by: tushar8408 <32977767+tushar8408@users.noreply.github.com> (cherry picked from commit 56aa55b991373632fec89cd5b408605457aa7686) --- .../model_management_endpoints.py | 42 ++- litellm/proxy/proxy_server.py | 39 ++- .../test_model_management_endpoints.py | 301 ++++++++++++++++++ .../test_team_model_name_translation.py | 178 +++++++++++ .../playground/chat_ui/AgentBuilderView.tsx | 154 +++++---- 5 files changed, 648 insertions(+), 66 deletions(-) create mode 100644 tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 722fcd30033..404ed4491e5 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -490,9 +490,45 @@ def _get_public_model_name( patch_data: updateDeployment, db_model: Deployment, ) -> str: - """Determine the public model name from patch or existing model.""" - if patch_data.model_name: - return patch_data.model_name + """Determine the public model name from patch or existing model. + + The top-level ``model_name`` is the rename channel. For team-scoped rows + the DB ``model_name`` column holds an internal routing key + (``model_name_{team_id}_{uuid}``), and ``/model/info`` historically leaked + it into the dashboard edit form, so a non-rename save (e.g. a TPM tweak) + would PATCH the internal name and the update path would treat it as a + rename -- overwriting ``team_public_model_name`` and rewriting the team ACL + (see issue #28382). + + Guard against that by ignoring an incoming ``model_name`` that matches the + internal shape, or is a no-op against the current DB column. Anything else + is a genuine rename and wins. We deliberately do NOT read + ``patch_data.model_info.team_public_model_name``: the dashboard passes the + existing ``model_info`` blob through untouched on a rename, so honoring it + would return the OLD public name and silently drop the rename. + + Precedence (highest first): + 1. patch_data.model_name -- a genuine rename: not internal-shape and not a + no-op against db_model.model_name. + 2. db_model.model_info.team_public_model_name -- existing public name. + 3. db_model.model_name -- last-resort fallback for legacy rows. + """ + team_id = (patch_data.model_info.team_id if patch_data.model_info else None) or ( + db_model.model_info.team_id if db_model.model_info else None + ) + + def _is_internal_shape(name: Optional[str]) -> bool: + if team_id is None or not name: + return False + return name.startswith(f"model_name_{team_id}_") + + incoming = patch_data.model_name + if ( + incoming + and not _is_internal_shape(incoming) + and incoming != db_model.model_name + ): + return incoming if db_model.model_info and db_model.model_info.team_public_model_name: return db_model.model_info.team_public_model_name diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b296792cd09..5587d37f1f7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11875,6 +11875,9 @@ async def model_info_v2( # Update total count to include agents search_total_count = len(all_models) + # Translate `model_name` to the public name for team-scoped rows. + all_models = [_translate_model_name_for_response(m) for m in all_models] + return _paginate_models_response( all_models=all_models, page=page, @@ -12309,6 +12312,33 @@ async def model_metrics_exceptions( return {"data": response, "exception_types": list(exception_types)} +def _translate_model_name_for_response(model: dict) -> dict: + """For team-scoped DB rows, replace `model_name` with the public name + in `model_info.team_public_model_name` before returning. The DB column + and the in-memory router index keep the internal mangled name + (`model_name_{team_id}_{uuid}`) as the routing key -- this swap is a + presentation-layer concern. Returns a shallow copy; never mutates. + + Without this swap the internal name leaks into `/v1/model/info` and + `/v2/model/info`, the dashboard binds its edit form to it, and a + non-rename save round-trips the internal name back -- corrupting + `team_public_model_name` and the team ACL (see issue #28382). + """ + if not isinstance(model, dict): + return model + model_info = model.get("model_info") or {} + if not isinstance(model_info, dict): + return model + team_public = model_info.get("team_public_model_name") + team_id = model_info.get("team_id") + if not team_public or not team_id: + return model + current = model.get("model_name") or "" + if not current.startswith(f"model_name_{team_id}_"): + return model + return {**model, "model_name": team_public} + + def _get_proxy_model_info(model: dict) -> dict: # provided model_info in config.yaml model_info = model.get("model_info", {}) @@ -12349,7 +12379,7 @@ def _get_proxy_model_info(model: dict) -> dict: deployment_dict=model, excluded_keys={"litellm_credential_name"} ) - return model + return _translate_model_name_for_response(model) @router.get( @@ -12489,8 +12519,11 @@ async def model_info_v1( # noqa: PLR0915 else: all_models = [] - for in_place_model in all_models: - in_place_model = _get_proxy_model_info(model=in_place_model) + # Reassign each entry: _get_proxy_model_info returns a (possibly new) + # dict via _translate_model_name_for_response, which does NOT mutate in + # place. Binding only the loop variable would drop the public-name swap + # for team-scoped rows and leak the internal routing key (#28382). + all_models = [_get_proxy_model_info(model=model) for model in all_models] verbose_proxy_logger.debug("all_models: %s", all_models) return {"data": all_models} diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 85c7c130b36..f16074a049b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -1129,6 +1129,307 @@ class TestTeamModelUpdate: ) assert "403" in str(exc_info.value) + def test_get_public_model_name_28382_dashboard_echo_preserves_public_name(self): + """Regression for #28382 - a non-rename dashboard PATCH echoes the + internal generated model_name (model_name_{team}_{uuid}) at the top + level. That internal-shape value must be ignored (not treated as a + rename), so _get_public_model_name falls through to the existing public + name instead of overwriting it with the internal one.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _get_public_model_name, + ) + from litellm.types.router import ModelInfo + + db_model = Deployment( + model_name="model_name_test-team_abc123", + litellm_params=LiteLLM_Params(model="azure/gpt-5.2-low-rpm-testing"), + model_info=ModelInfo( + team_id="test-team", + team_public_model_name="gpt-5.2-low-rpm-testing", + ), + ) + patch_data = updateDeployment( + model_name="model_name_test-team_abc123", + model_info=ModelInfo( + team_id="test-team", + team_public_model_name="gpt-5.2-low-rpm-testing", + ), + ) + + assert ( + _get_public_model_name(patch_data=patch_data, db_model=db_model) + == "gpt-5.2-low-rpm-testing" + ) + + def test_get_public_model_name_preserves_db_public_name_when_internal_name_unchanged( + self, + ): + """If patch_data.model_info has no team_public_model_name and + patch_data.model_name equals db_model.model_name (dashboard re-sending + the internal name without touching the public-name field), the + existing db_model.model_info.team_public_model_name must be preserved.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _get_public_model_name, + ) + from litellm.types.router import ModelInfo + + db_model = Deployment( + model_name="model_name_test-team_abc123", + litellm_params=LiteLLM_Params(model="azure/gpt-5.2-low-rpm-testing"), + model_info=ModelInfo( + team_id="test-team", + team_public_model_name="gpt-5.2-low-rpm-testing", + ), + ) + patch_data = updateDeployment( + model_name="model_name_test-team_abc123", + model_info=ModelInfo(team_id="test-team"), + ) + + assert ( + _get_public_model_name(patch_data=patch_data, db_model=db_model) + == "gpt-5.2-low-rpm-testing" + ) + + def test_get_public_model_name_allows_top_level_rename(self): + """A genuine rename via the top-level model_name field (no + patch_data.model_info.team_public_model_name supplied, and the new + name differs from the existing internal db model_name) must still + return the new name.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _get_public_model_name, + ) + from litellm.types.router import ModelInfo + + db_model = Deployment( + model_name="model_name_test-team_abc123", + litellm_params=LiteLLM_Params(model="azure/gpt-5.2-low-rpm-testing"), + model_info=ModelInfo( + team_id="test-team", + team_public_model_name="old-public-name", + ), + ) + patch_data = updateDeployment( + model_name="new-public-name", + model_info=ModelInfo(team_id="test-team"), + ) + + assert ( + _get_public_model_name(patch_data=patch_data, db_model=db_model) + == "new-public-name" + ) + + def test_get_public_model_name_top_level_rename_wins_over_stale_model_info(self): + """Regression (codex review): on a dashboard rename the UI sends the new + name in model_name but passes the existing model_info blob through + untouched -- so it still carries the OLD team_public_model_name. The + top-level rename must win; otherwise _update_existing_team_model_assignment + sees no change, never updates the team ACL, and the rename is silently + dropped while the UI optimistically shows the new name.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _get_public_model_name, + ) + from litellm.types.router import ModelInfo + + db_model = Deployment( + model_name="model_name_team-a_abc123", + litellm_params=LiteLLM_Params(model="azure/gpt-4.1"), + model_info=ModelInfo( + team_id="team-a", team_public_model_name="old-public-name" + ), + ) + patch_data = updateDeployment( + model_name="new-public-name", + model_info=ModelInfo( + team_id="team-a", + team_public_model_name="old-public-name", # stale, untouched by UI + ), + ) + + assert ( + _get_public_model_name(patch_data=patch_data, db_model=db_model) + == "new-public-name" + ) + + def test_get_public_model_name_falls_back_to_db_public_name(self): + """When patch_data carries no name hints at all (neither model_name + nor model_info.team_public_model_name), fall back to the existing + db_model.model_info.team_public_model_name.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _get_public_model_name, + ) + from litellm.types.router import ModelInfo + + db_model = Deployment( + model_name="model_name_test-team_abc123", + litellm_params=LiteLLM_Params(model="azure/gpt-5.2-low-rpm-testing"), + model_info=ModelInfo( + team_id="test-team", + team_public_model_name="gpt-5.2-low-rpm-testing", + ), + ) + patch_data = updateDeployment( + model_info=ModelInfo(team_id="test-team"), + ) + + assert ( + _get_public_model_name(patch_data=patch_data, db_model=db_model) + == "gpt-5.2-low-rpm-testing" + ) + + def test_get_public_model_name_last_resort_returns_db_model_name(self): + """Legacy rows may have no team_public_model_name anywhere; the + function must still return a string (the existing db_model.model_name) + rather than raising.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _get_public_model_name, + ) + from litellm.types.router import ModelInfo + + db_model = Deployment( + model_name="legacy-model", + litellm_params=LiteLLM_Params(model="azure/legacy"), + model_info=ModelInfo(team_id="test-team"), + ) + patch_data = updateDeployment( + model_info=ModelInfo(team_id="test-team"), + ) + + assert ( + _get_public_model_name(patch_data=patch_data, db_model=db_model) + == "legacy-model" + ) + + def test_get_public_model_name_ignores_different_internal_shape_name(self): + """A stale client may PATCH an internal-shaped model_name that does not + equal the current DB column (e.g. a different uuid). It must NOT be + treated as a rename -- fall through to the existing public name.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _get_public_model_name, + ) + from litellm.types.router import ModelInfo + + db_model = Deployment( + model_name="model_name_test-team_realuuid", + litellm_params=LiteLLM_Params(model="azure/gpt-5.2-low-rpm-testing"), + model_info=ModelInfo( + team_id="test-team", + team_public_model_name="gpt-5.2-low-rpm-testing", + ), + ) + patch_data = updateDeployment( + model_name="model_name_test-team_differentuuid", + model_info=ModelInfo(team_id="test-team"), + ) + + assert ( + _get_public_model_name(patch_data=patch_data, db_model=db_model) + == "gpt-5.2-low-rpm-testing" + ) + + def test_get_public_model_name_ignores_internal_shape_patch_public(self): + """If a corrupted row round-trips an internal-shaped value in + model_info.team_public_model_name, it must not be accepted as the + public name -- fall through to the existing db public name.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _get_public_model_name, + ) + from litellm.types.router import ModelInfo + + db_model = Deployment( + model_name="model_name_test-team_realuuid", + litellm_params=LiteLLM_Params(model="azure/gpt-5.2-low-rpm-testing"), + model_info=ModelInfo( + team_id="test-team", + team_public_model_name="gpt-5.2-low-rpm-testing", + ), + ) + patch_data = updateDeployment( + model_info=ModelInfo( + team_id="test-team", + team_public_model_name="model_name_test-team_realuuid", + ), + ) + + assert ( + _get_public_model_name(patch_data=patch_data, db_model=db_model) + == "gpt-5.2-low-rpm-testing" + ) + + @pytest.mark.asyncio + async def test_dashboard_edit_preserves_public_name_and_acl(self): + """End-to-end regression for #28382: PATCH payload shaped like the + dashboard's model-edit form (top-level model_name = internal generated + name, model_info.team_public_model_name = public name) must NOT trigger + a public-name rename, must NOT touch the team ACL, and must serialize + the public name back into model_info.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _update_team_model_in_db, + ) + from litellm.types.router import ModelInfo + + db_model = Deployment( + model_name="model_name_test-team_abc123", + litellm_params=LiteLLM_Params( + model="azure/gpt-5.2-low-rpm-testing", + custom_llm_provider="azure", + ), + model_info=ModelInfo( + id="model-id-123", + team_id="test-team", + team_public_model_name="gpt-5.2-low-rpm-testing", + ), + ) + patch_data = updateDeployment( + model_name="model_name_test-team_abc123", + litellm_params=None, + model_info=ModelInfo( + id="model-id-123", + team_id="test-team", + team_public_model_name="gpt-5.2-low-rpm-testing", + ), + ) + user_api_key_dict = UserAPIKeyAuth( + user_id="test_user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + prisma_client = MockPrismaClient(team_exists=True) + + with ( + patch( + "litellm.proxy.proxy_server.premium_user", + True, + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add" + ) as mock_team_model_add, + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_delete" + ) as mock_team_model_delete, + ): + result = await _update_team_model_in_db( + db_model=db_model, + patch_data=patch_data, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, # type: ignore + ) + + # team ACL must not be touched on a no-op edit + mock_team_model_add.assert_not_called() + mock_team_model_delete.assert_not_called() + + # the merged model_info written to the DB must keep the public name + model_info_json = result.get("model_info", "") + parsed_model_info = json.loads(model_info_json) + assert ( + parsed_model_info.get("team_public_model_name") == "gpt-5.2-low-rpm-testing" + ) + + # the internal model_name must not have been overwritten (caller + # intentionally clears patch_data.model_name so the DB row's name + # column is left alone) + assert result.get("model_name") == "model_name_test-team_abc123" + class TestModelInfoEndpoint: """Test the model_info endpoint for retrieving individual model information""" diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py new file mode 100644 index 00000000000..97e5c494916 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -0,0 +1,178 @@ +"""Coverage for team-scoped model-name translation in /model/info responses. + +These live in tests/test_litellm/proxy/proxy_server/ (not the top-level +test_proxy_server.py) because the CI coverage job collects this directory. +They exercise the read-path fix for issue #28382: `/v1`, `/v2`, and +`/model/info` must surface `model_info.team_public_model_name` for team-scoped +rows instead of the internal routing key `model_name_{team_id}_{uuid}`. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import litellm.proxy.proxy_server as ps +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.proxy_server import ( + _get_proxy_model_info, + _translate_model_name_for_response, +) + + +def _team_row() -> dict: + return { + "model_name": "model_name_team-abc-123_4a6b8", + "litellm_params": {"model": "azure/gpt-5.2-low-rpm-testing"}, + "model_info": { + "id": "byok-id-1", + "team_id": "team-abc-123", + "team_public_model_name": "team-claude-sonnet", + "db_model": True, + }, + } + + +def test_translate_swaps_internal_name_for_public(): + """Team-scoped row: model_name is swapped to the public name.""" + result = _translate_model_name_for_response(_team_row()) + assert result["model_name"] == "team-claude-sonnet" + + +def test_translate_leaves_global_row_untouched(): + """No team_id / team_public_model_name -> pass through unchanged.""" + model = { + "model_name": "gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": "normal-id-1", "db_model": False}, + } + assert _translate_model_name_for_response(model)["model_name"] == "gpt-4o" + + +def test_translate_leaves_non_internal_shape_untouched(): + """Team row whose model_name is not the internal routing key is not rewritten.""" + model = _team_row() + model["model_name"] = "already-public-name" + assert ( + _translate_model_name_for_response(model)["model_name"] == "already-public-name" + ) + + +def test_translate_handles_missing_or_non_dict_model_info(): + """Missing / None / non-dict model_info, and a non-dict model, must not raise.""" + # missing model_info + assert _translate_model_name_for_response({"model_name": "x"})["model_name"] == "x" + # model_info is None -> coerced to {} -> no team fields + assert ( + _translate_model_name_for_response({"model_name": "x", "model_info": None})[ + "model_name" + ] + == "x" + ) + # model_info is a truthy non-dict (e.g. a stray string) -> early return + assert ( + _translate_model_name_for_response( + {"model_name": "x", "model_info": "garbage"} + )["model_name"] + == "x" + ) + # model itself is not a dict + assert _translate_model_name_for_response("not-a-dict") == "not-a-dict" # type: ignore[arg-type] + + +def test_translate_does_not_mutate_input(): + """Returns a shallow copy; the router's in-memory list keeps the routing key.""" + model = _team_row() + result = _translate_model_name_for_response(model) + assert result is not model + assert model["model_name"] == "model_name_team-abc-123_4a6b8" + + +def test_get_proxy_model_info_returns_public_name_for_team_row(): + """`_get_proxy_model_info` must return the public name for a team-scoped + row. Because _translate_model_name_for_response returns a shallow copy + (it does not mutate), callers MUST use the return value -- the + `/v1/model/info` list path historically discarded it, leaking the internal + routing key (#28382).""" + # Mirror the (fixed) /v1/model/info list path: assign the return back. + all_models = [_get_proxy_model_info(model=m) for m in [_team_row()]] + assert all_models[0]["model_name"] == "team-claude-sonnet" + + +@pytest.mark.asyncio +async def test_model_info_v2_translates_team_model_name(monkeypatch): + """/v2/model/info must surface the public name for team-scoped rows. + Covers the translation step in model_info_v2 (the read-path call site).""" + router = MagicMock() + router.model_list = [_team_row()] + + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(ps.proxy_config, "get_config", AsyncMock(return_value={})) + monkeypatch.setattr( + ps, + "_apply_search_filter_to_models", + AsyncMock(side_effect=lambda all_models, **kw: (all_models, len(all_models))), + ) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + import litellm.proxy.agent_endpoints.model_list_helpers as mlh + + monkeypatch.setattr( + mlh, + "append_agents_to_model_info", + AsyncMock(side_effect=lambda models, **kw: models), + ) + + admin = UserAPIKeyAuth(user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN) + # Pass every query param explicitly: called directly (not through FastAPI), + # the fastapi.Query(...) defaults are Query objects, not their values. + resp = await ps.model_info_v2( + user_api_key_dict=admin, + model=None, + user_models_only=False, + include_team_models=False, + debug=False, + page=1, + size=50, + search=None, + modelId=None, + teamId=None, + sortBy=None, + sortOrder="asc", + ) + + names = [m["model_name"] for m in resp["data"]] + assert "team-claude-sonnet" in names + assert "model_name_team-abc-123_4a6b8" not in names + + +@pytest.mark.asyncio +async def test_model_info_v1_list_path_translates_team_model_name(monkeypatch): + """/v1/model/info list path (no litellm_model_id) must surface the public + name. Covers the list comprehension that assigns _get_proxy_model_info's + return back into all_models (#28382 review).""" + router = MagicMock() + router.get_model_names.return_value = ["team-claude-sonnet"] + router.get_model_access_groups.return_value = {} + router.get_model_list.return_value = [_team_row()] + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", [_team_row()]) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "get_key_models", lambda **kw: []) + monkeypatch.setattr(ps, "get_team_models", lambda **kw: []) + monkeypatch.setattr( + ps, "get_complete_model_list", lambda **kw: ["team-claude-sonnet"] + ) + + admin = UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[] + ) + resp = await ps.model_info_v1(user_api_key_dict=admin, litellm_model_id=None) + + names = [m["model_name"] for m in resp["data"]] + assert "team-claude-sonnet" in names + assert "model_name_team-abc-123_4a6b8" not in names diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/AgentBuilderView.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/AgentBuilderView.tsx index c47c201d074..8b0de5dda7c 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/AgentBuilderView.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/AgentBuilderView.tsx @@ -1,6 +1,14 @@ "use client"; -import { CommentOutlined, DeleteOutlined, ExperimentOutlined, LinkOutlined, PlusOutlined, RobotOutlined, SaveOutlined } from "@ant-design/icons"; +import { + CommentOutlined, + DeleteOutlined, + ExperimentOutlined, + LinkOutlined, + PlusOutlined, + RobotOutlined, + SaveOutlined, +} from "@ant-design/icons"; import { Button, Input, Modal, Select, Spin, Tabs } from "antd"; import React, { useCallback, useEffect, useState } from "react"; import CodeBlock from "@/app/(dashboard)/api-reference/components/CodeBlock"; @@ -64,9 +72,10 @@ function ConnectTabContent({ onCreateKey, }: ConnectTabContentProps) { const baseUrl = proxyBaseUrl ?? getConnectTabBaseUrl(proxySettings, customProxyBaseUrl); - const apiKeyForCurl = - createdKeyValue ? - createdKeyValue.startsWith("Bearer ") ? createdKeyValue : `Bearer ${createdKeyValue}` + const apiKeyForCurl = createdKeyValue + ? createdKeyValue.startsWith("Bearer ") + ? createdKeyValue + : `Bearer ${createdKeyValue}` : "Bearer sk-1234"; const curlExample = `curl -L -X POST '${baseUrl}/v1/chat/completions' \\ -H 'x-litellm-api-key: ${apiKeyForCurl}' \\ @@ -101,12 +110,7 @@ function ConnectTabContent({ Create a virtual key that can only call this agent. The key will be scoped to you (user_id) and restricted to the model {agentName}.

- {disabledPersonalKeyCreation && ( @@ -127,6 +131,14 @@ function getAgentModelId(agent: AgentModel): string | null { return info?.id ?? null; } +// Selection key that always resolves to a non-null string. Prefers the DB +// id (stable across renames and unique across teams) but falls back to +// `model_name` so config-file-defined agents — which have no `model_info.id` +// — remain selectable. +function getAgentSelectionKey(agent: AgentModel): string { + return getAgentModelId(agent) ?? agent.model_name; +} + function parseUnderlyingModel(litellmModel: string | undefined): string | undefined { if (!litellmModel || !litellmModel.startsWith("litellm_agent/")) return undefined; return litellmModel.slice("litellm_agent/".length) || undefined; @@ -191,22 +203,25 @@ export default function AgentBuilderView({ const [deleting, setDeleting] = useState(false); const effectiveApiKey = apiKey || accessToken || ""; - const selectedAgent = selectedId === NEW_AGENT_ID ? null : agentModels.find((a) => a.model_name === selectedId) ?? null; + const selectedAgent = + selectedId === NEW_AGENT_ID ? null : agentModels.find((a) => getAgentSelectionKey(a) === selectedId) ?? null; const isNewAgent = selectedId === NEW_AGENT_ID; const selectedAgentModelId = selectedAgent ? getAgentModelId(selectedAgent) : null; - const loadAgents = useCallback(async () => { - if (!accessToken || !userID || !userRole) return; + const loadAgents = useCallback(async (): Promise => { + if (!accessToken || !userID || !userRole) return []; setLoadingAgents(true); try { const list = await fetchAvailableAgentModels(accessToken, userID, userRole); setAgentModels(list); - if (!selectedId || (selectedId !== NEW_AGENT_ID && !list.some((a) => a.model_name === selectedId))) { - setSelectedId(list.length > 0 ? list[0].model_name : null); + if (!selectedId || (selectedId !== NEW_AGENT_ID && !list.some((a) => getAgentSelectionKey(a) === selectedId))) { + setSelectedId(list.length > 0 ? getAgentSelectionKey(list[0]) : null); } + return list; } catch (e) { console.error(e); NotificationsManager.fromBackend("Failed to load agents"); + return []; } finally { setLoadingAgents(false); } @@ -267,7 +282,13 @@ export default function AgentBuilderView({ setDraftMaxTokens(typeof p?.max_tokens === "number" ? p.max_tokens : 4096); const rawTools = selectedAgent.litellm_params?.tools; const tools: MCPToolEntry[] = Array.isArray(rawTools) - ? rawTools.filter((t): t is MCPToolEntry => t && typeof t === "object" && (t as MCPToolEntry).type === "mcp" && typeof (t as MCPToolEntry).server_url === "string") + ? rawTools.filter( + (t): t is MCPToolEntry => + t && + typeof t === "object" && + (t as MCPToolEntry).type === "mcp" && + typeof (t as MCPToolEntry).server_url === "string", + ) : []; setDraftTools(tools); } @@ -297,7 +318,7 @@ export default function AgentBuilderView({ } setSaving(true); try { - await modelCreateCall(accessToken, { + const response = await modelCreateCall(accessToken, { model_name: draftName.trim(), litellm_params: { model: `litellm_agent/${draftUnderlyingModel}`, @@ -308,9 +329,15 @@ export default function AgentBuilderView({ }, model_info: {}, }); - const newName = draftName.trim(); - await loadAgents(); - setSelectedId(newName); + // /model/new returns the row with `model_id` at the top level. + // Prefer that id over name-matching so we land on the just-created + // agent even when its public name collides with another team's. + const createdId: string | null = response?.model_id ?? response?.model_info?.id ?? null; + const list = await loadAgents(); + const created = createdId + ? list.find((a) => getAgentModelId(a) === createdId) ?? list.find((a) => a.model_name === draftName.trim()) + : list.find((a) => a.model_name === draftName.trim()); + setSelectedId(created ? getAgentSelectionKey(created) : list[0] ? getAgentSelectionKey(list[0]) : null); setActiveTab("chat"); } catch (e) { NotificationsManager.fromBackend("Failed to save agent"); @@ -342,8 +369,10 @@ export default function AgentBuilderView({ selectedAgentModelId, ); NotificationsManager.success("Agent updated successfully"); - await loadAgents(); - setSelectedId(draftName.trim()); + const list = await loadAgents(); + const stillSelected = list.find((a) => getAgentModelId(a) === selectedAgentModelId); + const target = stillSelected ?? list[0]; + setSelectedId(target ? getAgentSelectionKey(target) : null); } catch (e) { NotificationsManager.fromBackend("Failed to update agent"); } finally { @@ -387,9 +416,9 @@ export default function AgentBuilderView({ try { await modelDeleteCall(accessToken, selectedAgentModelId); NotificationsManager.success("Agent deleted"); - await loadAgents(); - const remaining = agentModels.filter((a) => a.model_name !== selectedAgent.model_name); - setSelectedId(remaining.length > 0 ? remaining[0].model_name : null); + const list = await loadAgents(); + const remaining = list.filter((a) => getAgentModelId(a) !== selectedAgentModelId); + setSelectedId(remaining.length > 0 ? getAgentSelectionKey(remaining[0]) : null); } catch (e) { NotificationsManager.fromBackend("Failed to delete agent"); } finally { @@ -401,9 +430,7 @@ export default function AgentBuilderView({ if (!accessToken || !userID || !userRole) { return ( -
- Sign in to use Agent Builder. -
+
Sign in to use Agent Builder.
); } @@ -412,24 +439,25 @@ export default function AgentBuilderView({
Agent Builder - {isNewAgent ? ( - - ) : ( - Build Agents that pass your compliance requirements. - )} + {isNewAgent ? ( + + ) : ( + Build Agents that pass your compliance requirements. + )}
- Agent Builder is experimental and may change or be removed without notice. We’d love your feedback—email us at{" "} + Agent Builder is experimental and may change or be removed without notice. We’d love your feedback—email us + at{" "} product@berri.ai @@ -452,21 +480,24 @@ export default function AgentBuilderView({
) : ( <> - {agentModels.map((agent) => ( - - ))} + {agentModels.map((agent) => { + const key = getAgentSelectionKey(agent); + return ( + + ); + })}