diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 7dc5386c413..59121d821f3 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -6,12 +6,14 @@ This is an enterprise feature and requires a premium license. import asyncio import re +import time from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass from functools import partial from itertools import chain from typing import TYPE_CHECKING, Final, NamedTuple, Protocol, overload +import regex from fastapi import ( APIRouter, Body, @@ -51,6 +53,8 @@ from litellm.proxy.management_endpoints.scim.scim_transformations import ( ScimTransformations, ) from litellm.proxy.management_endpoints.team_endpoints import ( + _CacheableTeamRow, # pyright: ignore[reportPrivateUsage] # shared team-cache protocol has no public export + _refresh_cached_team, # pyright: ignore[reportPrivateUsage] # shared team-cache refresh helper has no public export fetch_and_validate_organization, new_team, team_member_add, @@ -444,12 +448,20 @@ async def _get_scim_settings() -> SCIMSettings: SCIM_ORG_MAPPING_MATCH_TIMEOUT_SECONDS: Final = 1.0 +def _fullmatch_before_deadline(pattern: str, display_name: str, deadline: float) -> bool: + remaining: Final = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("scim organization mapping evaluation deadline exceeded") + return regex.fullmatch(pattern, display_name, timeout=remaining) is not None + + def _first_matching_organization_id(display_name: str, settings: SCIMSettings) -> str | None: + deadline: Final = time.monotonic() + SCIM_ORG_MAPPING_MATCH_TIMEOUT_SECONDS return next( ( mapping.organization_id for mapping in settings.organization_mappings - if re.fullmatch(mapping.group_display_name_pattern, display_name) + if _fullmatch_before_deadline(mapping.group_display_name_pattern, display_name, deadline) ), None, ) @@ -458,17 +470,16 @@ def _first_matching_organization_id(display_name: str, settings: SCIMSettings) - async def _resolve_scim_group_organization_id(display_name: str, settings: SCIMSettings) -> str | None: """First organization mapping whose pattern fully matches the group displayName, or None. - Matching runs off the event loop with a timeout so a backtracking-heavy - pattern cannot stall SCIM and unrelated proxy requests. + Matching runs off the event loop, and the regex engine enforces a hard + per-evaluation deadline that terminates the match itself, so a + backtracking-heavy pattern cannot stall the proxy or keep burning a + worker thread after the request has failed. """ if not settings.organization_mappings: return None try: - return await asyncio.wait_for( - asyncio.to_thread(_first_matching_organization_id, display_name, settings), - timeout=SCIM_ORG_MAPPING_MATCH_TIMEOUT_SECONDS, - ) - except asyncio.TimeoutError as e: + return await asyncio.to_thread(_first_matching_organization_id, display_name, settings) + except TimeoutError as e: raise HTTPException( status_code=400, detail={ @@ -478,6 +489,32 @@ async def _resolve_scim_group_organization_id(display_name: str, settings: SCIMS ) from e +async def _refresh_scim_updated_team_cache(updated_team: _CacheableTeamRow | None) -> None: + """Keep the cached team object used by auth checks in sync after a SCIM write, + so an organization or alias change takes effect immediately instead of after TTL expiry. + + Best-effort: the DB write has already committed, so a cache backend failure + must not fail the SCIM response and trigger an IdP retry of a successful write. + """ + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + + if updated_team is None: + return + try: + await _refresh_cached_team( + team_row=updated_team, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # best-effort refresh: the team row is committed; a stale cache entry expires at TTL + verbose_proxy_logger.warning( + "Failed to refresh cached team %s after SCIM update; " + "a stale team object may be served until its TTL expires: %s", + updated_team.team_id, + e, + ) + + async def _validate_mapped_organization_exists( prisma_client: PrismaClient, organization_id: str, existing_team: LiteLLM_TeamTable ) -> None: @@ -2583,6 +2620,7 @@ async def update_group( where={"team_id": group_id}, data=update_data, ) + await _refresh_scim_updated_team_cache(updated_team) # Handle user-team relationship changes current_members: Final = set(await _get_team_member_user_ids_from_team(existing_team)) @@ -2754,10 +2792,12 @@ async def _apply_group_patch_updates(group_id: str, update_data: dict[str, objec update_data["metadata"] = safe_dumps(update_data["metadata"]) if update_data: - return await TeamRepository(prisma_client).table.update( + updated_team: Final = await TeamRepository(prisma_client).table.update( where={"team_id": group_id}, data=update_data, ) + await _refresh_scim_updated_team_cache(updated_team) + return updated_team return await TeamRepository(prisma_client).table.find_unique(where={"team_id": group_id}) diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 2f6c2aa20c7..a20657db524 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -5813,3 +5813,81 @@ async def test_patch_group_rename_assigns_mapped_organization(mocker: MockerFixt def test_scim_settings_rejects_invalid_regex_pattern(): with pytest.raises(ValueError, match="Invalid regex pattern"): SCIMGroupOrganizationMapping(group_display_name_pattern="Engineering-[", organization_id="org-eng") + + +@pytest.mark.asyncio +async def test_resolve_organization_id_times_out_on_catastrophic_pattern(): + """A backtracking-heavy pattern must terminate with an actionable 400 instead of + matching forever.""" + from litellm.proxy.management_endpoints.scim.scim_v2 import ( + _resolve_scim_group_organization_id, + ) + + settings = SCIMSettings( + organization_mappings=[ + SCIMGroupOrganizationMapping(group_display_name_pattern="(a|a)+$", organization_id="org-eng") + ] + ) + + with pytest.raises(HTTPException) as exc_info: + await _resolve_scim_group_organization_id("a" * 64 + "!", settings) + + assert exc_info.value.status_code == 400 + assert "timed out" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_update_group_refreshes_cached_team(mocker: MockerFixture): # test-quality-ok: the observable is the shared cache refresh helper receiving the committed row; the helper's cache write is covered by team_endpoints tests + """A PUT that moves a team into a mapped organization must refresh the cached + team object auth checks read, so the move takes effect immediately.""" + from litellm.proxy._types import LiteLLM_TeamTable, Member + from litellm.proxy.management_endpoints.scim.scim_transformations import ( + ScimTransformations, + ) + + group_id = "team-1" + existing_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Sales", + members=["user1"], + members_with_roles=[Member(user_id="user1", role="user")], + metadata={}, + ) + scim_group_update = SCIMGroup( + schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"], + id=group_id, + displayName="Engineering-Platform", + members=[SCIMMember(value="user1")], + ) + + mock_prisma_client = _mock_group_prisma_client(mocker, existing_team=existing_team) + mocker.patch( # test-quality-ok: stubs the collaborator to pin the org mapping the endpoint computes + "litellm.proxy.management_endpoints.scim.scim_v2._get_scim_settings", + AsyncMock(return_value=_org_mapping_settings()), + ) + mocker.patch( # test-quality-ok: stubs the collaborator to pin the org mapping the endpoint computes + "litellm.proxy.management_endpoints.scim.scim_v2._validate_mapped_organization_exists", + AsyncMock(), + ) + mocker.patch( # test-quality-ok: stubs the collaborator to pin the org mapping the endpoint computes + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + AsyncMock(), + ) + mocker.patch( # test-quality-ok: stubs the collaborator to pin the org mapping the endpoint computes + "litellm.proxy.management_endpoints.scim.scim_v2._recompute_scim_member_roles", + AsyncMock(), + ) + mocker.patch.object( # test-quality-ok: stubs the collaborator to pin the org mapping the endpoint computes + ScimTransformations, + "transform_litellm_team_to_scim_group", + AsyncMock(return_value=scim_group_update), + ) + refresh_mock = mocker.patch( # test-quality-ok: asserts the cache refresh the endpoint must trigger + "litellm.proxy.management_endpoints.scim.scim_v2._refresh_cached_team", + AsyncMock(), + ) + + await update_group(group_id=group_id, group=scim_group_update) + + updated_team = mock_prisma_client.db.litellm_teamtable.update.return_value + assert refresh_mock.await_args.kwargs["team_row"] is updated_team