fix(proxy): stop LDAP reads when sync jobs time out

docs/ldap_sync.md

@@ -54,4 +54,4 @@ Failed or incomplete group searches, referrals, duplicate group IDs and invalid
 
 The LDAP identity must have permission to read every selected group, its complete membership attributes, and the user identity attributes. LDAP can omit attributes hidden by access controls just as it omits empty attributes, so the sync cannot distinguish every permission restriction from an empty group. Verify the bind account and dry-run output before enabling writes.
 
-Runs have a timeout shorter than their shared Redis lease. A timeout cancels further work and the next interval retries from current state. Cancellation does not undo groups already committed earlier in the run. Stopping the sync (`enabled: false` or removing the block and restarting) leaves current teams and memberships in place.
+Runs have a timeout shorter than their shared Redis lease. A timeout cancels further writes and signals the LDAP reader to stop between entries and lookups. LDAP socket timeouts and server search time limits use `timeout_seconds` as an upper bound and are capped further for short sync intervals, and the next interval retries from current state. Cancellation does not undo groups already committed earlier in the run. Stopping the sync (`enabled: false` or removing the block and restarting) leaves current teams and memberships in place.

litellm/proxy/_types.py

@@ -2482,7 +2482,8 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
     """
 
     ldap_sync: LDAPSyncSettings | None = Field(
-        None, description="Periodically synchronize LDAP groups to teams and existing users' memberships. Requires Redis."
+        None,
+        description="Periodically synchronize LDAP groups to teams and existing users' memberships. Requires Redis.",
     )
 
     request_tag_mappings: Sequence[RequestTagMapping] | None = Field(

litellm/proxy/common_utils/ldap_directory.py

@@ -1,8 +1,9 @@
 import importlib
 import ssl
-from collections.abc import Mapping, Sequence
+from collections.abc import Iterable, Mapping, Sequence
 from dataclasses import dataclass
-from typing import Final, Protocol, TypeAlias, cast
+from threading import Event
+from typing import Final, Protocol, TypeAlias, runtime_checkable
 from urllib.parse import urlsplit
 
 from pydantic import BaseModel, ConfigDict, StrictStr, TypeAdapter
@@ -10,6 +11,7 @@ from pydantic import BaseModel, ConfigDict, StrictStr, TypeAdapter
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.proxy.ldap_sync import LDAPSyncSettings
 
+_ITEMS: Final = TypeAdapter(Iterable[object])
 _STRINGS: Final = TypeAdapter(tuple[StrictStr, ...])
 _GROUP_CLASSES: Final = frozenset({"group", "groupofnames", "groupofuniquenames", "posixgroup"})
 
@@ -50,6 +52,7 @@ class _PagedSearch(Protocol):
         search_scope: str,
         attributes: Sequence[str],
         paged_size: int,
+        time_limit: int,
         generator: bool,
     ) -> object: ...
 
@@ -73,9 +76,18 @@ class LDAPConnection(Protocol):
     def start_tls(self) -> bool: ...
     def bind(self) -> bool: ...
     def unbind(self) -> bool: ...
-    def search(self, *, search_base: str, search_filter: str, search_scope: str, attributes: Sequence[str]) -> bool: ...
+    def search(
+        self,
+        *,
+        search_base: str,
+        search_filter: str,
+        search_scope: str,
+        attributes: Sequence[str],
+        time_limit: int,
+    ) -> bool: ...
 
 
+@runtime_checkable
 class _LDAPLibrary(Protocol):
     def Tls(self, *, validate: ssl.VerifyMode, ca_certs_file: str | None) -> object: ...
     def Server(self, host: str, *, port: int, use_ssl: bool, tls: object, connect_timeout: int) -> object: ...
@@ -104,9 +116,9 @@ def _secret(value: str | None) -> str | None:
 
 
 def open_ldap_connection(settings: LDAPSyncSettings) -> LDAPConnection:
-    library: Final = cast(  # cast-ok: ldap3 has no types, this protocol describes its public constructors
-        _LDAPLibrary, importlib.import_module("ldap3")
-    )
+    library: Final = importlib.import_module("ldap3")
+    if not isinstance(library, _LDAPLibrary):
+        raise TypeError("ldap3 does not expose the required LDAP constructors")
     parsed: Final = urlsplit(settings.server_url)
     tls: Final = library.Tls(validate=ssl.CERT_REQUIRED, ca_certs_file=settings.ca_cert_path)
     server: Final = library.Server(
@@ -152,7 +164,25 @@ def _escape_filter(value: str) -> str:
     )
 
 
-def _resolve_member(connection: LDAPConnection, settings: LDAPSyncSettings, value: str) -> str | DirectoryError:
+def _check_cancelled(stop: Event) -> None:
+    if stop.is_set():
+        raise InterruptedError("LDAP sync cancelled")
+
+
+def _checked_items(raw: object, stop: Event) -> Iterable[object]:
+    items: Final = iter(_ITEMS.validate_python(raw))
+    while True:
+        _check_cancelled(stop)
+        try:
+            yield next(items)
+        except StopIteration:
+            return
+
+
+def _resolve_member(
+    connection: LDAPConnection, settings: LDAPSyncSettings, value: str, stop: Event
+) -> str | DirectoryError:
+    _check_cancelled(stop)
     usernames: Final = settings.member_attribute.lower() == "memberuid"
     _ = connection.search(
         search_base=settings.user_search_base if usernames else value,
@@ -161,6 +191,7 @@ def _resolve_member(connection: LDAPConnection, settings: LDAPSyncSettings, valu
         else "(objectClass=*)",
         search_scope="SUBTREE" if usernames else "BASE",
         attributes=(settings.user_match_attribute, "objectClass"),
+        time_limit=settings.timeout_seconds,
     )
     if connection.result.get("result") != 0:
         return DirectoryError("LDAP member lookup did not complete successfully")
@@ -177,8 +208,9 @@ def _resolve_member(connection: LDAPConnection, settings: LDAPSyncSettings, valu
 
 
 def _read_group(
-    connection: LDAPConnection, settings: LDAPSyncSettings, entry: _Entry
+    connection: LDAPConnection, settings: LDAPSyncSettings, entry: _Entry, stop: Event
 ) -> DirectoryGroup | DirectoryError:
+    _check_cancelled(stop)
     team_id: Final = _one(entry, settings.group_id_attribute)
     name: Final = _one(entry, settings.group_name_attribute)
     if team_id is None or name is None:
@@ -186,7 +218,7 @@ def _read_group(
     if any(";range=" in key.lower() for key in entry.attributes):
         return DirectoryError("LDAP ranged membership responses are not supported", entry.dn)
     identities: Final = tuple(
-        _resolve_member(connection, settings, value) for value in _values(entry, settings.member_attribute)
+        _resolve_member(connection, settings, value, stop) for value in _values(entry, settings.member_attribute)
     )
     error: Final = next((identity for identity in identities if isinstance(identity, DirectoryError)), None)
     if error is not None:
@@ -194,33 +226,45 @@ def _read_group(
     return DirectoryGroup(team_id, name, tuple(identity for identity in identities if isinstance(identity, str)))
 
 
-def read_directory(settings: LDAPSyncSettings, connection: LDAPConnection | None = None) -> DirectoryResult:
+def read_directory(
+    settings: LDAPSyncSettings, connection: LDAPConnection | None = None, stop: Event | None = None
+) -> DirectoryResult:
+    cancellation: Final = stop if stop is not None else Event()
     try:
+        _check_cancelled(cancellation)
         client: Final = connection if connection is not None else open_ldap_connection(settings)
         try:
             _ = client.open()
+            _check_cancelled(cancellation)
             if settings.start_tls and urlsplit(settings.server_url).scheme == "ldap" and not client.start_tls():
                 return DirectoryError("LDAP TLS negotiation failed")
+            _check_cancelled(cancellation)
             if not client.bind():
                 return DirectoryError("LDAP bind failed")
+            _check_cancelled(cancellation)
             raw: Final = client.extend.standard.paged_search(
                 search_base=settings.group_search_base,
                 search_filter=settings.group_filter,
                 search_scope="SUBTREE",
                 attributes=(settings.group_id_attribute, settings.group_name_attribute, settings.member_attribute),
                 paged_size=500,
-                generator=False,
+                time_limit=settings.timeout_seconds,
+                generator=True,
             )
+            items: Final = tuple(_checked_items(raw, cancellation))
+            _check_cancelled(cancellation)
             if client.result.get("result") != 0:
                 return DirectoryError("LDAP group search did not complete successfully")
-            entries: Final = _entries(raw)
+            entries: Final = _entries(items)
             if isinstance(entries, DirectoryError):
                 return entries
             group_ids: Final = tuple(_one(entry, settings.group_id_attribute) for entry in entries)
             if len(group_ids) != len(frozenset(group_ids)):
                 return DirectoryError("LDAP group IDs are not unique within the selected groups")
-            return tuple(_read_group(client, settings, entry) for entry in entries)
+            groups: Final = tuple(_read_group(client, settings, entry, cancellation) for entry in entries)
+            _check_cancelled(cancellation)
+            return groups
         finally:
             _ = client.unbind()
-    except Exception as exc:
+    except Exception as exc:  # noqa: BLE001  # SDK errors must abort safely without exposing bind credentials
         return DirectoryError(f"LDAP read failed ({type(exc).__name__})")

litellm/proxy/common_utils/ldap_sync.py

@@ -1,10 +1,10 @@
 import asyncio
 import math
-from collections.abc import Mapping
+from collections.abc import Awaitable, Callable, Mapping
 from datetime import datetime, timedelta, timezone
-from typing import Final
+from threading import Event
+from typing import Final, Protocol
 
-from apscheduler.schedulers.asyncio import AsyncIOScheduler
 from fastapi import Request
 from pydantic import TypeAdapter
 
@@ -30,6 +30,21 @@ LDAP_SYNC_JOB_ID: Final = "ldap_team_sync"
 _SETTINGS: Final = TypeAdapter(Mapping[str, object])
 
 
+class LDAPScheduler(Protocol):
+    def add_job(
+        self,
+        func: Callable[[], Awaitable[None]],
+        trigger: str,
+        *,
+        minutes: float,
+        next_run_time: datetime,
+        id: str,
+        replace_existing: bool,
+        max_instances: int,
+        coalesce: bool,
+    ) -> object: ...
+
+
 async def resolve_user_ids(
     users: UserRepository, group: DirectoryGroup, settings: LDAPSyncSettings
 ) -> frozenset[str] | DirectoryError:
@@ -38,9 +53,18 @@ async def resolve_user_ids(
         if settings.litellm_match_field == "user_email"
         else await users.find_many_in(settings.litellm_match_field, group.identities)
     )
-    identities: Final = tuple(getattr(row, settings.litellm_match_field) for row in rows)
+    identities: Final = tuple(
+        row.user_email
+        if settings.litellm_match_field == "user_email"
+        else row.user_id
+        if settings.litellm_match_field == "user_id"
+        else row.sso_user_id
+        for row in rows
+    )
     normalized: Final = tuple(
-        identity.casefold() if settings.litellm_match_field == "user_email" else identity for identity in identities
+        identity.casefold() if settings.litellm_match_field == "user_email" else identity
+        for identity in identities
+        if identity is not None
     )
     if len(normalized) != len(frozenset(normalized)):
         return DirectoryError("Multiple LiteLLM users match one LDAP identity")
@@ -67,6 +91,7 @@ class LDAPSync:
         proxy_logging_obj: ProxyLogging,
     ) -> None:
         self.settings = settings
+        self.lease_seconds = math.ceil(settings.interval_minutes * 60) - 1
         self.prisma_client = prisma_client
         self.pod_lock_manager = pod_lock_manager
         self.user_api_key_cache = user_api_key_cache
@@ -81,7 +106,7 @@ class LDAPSync:
             return
         team: Final = await self.teams.find_by_id(group.team_id, id_field="team_id")
         if self.settings.dry_run:
-            existing: Final = (
+            existing: Final[frozenset[str]] = (
                 frozenset(member.user_id for member in team.members_with_roles if member.user_id is not None)
                 if team is not None
                 else frozenset()
@@ -126,7 +151,14 @@ class LDAPSync:
         )
 
     async def _sync(self) -> None:
-        snapshot: Final = await asyncio.to_thread(read_directory, self.settings)
+        cancellation: Final = Event()
+        directory_settings: Final = self.settings.model_copy(
+            update={"timeout_seconds": min(self.settings.timeout_seconds, max(1, self.lease_seconds // 4))}
+        )
+        try:
+            snapshot: Final = await asyncio.to_thread(read_directory, directory_settings, stop=cancellation)
+        finally:
+            cancellation.set()
         if isinstance(snapshot, DirectoryError):
             verbose_proxy_logger.warning("LDAP sync aborted: %s", snapshot.reason)
             return
@@ -136,30 +168,29 @@ class LDAPSync:
                 continue
             try:
                 await self.sync_group(group)
-            except Exception as exc:
+            except Exception as exc:  # noqa: BLE001  # isolate failed groups so later groups can still synchronize
                 verbose_proxy_logger.warning(
                     "LDAP sync failed team_id=%s (%s); retrying next interval", group.team_id, type(exc).__name__
                 )
 
     async def run(self) -> None:
-        lease_seconds: Final = math.ceil(self.settings.interval_minutes * 60) - 1
         if (
             await self.pod_lock_manager.acquire_lock(
                 cronjob_id=LDAP_SYNC_JOB_ID,
-                ttl=lease_seconds,
+                ttl=self.lease_seconds,
                 allow_reentrant=False,
             )
             is not True
         ):
             return
         try:
-            await asyncio.wait_for(self._sync(), timeout=min(600, lease_seconds / 2))
+            await asyncio.wait_for(self._sync(), timeout=min(600, self.lease_seconds / 2))
         except asyncio.TimeoutError:
             verbose_proxy_logger.warning("LDAP sync timed out; retrying next interval")
 
 
 def initialize_ldap_sync(
-    scheduler: AsyncIOScheduler,
+    scheduler: LDAPScheduler,
     general_settings: Mapping[str, object],
     prisma_client: PrismaClient | None,
     pod_lock_manager: PodLockManager,

litellm/proxy/management_helpers/team_roster_sync.py

@@ -43,6 +43,7 @@ from litellm.proxy.management_helpers.access_group_team_sync import (
 from litellm.proxy.utils import PrismaClient, ProxyLogging
 from litellm.repositories.chunked_in import delete_many_in, find_many_in, update_many_in
 from litellm.repositories.team_repository import TeamRepository
+from litellm.types.proxy.management_endpoints.ui_sso import DefaultTeamSSOParams
 
 if TYPE_CHECKING:
     from prisma import Prisma
@@ -52,6 +53,7 @@ if TYPE_CHECKING:
 
 ROSTER_TX_TIMEOUT: Final = timedelta(seconds=60)
 _JSON_OBJECT: Final = TypeAdapter(dict[str, object])
+_TEAM_DEFAULTS: Final = TypeAdapter(DefaultTeamSSOParams | dict[str, object] | None)
 
 
 def new_team_request_with_defaults(
@@ -60,11 +62,11 @@ def new_team_request_with_defaults(
     members_with_roles: Sequence[Member],
     metadata: Mapping[str, object] | None = None,
 ) -> NewTeamRequest:
-    default_params: Final = litellm.default_team_params
-    defaults: Final[Mapping[str, object]] = (
-        deepcopy(default_params)
-        if isinstance(default_params, dict)
-        else default_params.model_dump(exclude_none=True)
+    default_params: Final = _TEAM_DEFAULTS.validate_python(litellm.default_team_params)
+    defaults: Final = _JSON_OBJECT.validate_python(
+        default_params.model_dump(exclude_none=True)
+        if isinstance(default_params, DefaultTeamSSOParams)
+        else deepcopy(default_params)
         if default_params is not None
         else {}
     )
@@ -75,7 +77,10 @@ def new_team_request_with_defaults(
             "team_id": team_id,
             "team_alias": team_alias,
             "members_with_roles": members_with_roles,
-            "metadata": {**(default_metadata if isinstance(default_metadata, dict) else {}), **(metadata or {})},
+            "metadata": {
+                **(_JSON_OBJECT.validate_python(default_metadata) if isinstance(default_metadata, dict) else {}),
+                **(metadata or {}),
+            },
         }
     )
 

tests/unit/proxy/common_utils/test_ldap_directory.py

@@ -1,10 +1,17 @@
-from collections.abc import Mapping, Sequence
+import ssl
+from collections.abc import Iterable, Mapping, Sequence
+from threading import Event
 from typing import Final
 
 import pytest
 from pydantic import ValidationError
 
-from litellm.proxy.common_utils.ldap_directory import DirectoryError, DirectoryGroup, read_directory
+from litellm.proxy.common_utils.ldap_directory import (
+    DirectoryError,
+    DirectoryGroup,
+    open_ldap_connection,
+    read_directory,
+)
 from litellm.types.proxy.ldap_sync import LDAPSyncSettings
 
 SETTINGS: Final = LDAPSyncSettings(
@@ -33,10 +40,11 @@ def user(mail: str, **attributes: object) -> Mapping[str, object]:
 class Connection:
     def __init__(
         self,
-        groups: Sequence[Mapping[str, object]],
+        groups: Iterable[Mapping[str, object]],
         users: Mapping[str, Sequence[Mapping[str, object]]] | None = None,
         result_code: int = 0,
         tls_ok: bool = True,
+        stop_after_search: Event | None = None,
     ) -> None:
         self.groups = groups
         self.users = users or {}
@@ -48,6 +56,7 @@ class Connection:
         self.searches: list[tuple[str, str, str]] = []
         self.closed = False
         self.tls_ok = tls_ok
+        self.stop_after_search = stop_after_search
 
     def open(self) -> None:
         pass
@@ -70,14 +79,25 @@ class Connection:
         search_scope: str,
         attributes: Sequence[str],
         paged_size: int,
+        time_limit: int,
         generator: bool,
     ) -> object:
         return self.groups
 
-    def search(self, *, search_base: str, search_filter: str, search_scope: str, attributes: Sequence[str]) -> bool:
+    def search(
+        self,
+        *,
+        search_base: str,
+        search_filter: str,
+        search_scope: str,
+        attributes: Sequence[str],
+        time_limit: int,
+    ) -> bool:
         self.searches.append((search_base, search_filter, search_scope))
         self.response = self.users.get(search_base if search_scope == "BASE" else search_filter, ())
         self.result = {"result": self.result_code}
+        if self.stop_after_search is not None:
+            self.stop_after_search.set()
         return self.result_code == 0
 
 
@@ -187,3 +207,54 @@ def test_disabled_configuration_needs_no_connection_details():
 def test_invalid_configuration_is_rejected(updates: Mapping[str, object]):
     with pytest.raises(ValidationError):
         LDAPSyncSettings.model_validate({**SETTINGS.model_dump(), **updates})
+
+
+def test_cancelled_reader_stops_before_opening_ldap():
+    cancellation = Event()
+    cancellation.set()
+    connection = Connection([group(member=[])])
+    result = read_directory(SETTINGS, connection, cancellation)
+    assert isinstance(result, DirectoryError)
+    assert not connection.closed
+    assert connection.searches == []
+
+
+def test_cancellation_after_member_read_stops_remaining_lookups_and_closes_connection():
+    cancellation = Event()
+    connection = Connection(
+        [group(member=[ALICE, BOB])],
+        {ALICE: [user("alice@example.org")], BOB: [user("bob@example.org")]},
+        stop_after_search=cancellation,
+    )
+    result = read_directory(SETTINGS, connection, cancellation)
+    assert isinstance(result, DirectoryError)
+    assert len(connection.searches) == 1
+    assert connection.closed
+
+
+def test_cancellation_during_paged_snapshot_discards_results_before_membership_reads():
+    cancellation = Event()
+
+    def pages() -> Iterable[Mapping[str, object]]:
+        cancellation.set()
+        yield group(member=[ALICE])
+        pytest.fail("LDAP snapshot should stop consuming pages after cancellation")
+
+    connection = Connection(pages())
+    result = read_directory(SETTINGS, connection, cancellation)
+    assert isinstance(result, DirectoryError)
+    assert connection.searches == []
+    assert connection.closed
+
+
+@pytest.mark.parametrize(
+    "url, ssl_enabled, port",
+    [("ldaps://directory.example.org", True, 636), ("ldap://directory.example.org:1389", False, 1389)],
+)
+def test_real_ldap_client_constructor_verifies_tls_and_disables_referrals(url: str, ssl_enabled: bool, port: int):
+    pytest.importorskip("ldap3")
+    client = open_ldap_connection(SETTINGS.model_copy(update={"server_url": url}))
+    assert client.server.tls.validate == ssl.CERT_REQUIRED
+    assert client.server.ssl is ssl_enabled
+    assert client.server.port == port
+    assert client.auto_referrals is False