fix(auto-router): preserve JEV transport across dashboard edits

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/management_endpoints/model_management_endpoints.py

@@ -22,7 +22,7 @@ from types import MappingProxyType
 from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeVar, cast, runtime_checkable
 
 from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
-from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator
+from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError, field_validator
 
 import litellm
 from litellm._logging import verbose_proxy_logger
@@ -289,7 +289,11 @@ def _strategy_router_write_violation(
     if incoming_params is None:
         return None
     config_violation: Final = validate_complexity_router_config_write(
-        complexity_router_config=incoming_params.complexity_router_config
+        complexity_router_config=(
+            _effective_complexity_router_config(incoming_params, existing_params)
+            if incoming_params.complexity_router_config is not None
+            else None
+        )
     )
     if config_violation is not None:
         return config_violation
@@ -350,11 +354,33 @@ WHERE model_id <> $1
 def _effective_complexity_router_config(
     incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None
 ) -> object:
-    """The complexity config a write leaves on the row: the incoming one when the write carries it, else the stored one."""
     incoming: Final = None if incoming_params is None else incoming_params.complexity_router_config
-    if incoming is not None or existing_params is None:
+    existing: Final = None if existing_params is None else existing_params.complexity_router_config
+    if incoming is None:
+        return existing
+    if existing is None or incoming.get("classifier_type") != "jev" or existing.get("classifier_type") != "jev":
         return incoming
-    return existing_params.complexity_router_config
+    incoming_jev: Final[object] = incoming.get("jev_classifier_config")
+    existing_jev: Final[object] = existing.get("jev_classifier_config")
+    if not isinstance(incoming_jev, Mapping) or not isinstance(existing_jev, Mapping):
+        return incoming
+    supplied: Final = TypeAdapter(dict[str, object]).validate_python(incoming_jev)
+    stored: Final = TypeAdapter(dict[str, object]).validate_python(existing_jev)
+    same_base: Final = "api_base" not in supplied or supplied["api_base"] == stored.get("api_base")
+    transport: Final = MappingProxyType(
+        {
+            key: value
+            for key, value in stored.items()
+            if key in ("api_key", "api_base") and (key != "api_key" or same_base)
+        }
+    )
+    return {  # mutable-ok: persisted JSON requires concrete nested dicts
+        **incoming,
+        "jev_classifier_config": {  # mutable-ok: json.dumps cannot serialize MappingProxyType
+            **transport,
+            **supplied,
+        },
+    }
 
 
 def _effective_model(
@@ -886,7 +912,12 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr
     if updated_patch.litellm_params:
         # Encrypt any sensitive values
         encrypted_params: Final = {
-            k: encrypt_value_helper(v) for k, v in updated_patch.litellm_params.model_dump(exclude_none=True).items()
+            k: (
+                _effective_complexity_router_config(updated_patch.litellm_params, db_model.litellm_params)
+                if k == "complexity_router_config"
+                else encrypt_value_helper(v)
+            )
+            for k, v in updated_patch.litellm_params.model_dump(exclude_none=True).items()
         }
 
         merged_litellm_params.update(encrypted_params)
@@ -2528,14 +2559,21 @@ async def update_model(
             _new_litellm_params_dict: Final = model_params.litellm_params.dict(exclude_none=True)
 
             ### ENCRYPT PARAMS ###
-            for k, v in _new_litellm_params_dict.items():
-                encrypted_value = encrypt_value_helper(value=v)
-                model_params.litellm_params[k] = encrypted_value
+            encrypted_params: Final = MappingProxyType(
+                {
+                    k: (
+                        _effective_complexity_router_config(model_params.litellm_params, deployment.litellm_params)
+                        if k == "complexity_router_config"
+                        else encrypt_value_helper(value=v)
+                    )
+                    for k, v in _new_litellm_params_dict.items()
+                }
+            )
 
             ### MERGE WITH EXISTING DATA ###
             _mp: Final[dict[str, object]] = model_params.litellm_params.dict()
             merged_dictionary: Final = {
-                key: _existing_litellm_params_dict[key] if value is None else value
+                key: _existing_litellm_params_dict[key] if value is None else encrypted_params[key]
                 for key, value in _mp.items()
                 if value is not None or _existing_litellm_params_dict.get(key) is not None
             }

tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py

@@ -17,6 +17,7 @@ from litellm.proxy._types import (
     LiteLLM_TeamTable,
     LitellmUserRoles,
     Member,
+    ProxyException,
     ReconcileOutcome,
     UserAPIKeyAuth,
 )
@@ -27,6 +28,8 @@ from litellm.proxy.management_endpoints.model_management_endpoints import (
     _raise_if_rate_limits_required_but_missing,
     clear_cache,
     delete_team_models,
+    patch_model,
+    update_model,
 )
 from litellm.proxy.utils import PrismaClient
 from litellm.router import Router
@@ -6602,6 +6605,65 @@ class TestTeamMemberAutoRouterWrites:
         assert saved_info["team_id"] == "member-team"
         assert saved_info["access_groups"] == ["retained-admin-group"]
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("endpoint", ["patch", "legacy"])
+    @pytest.mark.parametrize("change", ["save", "rotate", "move", "move-without-key", "reset", "heuristic"])
+    async def test_jev_dashboard_save_preserves_server_transport(self, endpoint: str, change: str) -> None:
+        original: Final = self._row()
+        transport: Final = {"api_key": "synthetic-original-jev-key", "api_base": "https://jev.example.com"}
+        stored_config: Final = {
+            "classifier_type": "jev",
+            "tiers": {"SIMPLE": "allowed"},
+            "jev_classifier_config": {**transport, "instructions": "Old instructions", "timeout_ms": 6100},
+        }
+        row: Final = original.model_copy(
+            update={
+                "litellm_params": {
+                    "model": "auto_router/complexity_router",
+                    "complexity_router_config": stored_config,
+                },
+            }
+        )
+        database: Final = self._database(self._team(), row)
+        overrides: Final = {
+            "save": {},
+            "rotate": {"api_key": "synthetic-replacement-jev-key"},
+            "move": {"api_base": "https://new-jev.example.com", "api_key": "synthetic-replacement-jev-key"},
+            "move-without-key": {"api_base": "https://new-jev.example.com"},
+            "reset": {"api_key": None, "api_base": None},
+            "heuristic": {},
+        }[change]
+        config: Final = {
+            "tiers": {"SIMPLE": "allowed"},
+            "classifier_type": "heuristic" if change == "heuristic" else "jev",
+            **({} if change == "heuristic" else {"jev_classifier_config": {"timeout_ms": 8100, **overrides}}),
+        }
+        request: Final = updateDeployment(
+            litellm_params=updateLiteLLMParams(complexity_router_config=config),
+            model_info=ModelInfo(id=row.model_id),
+        )
+        actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
+        with self._environment(database, row):
+            operation: Final = (
+                patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor)
+            )
+            if change == "move-without-key":
+                with pytest.raises(ProxyException, match="api_base requires"):
+                    await operation
+                database.db.litellm_proxymodeltable.update.assert_not_awaited()
+                return
+            await operation
+        written: Final = database.db.litellm_proxymodeltable.update.await_args.kwargs["data"]
+        saved: Final = json.loads(written["litellm_params"])["complexity_router_config"]
+        expected: Final = (
+            config
+            if change == "heuristic"
+            else {**config, "jev_classifier_config": {**transport, "timeout_ms": 8100, **overrides}}
+        )
+        assert saved == expected
+        assert row.litellm_params["complexity_router_config"] == stored_config
+        assert request.litellm_params.complexity_router_config == config
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize("endpoint", ["patch", "legacy"])
     @pytest.mark.parametrize("access", ["owner", "peer", "limited-key"])

ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts

@@ -48,6 +48,40 @@ const hydratedState: KeywordMatchingState = {
 };
 
 describe("buildUpdatedComplexityRouterConfig keyword matching", () => {
+  it.each([false, true])("omits masked JEV credentials from dashboard saves, edited: %s", (edited) => {
+    const stored = {
+      classifier_type: "jev" as const,
+      tiers: FORM_VALUE.tiers,
+      jev_classifier_config: {
+        model: "jev-configured",
+        timeout_ms: 6100,
+        instructions: "Existing instructions",
+        api_key: "sk-s****************cret",
+        api_base: "https://jev.example.com",
+      },
+    };
+    const hydrated = hydrateComplexityRouterConfig(stored, undefined);
+    expect(hydrated.jev_classifier_config).not.toHaveProperty("api_key");
+    expect(hydrated.jev_classifier_config).not.toHaveProperty("api_base");
+    const value = edited
+      ? {
+          ...hydrated,
+          jev_classifier_config: { model: "jev-updated", timeout_ms: 8100, instructions: "" },
+        }
+      : hydrated;
+    const saved = buildUpdatedComplexityRouterConfig(stored, value);
+    expect(saved.jev_classifier_config).toEqual({
+      ...(edited
+        ? { model: "jev-updated", timeout_ms: 8100 }
+        : { model: "jev-configured", timeout_ms: 6100, instructions: "Existing instructions" }),
+    });
+    for (const classifierType of ["llm", "heuristic"] as const) {
+      expect(
+        buildUpdatedComplexityRouterConfig(saved, transitionClassifierType(value, classifierType)),
+      ).not.toHaveProperty("jev_classifier_config");
+    }
+  });
+
   it("hydrates nullable JEV instructions without resetting the server configuration", () => {
     const stored = {
       classifier_type: "jev" as const,