From a73e813496596a88ee283f7e37fc21689d4a70e8 Mon Sep 17 00:00:00 2001
From: JackDanger <jdanger@launchdarkly.com>
Date: Sun, 18 Oct 2026 05:49:49 +0000
Subject: [PATCH] test(proxy): cover the byok repoint and wif interdiction
 flows around the api_base gate

Generated-by: termag
---
 .../test_r2_9_byok_flows.py                   | 113 ++++++++++++++++++
 1 file changed, 113 insertions(+)
 create mode 100644 tests/unit/proxy/management_endpoints/test_r2_9_byok_flows.py

diff --git a/tests/unit/proxy/management_endpoints/test_r2_9_byok_flows.py b/tests/unit/proxy/management_endpoints/test_r2_9_byok_flows.py
new file mode 100644
index 00000000000..ea8651c216a
--- /dev/null
+++ b/tests/unit/proxy/management_endpoints/test_r2_9_byok_flows.py
@@ -0,0 +1,113 @@
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from litellm.types.router import GenericLiteLLMParams
+
+async def test_team_admin_cannot_move_api_base_keeping_server_key():
+    import asyncio
+    from unittest.mock import AsyncMock, MagicMock
+
+    from fastapi import HTTPException
+    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+    from litellm.types.router import GenericLiteLLMParams
+
+    from litellm.proxy.management_endpoints.model_management_endpoints import (
+        ModelManagementAuthChecks,
+    )
+
+    stored = GenericLiteLLMParams(api_key="server-owned-key")
+    incoming = GenericLiteLLMParams(api_base="https://attacker.example")
+    model_params = MagicMock()
+    model_params.litellm_params = stored
+
+    with pytest.raises(Exception) as exc:
+        await ModelManagementAuthChecks._reject_non_admin_wif_write(
+            model_params=model_params,
+            incoming_params=incoming.model_dump(exclude_none=True),
+            user_api_key_dict=UserAPIKeyAuth(
+                api_key="sk-team-admin",
+                user_role=LitellmUserRoles.INTERNAL_USER,
+                user_id="team-admin-1",
+            ),
+            prisma_client=AsyncMock(),
+        )
+    assert int(getattr(exc.value, "code", 0)) == 403
+    assert "api_key stays server-owned" in str(getattr(exc.value, "message", getattr(exc.value, "detail", "")))
+
+
+@pytest.mark.asyncio
+
+async def test_team_admin_bring_your_own_key_move_allowed():
+    from unittest.mock import AsyncMock
+
+    import litellm.proxy.management_endpoints.model_management_endpoints as mme
+    from litellm.proxy.management_endpoints.model_management_endpoints import (
+        ModelManagementAuthChecks,
+    )
+    from litellm.types.router import GenericLiteLLMParams
+    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+
+    stored = GenericLiteLLMParams(api_key="server-owned-key")
+    incoming = GenericLiteLLMParams(
+        api_base="https://my-team-proxy.example", api_key="my-team-key"
+    )
+    model_params = MagicMock()
+    model_params.litellm_params = stored
+
+    with (
+        patch.object(mme, "effective_server_owned_wif_fields", AsyncMock(return_value=())),
+    ):
+        await ModelManagementAuthChecks._reject_non_admin_wif_write(
+            model_params=model_params,
+            incoming_params=incoming.model_dump(exclude_none=True),
+            user_api_key_dict=UserAPIKeyAuth(
+                api_key="sk-team",
+                user_role=LitellmUserRoles.INTERNAL_USER,
+                user_id="team-admin-1",
+            ),
+            prisma_client=AsyncMock(),
+        )
+
+
+@pytest.mark.asyncio
+
+async def test_team_admin_wif_association_still_guarded():
+    """Under WIF the original server-owned-token rule keeps firing."""
+    from unittest.mock import AsyncMock
+
+    import litellm.proxy.management_endpoints.model_management_endpoints as mme
+    from litellm.proxy.management_endpoints.model_management_endpoints import (
+        ModelManagementAuthChecks,
+    )
+    from litellm.types.router import GenericLiteLLMParams
+    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+
+    stored = GenericLiteLLMParams(api_key="server-key")
+    incoming = GenericLiteLLMParams(api_base="https://moved.example")
+    model_params = MagicMock()
+    model_params.litellm_params = stored
+
+    from litellm.types.utils import CredentialItem
+
+    with (
+        patch.object(
+            mme, "effective_server_owned_wif_fields", AsyncMock(return_value=("oidc_token_url",))
+        ),
+    ):
+        try:
+            await ModelManagementAuthChecks._reject_non_admin_wif_write(
+                model_params=model_params,
+                incoming_params=incoming.model_dump(exclude_none=True),
+                user_api_key_dict=UserAPIKeyAuth(
+                    api_key="sk-team", user_role=LitellmUserRoles.INTERNAL_USER, user_id="t"
+                ),
+                prisma_client=AsyncMock(),
+            )
+        except Exception as guard_error:
+            # On this stack both the static-key harvest gate and the WIF
+            # rule can fire for the same shape; accept either interdiction.
+            assert ("api_key stays server-owned" in str(guard_error)) or (
+                "workload identity" in str(guard_error)
+            )
+            return
+        raise AssertionError("expected a server-owned-key interdiction to fire")