scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-06 07:48:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 112

fix(proxy): refuse a stale gateway team pick and require an explicit choice

Browse source
This commit is contained in:
mateo-berri 2026-10-28 22:12:15 -07:00
parent 9726644b8c
commit 23df4f3b84
5 changed files with 64 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/proxy/anthropic_endpoints/gateway_endpoints.py
View file

@ -261,12 +261,6 @@ async def device_authorization(request: Request) -> JSONResponse:
return JSONResponse(content=body.model_dump(exclude_none=True))
def _session_team_id(chosen_team_id: str | None, teams: tuple[str, ...]) -> str | None:
if chosen_team_id in teams:
return chosen_team_id
return teams[0] if teams else None
def _validate_login(flow: Mapping[str, object]) -> _GatewayLogin | _OAuthError:
from litellm.proxy.management_endpoints.ui_sso import selected_cli_sso_team_detail
@ -278,8 +272,15 @@ def _validate_login(flow: Mapping[str, object]) -> _GatewayLogin | _OAuthError:
status_code=400, error="invalid_grant", description="The login session is malformed; sign in again"
)
chosen_team_id: Final = flow.get(CLI_SSO_CHOSEN_TEAM_ID)
team_id: Final = _session_team_id(chosen_team_id if isinstance(chosen_team_id, str) else None, session_data.teams)
chosen: Final = flow.get(CLI_SSO_CHOSEN_TEAM_ID)
chosen_team_id: Final = chosen if isinstance(chosen, str) else None
if chosen_team_id is not None and chosen_team_id not in session_data.teams:
return _OAuthError(
status_code=400,
error="invalid_grant",
description="The team picked at sign-in is no longer one of your teams; sign in again",
)
team_id: Final = chosen_team_id or next(iter(session_data.teams), None)
selected_team: Final = selected_cli_sso_team_detail(team_details=session_data.team_details, team_id=team_id)
if selected_team is None:
return _OAuthError(

7
litellm/proxy/common_utils/html_forms/native_client_consent.py
View file

@ -74,7 +74,7 @@ button {{ flex: 1; padding: 10px; border-radius: 6px; font-size: 15px; cursor: p
"""
def team_select_field(teams: Sequence[tuple[str, str]]) -> str:
def team_select_field(teams: Sequence[tuple[str, str]], *, require_choice: bool = False) -> str:
if not teams:
return ""
if len(teams) == 1:
@ -83,9 +83,12 @@ def team_select_field(teams: Sequence[tuple[str, str]]) -> str:
f'<input type="hidden" name="team_id" value="{escape(team_id)}">'
f"<p>Requests are attributed to team <strong>{escape(team_label)}</strong>.</p>"
)
placeholder: Final = '<option value="" disabled selected>Choose a team</option>' if require_choice else ""
options: Final = "".join(
f'<option value="{escape(team_id)}">{escape(team_label)}</option>' for team_id, team_label in teams
)
required: Final = " required" if require_choice else ""
return (
f'<label for="team_id">Attribute requests to team</label><select id="team_id" name="team_id">{options}</select>'
'<label for="team_id">Attribute requests to team</label>'
f'<select id="team_id" name="team_id"{required}>{placeholder}{options}</select>'
)

7
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -16,6 +16,7 @@ import json
import os
import re
import secrets
from collections import Counter
from collections.abc import Awaitable, Callable, Mapping, Sequence
from copy import deepcopy
from html import escape
@ -748,7 +749,7 @@ def _render_cli_sso_verification_page(
<p>{instructions}</p>
<form method="post" action="{escaped_verify_url}">
<input type="hidden" name="browser_complete_token" value="{escaped_browser_complete_token}" />
{team_select_field(team_choices)}
{team_select_field(team_choices, require_choice=True)}
<label for="user_code">Verification code</label>
<input id="user_code" name="user_code" autocomplete="one-time-code"{user_code_value_attr} required autofocus />
<button type="submit">Continue</button>
@ -2456,9 +2457,11 @@ def _browser_team_choices(
) -> tuple[tuple[str, str], ...]:
if not flow.browser_team_choice:
return ()
return tuple(
labeled: Final = tuple(
(detail.team_id, detail.team_alias or detail.team_id) for detail in team_details if detail.team_id is not None
)
label_counts: Final = Counter(label for _, label in labeled)
return tuple((team_id, label if label_counts[label] == 1 else f"{label} ({team_id})") for team_id, label in labeled)
def selected_cli_sso_team_detail(team_details: object, team_id: str | None) -> CliSsoTeamDetail | None:

56
tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py
View file

@ -22,6 +22,7 @@ from fastapi.testclient import TestClient
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import ProxyException
from litellm.proxy.anthropic_endpoints import gateway_endpoints
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
from litellm.proxy.management_endpoints.ui_sso import (
_get_cli_sso_flow_cache_key,
_hash_cli_sso_secret,
@ -272,6 +273,33 @@ def _login_id(device_code: str) -> str:
return device_code.partition(".")[0]
def _request_token_after_browser_choice(
client: TestClient, cache: DualCache, chosen_team_id: str | None
) -> httpx.Response:
device_code: Final = _start_device_flow(client)
_complete_flow(
cache,
device_code,
session_data={
**_COMPLETED_SESSION,
"teams": ["team-a", "team-b"],
"team_details": [
*_COMPLETED_SESSION["team_details"],
{
"team_id": "team-b",
"team_alias": "Team B",
"team_models": ["claude-opus-4-5"],
"team_model_aliases": None,
},
],
},
)
flow_key: Final = _get_cli_sso_flow_cache_key(_login_id(device_code))
flow: Final = cache.get_cache(key=flow_key)
cache.set_cache(key=flow_key, value={**flow, "chosen_team_id": chosen_team_id}, ttl=600)
return _request_token(client, device_code)
def _complete_flow(cache: DualCache, device_code: str, session_data: Mapping[str, object] = _COMPLETED_SESSION) -> None:
key: Final = _get_cli_sso_flow_cache_key(_login_id(device_code))
flow: Final = cache.get_cache(key=key)
@ -405,28 +433,22 @@ def test_token_success_mints_bearer_and_is_single_use():
def test_token_mints_for_the_team_chosen_in_the_browser(
chosen_team_id: str | None, minted_team_id: str, minted_alias: str
):
two_team_session: Final = {
**_COMPLETED_SESSION,
"teams": ["team-a", "team-b"],
"team_details": [
*_COMPLETED_SESSION["team_details"],
{"team_id": "team-b", "team_alias": "Team B", "team_models": ["claude-opus-4-5"], "team_model_aliases": None},
],
}
with _gateway_env() as (client, cache):
device_code: Final = _start_device_flow(client)
_complete_flow(cache, device_code, session_data=two_team_session)
flow_key: Final = _get_cli_sso_flow_cache_key(_login_id(device_code))
flow: Final = cache.get_cache(key=flow_key)
cache.set_cache(key=flow_key, value={**flow, "chosen_team_id": chosen_team_id}, ttl=600)
with patch(_MINT, return_value="sk-litellm-session-token") as mint:
resp: Final = _request_token(client, device_code)
resp: Final = _request_token_after_browser_choice(client, cache, chosen_team_id)
session: Final = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(resp.json()["access_token"])
assert resp.status_code == 200
assert mint.call_args.kwargs["team_id"] == minted_team_id
assert mint.call_args.kwargs["team_alias"] == minted_alias
assert session is not None
assert (session.team_id, session.team_alias) == (minted_team_id, minted_alias)
assert _opened_refresh(resp.json()["refresh_token"]).team_id == minted_team_id
def test_token_refuses_a_browser_choice_that_is_no_longer_one_of_the_session_teams():
with _gateway_env() as (client, cache):
resp: Final = _request_token_after_browser_choice(client, cache, "team-gone")
assert resp.status_code == 400
assert resp.json()["error"] == "invalid_grant"
def test_token_teamless_user_mints_without_a_team():
with _gateway_env() as (client, cache):
device_code = _start_device_flow(client)

10
tests/unit/proxy/management_endpoints/test_ui_sso.py
View file

@ -4022,14 +4022,14 @@ class TestCLIKeyRegenerationFlow:
mock_prisma.db.litellm_teamtable.find_many = AsyncMock(
return_value=[
MagicMock(model_dump=lambda team_id=team_id, alias=alias: {"team_id": team_id, "team_alias": alias})
for team_id, alias in (("team-ops", "Ops <eu>"), ("team-claude", None))
for team_id, alias in (("team-ops", "Ops <eu>"), ("team-claude", None), ("team-ops-us", "Ops <eu>"))
]
)
with (
patch(
"litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db",
return_value=LiteLLM_UserTable(
user_id="dev-1", user_role="internal_user", teams=["team-claude", "team-ops"]
user_id="dev-1", user_role="internal_user", teams=["team-claude", "team-ops", "team-ops-us"]
),
),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
@ -4047,9 +4047,11 @@ class TestCLIKeyRegenerationFlow:
assert 'name="team_id"' not in page
return
assert (
'<select id="team_id" name="team_id">'
'<option value="team-ops">Ops &lt;eu&gt;</option>'
'<select id="team_id" name="team_id" required>'
'<option value="" disabled selected>Choose a team</option>'
'<option value="team-ops">Ops &lt;eu&gt; (team-ops)</option>'
'<option value="team-claude">team-claude</option>'
'<option value="team-ops-us">Ops &lt;eu&gt; (team-ops-us)</option>'
"</select>"
) in page