mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(auth): backfill organization_alias from the team's org for JWT and key auth
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
jesus
parent
1af7a403c6
commit
0e4fc7c86d
2 changed files with 207 additions and 0 deletions
|
|
@ -54,6 +54,7 @@ from litellm.proxy.auth.auth_checks import (
|
|||
get_end_user_object,
|
||||
get_jwt_key_mapping_object,
|
||||
get_object_permission,
|
||||
get_org_object,
|
||||
get_project_object,
|
||||
get_team_object,
|
||||
get_user_object,
|
||||
|
|
@ -2344,6 +2345,27 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached
|
|||
)
|
||||
|
||||
|
||||
async def _resolve_org_alias(
|
||||
org_id: str,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> str | None:
|
||||
try:
|
||||
org_object: Final = await get_org_object(
|
||||
org_id=org_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # fail-safe: attribution must never break auth
|
||||
verbose_proxy_logger.debug("org alias attribution: org lookup failed: %s", e)
|
||||
return None
|
||||
return org_object.organization_alias if org_object is not None else None
|
||||
|
||||
|
||||
def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseException) -> bool:
|
||||
"""Whether the token's own team fields may stand in for a team that failed to
|
||||
resolve, without widening access.
|
||||
|
|
@ -2601,6 +2623,15 @@ async def _run_centralized_common_checks(
|
|||
if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None:
|
||||
user_api_key_auth_obj.org_id = team_object.organization_id
|
||||
|
||||
if user_api_key_auth_obj.org_id is not None and user_api_key_auth_obj.organization_alias is None:
|
||||
user_api_key_auth_obj.organization_alias = await _resolve_org_alias(
|
||||
org_id=user_api_key_auth_obj.org_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
# common_checks identifies admin via user_object, not the token
|
||||
# (non_proxy_admin_allowed_routes_check). JWT admin shortcut and
|
||||
# master_key tokens get admin from the token; the DB row for the
|
||||
|
|
|
|||
|
|
@ -5288,6 +5288,182 @@ async def test_centralized_common_checks_backfills_org_id_from_team(key_org_id,
|
|||
setattr(_proxy_server_mod, k, v)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_centralized_common_checks_backfills_org_alias_from_team_org():
|
||||
import litellm.proxy.proxy_server as _proxy_server_mod
|
||||
from fastapi import Request
|
||||
from starlette.datastructures import URL
|
||||
|
||||
from litellm.proxy._types import LiteLLM_OrganizationTable, LiteLLM_TeamTableCachedObj
|
||||
|
||||
token = UserAPIKeyAuth(api_key="sk-test", user_id="u", team_id="t1")
|
||||
request = Request(scope={"type": "http"})
|
||||
request._url = URL(url="/chat/completions")
|
||||
|
||||
fetched_team = LiteLLM_TeamTableCachedObj(team_id="t1", organization_id="org-1")
|
||||
fetched_org = LiteLLM_OrganizationTable(
|
||||
organization_id="org-1",
|
||||
organization_alias="acme-org",
|
||||
budget_id="budget-1",
|
||||
created_by="test",
|
||||
updated_by="test",
|
||||
)
|
||||
|
||||
attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None)
|
||||
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
|
||||
try:
|
||||
for k, v in attrs.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
with (
|
||||
patch( # test-quality-ok: centralized checks read this auth module seam
|
||||
"litellm.proxy.auth.user_api_key_auth.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=fetched_team,
|
||||
),
|
||||
patch( # test-quality-ok: centralized checks read this auth module seam
|
||||
"litellm.proxy.auth.user_api_key_auth.get_org_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=fetched_org,
|
||||
) as mock_get_org,
|
||||
patch( # test-quality-ok: centralized checks read this auth module seam
|
||||
"litellm.proxy.auth.user_api_key_auth.common_checks",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_checks,
|
||||
):
|
||||
await _run_centralized_common_checks(
|
||||
user_api_key_auth_obj=token,
|
||||
request=request,
|
||||
request_data={"model": "gpt-4o"},
|
||||
route="/chat/completions",
|
||||
)
|
||||
|
||||
mock_checks.assert_awaited_once()
|
||||
assert token.organization_alias == "acme-org"
|
||||
assert mock_checks.await_args.kwargs["valid_token"].organization_alias == "acme-org"
|
||||
mock_get_org.assert_awaited_once_with(
|
||||
org_id="org-1",
|
||||
prisma_client=None,
|
||||
user_api_key_cache=attrs["user_api_key_cache"],
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=attrs["proxy_logging_obj"],
|
||||
)
|
||||
finally:
|
||||
for k, v in originals.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"org_id,organization_alias,team_org_id,expected_alias",
|
||||
[
|
||||
("org-pinned", "pinned", "org-from-team", "pinned"),
|
||||
(None, None, None, None),
|
||||
],
|
||||
)
|
||||
async def test_centralized_common_checks_preserves_or_omits_org_alias_without_org_lookup(
|
||||
org_id, organization_alias, team_org_id, expected_alias
|
||||
):
|
||||
import litellm.proxy.proxy_server as _proxy_server_mod
|
||||
from fastapi import Request
|
||||
from starlette.datastructures import URL
|
||||
|
||||
from litellm.proxy._types import LiteLLM_TeamTableCachedObj
|
||||
|
||||
token = UserAPIKeyAuth(
|
||||
api_key="sk-test",
|
||||
user_id="u",
|
||||
team_id="t1",
|
||||
org_id=org_id,
|
||||
organization_alias=organization_alias,
|
||||
)
|
||||
request = Request(scope={"type": "http"})
|
||||
request._url = URL(url="/chat/completions")
|
||||
|
||||
fetched_team = LiteLLM_TeamTableCachedObj(team_id="t1", organization_id=team_org_id)
|
||||
|
||||
attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None)
|
||||
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
|
||||
try:
|
||||
for k, v in attrs.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
with (
|
||||
patch( # test-quality-ok: centralized checks read this auth module seam
|
||||
"litellm.proxy.auth.user_api_key_auth.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=fetched_team,
|
||||
),
|
||||
patch( # test-quality-ok: centralized checks read this auth module seam
|
||||
"litellm.proxy.auth.user_api_key_auth.get_org_object",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_get_org,
|
||||
patch( # test-quality-ok: centralized checks read this auth module seam
|
||||
"litellm.proxy.auth.user_api_key_auth.common_checks",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
):
|
||||
await _run_centralized_common_checks(
|
||||
user_api_key_auth_obj=token,
|
||||
request=request,
|
||||
request_data={"model": "gpt-4o"},
|
||||
route="/chat/completions",
|
||||
)
|
||||
|
||||
assert token.organization_alias == expected_alias
|
||||
mock_get_org.assert_not_awaited()
|
||||
finally:
|
||||
for k, v in originals.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_centralized_common_checks_org_alias_lookup_failure_is_swallowed():
|
||||
import litellm.proxy.proxy_server as _proxy_server_mod
|
||||
from fastapi import Request
|
||||
from starlette.datastructures import URL
|
||||
|
||||
from litellm.proxy._types import LiteLLM_TeamTableCachedObj
|
||||
|
||||
token = UserAPIKeyAuth(api_key="sk-test", user_id="u", team_id="t1")
|
||||
request = Request(scope={"type": "http"})
|
||||
request._url = URL(url="/chat/completions")
|
||||
|
||||
fetched_team = LiteLLM_TeamTableCachedObj(team_id="t1", organization_id="org-1")
|
||||
|
||||
attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None)
|
||||
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
|
||||
try:
|
||||
for k, v in attrs.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
with (
|
||||
patch( # test-quality-ok: centralized checks read this auth module seam
|
||||
"litellm.proxy.auth.user_api_key_auth.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=fetched_team,
|
||||
),
|
||||
patch( # test-quality-ok: centralized checks read this auth module seam
|
||||
"litellm.proxy.auth.user_api_key_auth.get_org_object",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=Exception("db down"),
|
||||
),
|
||||
patch( # test-quality-ok: centralized checks read this auth module seam
|
||||
"litellm.proxy.auth.user_api_key_auth.common_checks",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_checks,
|
||||
):
|
||||
await _run_centralized_common_checks(
|
||||
user_api_key_auth_obj=token,
|
||||
request=request,
|
||||
request_data={"model": "gpt-4o"},
|
||||
route="/chat/completions",
|
||||
)
|
||||
|
||||
mock_checks.assert_awaited_once()
|
||||
assert token.organization_alias is None
|
||||
finally:
|
||||
for k, v in originals.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cli_session_token_org_backfilled_from_team(monkeypatch):
|
||||
"""LIT-4688 root cause: CLI session tokens (from /sso/cli/poll) are minted
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue