fix(sso): concurrent first sign-ins land on the winner's account, worker-kill test stops waiting on the reap (#46855)

* fix(sso): concurrent first sign-ins land on the account the winner created, worker-kill test pins half its streams to each worker

* test(integration): a worker kill returns once the worker is dead, the reap waits at block exit

uvicorn can take a full worker healthcheck to reap a worker that died mid-ping, so kill() held
the caller that long and peers held on the surviving worker hit their client timeout

* fix(sso): a sign-in that loses the first-insert race joins the default teams

The losing sign-in could read the winner's row before new_user added the default teams, so its
CLI key was minted with no team. The race tests now run the real new_user against a fake user table
This commit is contained in:
Mateo Wang 2026-10-16 21:58:05 +00:00 • committed by GitHub
parent d76ec3a47e
commit 29cd11bec6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 195 additions and 84 deletions

View file

@ -425,7 +425,7 @@ async def add_new_user_to_default_team(
user_id: str,
user_email: str | None,
user_api_key_dict: UserAPIKeyAuth,
teams: list[str] | list[NewUserRequestTeam],
teams: Sequence[str | NewUserRequestTeam],
prisma_client: "PrismaClient",
):
tasks: Final[list[Awaitable[object]]] = []

View file

@ -89,6 +89,7 @@ from litellm.proxy._types import (
Member,
NewTeamRequest,
NewUserRequest,
NewUserRequestTeam,
NewUserResponse,
ProxyErrorTypes,
ProxyException,
@ -118,7 +119,11 @@ from litellm.proxy.common_utils.html_forms.jwt_display_template import (
from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form
from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
from litellm.proxy.management_endpoints.internal_user_endpoints import (
add_new_user_to_default_team,
check_if_default_team_set,
new_user,
)
from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO, LiteLLMGoogleSSO, create_generic_sso_provider
from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import (
id_jag_assertion_capture_gap,
@ -2825,6 +2830,52 @@ async def insert_sso_user(
return response
DefaultTeamJoin: TypeAlias = Callable[
[LiteLLM_UserTable, Sequence[str | NewUserRequestTeam], PrismaClient], Awaitable[None]
]
async def _join_default_teams_as_admin(
user: LiteLLM_UserTable, teams: Sequence[str | NewUserRequestTeam], prisma_client: PrismaClient
) -> None:
await add_new_user_to_default_team(
user_id=user.user_id,
user_email=user.user_email,
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
teams=teams,
prisma_client=prisma_client,
)
def _default_teams_not_joined(user: LiteLLM_UserTable) -> tuple[str | NewUserRequestTeam, ...]:
joined: Final = frozenset(user.teams)
return tuple(
team
for team in check_if_default_team_set() or ()
if (team if isinstance(team, str) else team.team_id) not in joined
)
async def _insert_sso_user_or_find_concurrent_insert(
insert: Awaitable[NewUserResponse],
user_id: str | None,
prisma_client: PrismaClient,
join_default_teams: DefaultTeamJoin = _join_default_teams_as_admin,
) -> NewUserResponse | LiteLLM_UserTable:
try:
return await insert
except Exception:
users: Final = UserRepository(prisma_client, use_writer=True)
concurrent_insert: Final = None if user_id is None else await users.find_by_id(user_id)
if concurrent_insert is None:
raise
unjoined: Final = _default_teams_not_joined(concurrent_insert)
if not unjoined:
return concurrent_insert
await join_default_teams(concurrent_insert, unjoined, prisma_client)
return await users.find_by_id(concurrent_insert.user_id) or concurrent_insert
@router.get(
"/sso/get/ui_settings",
tags=["experimental"],
@ -3471,9 +3522,10 @@ class SSOAuthenticationHandler:
else:
verbose_proxy_logger.info("user not in DB, inserting user into LiteLLM DB")
# user not in DB, insert User into LiteLLM DB
user_info = await insert_sso_user(
result_openid=result,
user_defined_values=user_defined_values,
user_info = await _insert_sso_user_or_find_concurrent_insert(
insert=insert_sso_user(result_openid=result, user_defined_values=user_defined_values),
user_id=None if user_defined_values is None else user_defined_values.get("user_id"),
prisma_client=prisma_client,
)
return user_info
except Exception as e:

View file

@ -26,7 +26,7 @@ Reuse the existing canned provider handlers through `_support/upstream.py`. It r
Owned proxies disable `.env` loading by default. A contract testing dotenv behavior can explicitly set `PYTHON_DOTENV_DISABLED=0` in its owned environment. The runner drops inherited `*_API_KEY` variables, so a contract that needs one sets a synthetic value in its owned environment. An explicit scratch writer database drops an inherited reader with a different database name; pass both database URLs when the contract needs a paired reader. Database relays close listeners, active connections and their event loops on teardown. Upstream startup failures reap the child process before returning the failure. Custom subprocesses, direct SQL and background tasks still need an explicit owner in the test
Owned proxy workers configure a 32-thread executor before application startup, preserving additional worker startup hooks. Held synchronous provider calls must not consume a CPU-dependent executor limit before the rest of a fault burst can start. The proxy entrypoint imports LiteLLM inside `main`, so spawned workers begin supervision before importing the application. Worker healthchecks have a five-second deadline by default, so uvicorn kills a worker that stays paused longer than that. Tests never signal proxy processes themselves. `OwnedProxy.kill_worker` SIGKILLs a worker and returns once its replacement answers readiness with its own pid, `replacing_worker` does the same for a kill that has to land mid-burst, and `suspended` pauses a worker or the supervisor, which only works on a proxy launched with `healthcheck=WorkerHealthcheck.LENIENT`. Helper processes that are not the proxy stop through `stop_child`, `stop_processes` and `stop_process_group` and pause through `paused`. `scripts/check_test_quality.py` rule TQ010 rejects a raw `os.kill`, SIGKILL / SIGSTOP / SIGCONT, `.kill()`, `.terminate()`, `.suspend()` or `.send_signal()` anywhere under `tests/integration` outside `_support`, and TQ011 rejects an assertion that caps an elapsed time with a fixed number. Bound timings with `_support/timing.py` instead, relative to a measured control or to the injected delay. `run.py` measures how long this box takes to import the proxy application before it starts pytest, records it in `boot-calibration.json` beside the results and exports it as `INTEGRATION_IMPORT_SECONDS`. Readiness deadlines, graceful stops and the default per-test timeout are multiples of that measurement (`_support/budget.py`), so a slower box gets longer budgets and a two-worker proxy gets twice the readiness deadline of a one-worker one. Readiness requires responsive worker identities that are still alive and a positively identified owned listener. Unreadable socket records cannot establish ownership. Crash tests own their client tasks with `TaskGroup`, release held peers in `finally` and verify that the intended worker dies. A test requiring traffic on both workers prepares that traffic explicitly instead of assuming OS load balancing, for example by sending each half while the other worker is suspended
Owned proxy workers configure a 32-thread executor before application startup, preserving additional worker startup hooks. Held synchronous provider calls must not consume a CPU-dependent executor limit before the rest of a fault burst can start. The proxy entrypoint imports LiteLLM inside `main`, so spawned workers begin supervision before importing the application. Worker healthchecks have a five-second deadline by default, so uvicorn kills a worker that stays paused longer than that. Tests never signal proxy processes themselves. `OwnedProxy.kill_worker` SIGKILLs a worker and returns once its replacement answers readiness with its own pid, `replacing_worker` does the same for a kill that has to land mid-burst, and `suspended` pauses a worker or the supervisor, which only works on a proxy launched with `healthcheck=WorkerHealthcheck.LENIENT`. The `kill()` that `replacing_worker` yields returns as soon as the worker is dead, so a test can release requests it holds on the surviving worker right away. Uvicorn may take a full worker healthcheck to notice the death, so the block waits for the supervisor to reap the worker and for the replacement only when it exits. Helper processes that are not the proxy stop through `stop_child`, `stop_processes` and `stop_process_group` and pause through `paused`. `scripts/check_test_quality.py` rule TQ010 rejects a raw `os.kill`, SIGKILL / SIGSTOP / SIGCONT, `.kill()`, `.terminate()`, `.suspend()` or `.send_signal()` anywhere under `tests/integration` outside `_support`, and TQ011 rejects an assertion that caps an elapsed time with a fixed number. Bound timings with `_support/timing.py` instead, relative to a measured control or to the injected delay. `run.py` measures how long this box takes to import the proxy application before it starts pytest, records it in `boot-calibration.json` beside the results and exports it as `INTEGRATION_IMPORT_SECONDS`. Readiness deadlines, graceful stops and the default per-test timeout are multiples of that measurement (`_support/budget.py`), so a slower box gets longer budgets and a two-worker proxy gets twice the readiness deadline of a one-worker one. Readiness requires responsive worker identities that are still alive and a positively identified owned listener. Unreadable socket records cannot establish ownership. Crash tests own their client tasks with `TaskGroup`, release held peers in `finally` and verify that the intended worker dies. A test requiring traffic on both workers prepares that traffic explicitly instead of assuming OS load balancing, for example by sending each half while the other worker is suspended
Any port the harness picks before the server that listens on it has started comes from `allocate_port()` in `_support/ports.py`. That covers owned proxies, scripted upstreams, fake providers, OTLP sinks, database relays and Redis. A server that binds port 0 itself and keeps that socket, such as `wire_server` without a port, needs no allocation. Each pytest process claims a block of 200 ports between 10000 and 32767 by holding an exclusive lock file per block in the system temp directory, so two runs on one machine, such as parallel worktrees, never share a block, and the OS drops the lock when the process exits. The range sits above the fixed service ports and below the Linux and macOS ephemeral ranges, so a port-0 bind or an outbound connection never lands in a block. Ports inside a block go out in order and are never handed out twice in a run, a port that something outside the harness already holds is skipped, and a process that uses up its block claims the next free one. A test that restarts a server on the same port allocates that port up front instead of reading it back from the first server. TQ012 in `scripts/check_test_quality.py` rejects reading a port back from a port-0 bind anywhere under `tests/integration`

View file

@ -150,9 +150,9 @@ def _is_reaped(process: psutil.Process) -> bool:
return False
def _wait_until_reaped(process: psutil.Process, supervisor: psutil.Process, seconds: float) -> None:
def _wait_until_reaped(process: psutil.Process, seconds: float) -> None:
deadline: Final = time.monotonic() + seconds
while not _is_reaped(process) and not _is_stopped(supervisor):
while not _is_reaped(process):
assert time.monotonic() < deadline, f"The supervisor did not reap worker {process.pid} within {seconds:.0f}s"
time.sleep(0.05)
@ -160,8 +160,6 @@ def _wait_until_reaped(process: psutil.Process, supervisor: psutil.Process, seco
@dataclass(frozen=True, slots=True)
class WorkerKill:
victim: psutil.Process
supervisor: psutil.Process
reap_seconds: float
fired: threading.Event
def freeze(self) -> psutil.Process:
@ -173,7 +171,6 @@ class WorkerKill:
with suppress(psutil.NoSuchProcess):
self.victim.kill()
_wait_until_gone(self.victim, STRICT_HEALTHCHECK_SECONDS)
_wait_until_reaped(self.victim, self.supervisor, self.reap_seconds)
self.fired.set()
@ -250,22 +247,22 @@ class OwnedProxy:
assert pid in before, f"{pid} is not a serving worker of this proxy: {sorted(before)}"
return psutil.Process(pid), before
def _worker_kill(self, victim: psutil.Process) -> WorkerKill:
reap_seconds: Final = healthcheck_seconds(self.healthcheck) + STRICT_HEALTHCHECK_SECONDS
return WorkerKill(victim, psutil.Process(self.process.pid), reap_seconds, threading.Event())
def _replacement_of(self, victim: psutil.Process, before: frozenset[int]) -> int:
_wait_until_reaped(victim, healthcheck_seconds(self.healthcheck) + STRICT_HEALTHCHECK_SECONDS)
return self._await_replacement(before)
@contextmanager
def replacing_worker(self, pid: int) -> Iterator[WorkerKill]:
victim, before = self._victim(pid)
kill: Final = self._worker_kill(victim)
kill: Final = WorkerKill(victim, threading.Event())
yield kill
assert kill.fired.is_set(), f"replacing_worker({pid}) exited without killing the worker"
self._await_replacement(before)
self._replacement_of(victim, before)
def kill_worker(self, pid: int) -> int:
victim, before = self._victim(pid)
self._worker_kill(victim)()
return self._await_replacement(before)
WorkerKill(victim, threading.Event())()
return self._replacement_of(victim, before)
@contextmanager
def suspended(self, pid: int) -> Iterator[psutil.Process]:

View file

@ -1,21 +1,19 @@
import json
import re
import threading
import uuid
from collections.abc import Callable, Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import Future, ThreadPoolExecutor
from contextlib import ExitStack
from hashlib import sha256
from pathlib import Path
from typing import Final
import httpx
import psutil
import pytest
import yaml
from integration._support.client import Gateway, eventually, held, object_value
from integration._support.database import read_rows
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from integration._support.process import WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process
from integration._support.vertex import service_account_json
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue, TypeAdapter
@ -38,7 +36,6 @@ _GEMINI_USAGE: Final = {
"totalTokenCount": _PROMPT_TOKENS + _COMPLETION_TOKENS,
}
_ZERO_USAGE: Final = {"promptTokenCount": 0, "candidatesTokenCount": 0, "totalTokenCount": 0}
_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
_BURST_SELECT: Final = (
'SELECT api_key, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs"'
" WHERE api_key = ANY(string_to_array(%s, ','))"
@ -67,6 +64,7 @@ _UPSTREAM_ASKS_FOR_USAGE: Final = {
},
}
_SETTING_VALUES: Final = {"always_off": False, "always_on": True}
_STREAMS_PER_WORKER: Final = 6
def _openai_frame(chunk: Mapping[str, JsonValue], response_id: str) -> bytes:
@ -273,20 +271,15 @@ def _stream_or_dropped(client: httpx.Client, key: str) -> httpx.Response | None:
return None
def _local_port(client: httpx.Client) -> int:
with client.stream("GET", "/health/liveliness") as response:
port: Final = int(response.extensions["network_stream"].get_extra_info("client_addr")[1])
response.read()
assert response.status_code == 200, response.text
return port
def _accepted_client_ports(pid: int, proxy_port: int) -> frozenset[int]:
return frozenset(
connection.raddr.port
for connection in psutil.Process(pid).net_connections(kind="tcp")
if connection.raddr and connection.laddr.port == proxy_port
def _held_streams(
stack: ExitStack, pool: ThreadPoolExecutor, url: str, keys: Sequence[str], wire: Wire, held: int
) -> tuple[Future[httpx.Response | None], ...]:
clients: Final = tuple(stack.enter_context(httpx.Client(base_url=url, timeout=60, trust_env=False)) for _ in keys)
futures: Final = tuple(
pool.submit(_stream_or_dropped, client, key) for client, key in zip(clients, keys, strict=True)
)
eventually(wire.received.qsize, lambda size: size >= held, seconds=60)
return futures
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
@ -296,35 +289,27 @@ def test_worker_killed_mid_stream_leaves_the_sibling_billing_each_stream_once(ga
wire: Final = stack.enter_context(wire_server(_peer(gate)))
stack.callback(gate.set)
config: Final = _owned_config(tmp_path / "worker-kill.yaml", gateway, wire, {})
owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2))
workers: Final = eventually(
lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())),
lambda pids: len(pids) == 2,
seconds=30,
owned: Final = stack.enter_context(
owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2, healthcheck=WorkerHealthcheck.LENIENT)
)
victim, survivor = sorted(eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=30))
scenario: Final = stack.enter_context(owned.gateway.scenario())
proxy_url: Final = owned.gateway.client.base_url
clients: Final = tuple(
stack.enter_context(httpx.Client(base_url=str(proxy_url), timeout=60, trust_env=False)) for _ in range(12)
)
keys: Final = tuple(scenario.key(models=[_ADAPTER_MODEL]) for _ in clients)
ports: Final = tuple(_local_port(client) for client in clients)
pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients)))
futures: Final = tuple(
pool.submit(_stream_or_dropped, client, key) for client, key in zip(clients, keys, strict=True)
)
eventually(wire.received.qsize, lambda size: size >= len(clients), seconds=60)
shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & frozenset(ports) for pid in workers}
assert sum(map(len, shares.values())) == len(clients), shares
victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid]))
victim_keys: Final = tuple(scenario.key(models=[_ADAPTER_MODEL]) for _ in range(_STREAMS_PER_WORKER))
survivor_keys: Final = tuple(scenario.key(models=[_ADAPTER_MODEL]) for _ in range(_STREAMS_PER_WORKER))
proxy_url: Final = str(owned.gateway.client.base_url)
pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=2 * _STREAMS_PER_WORKER))
with owned.suspended(survivor):
on_victim: Final = _held_streams(stack, pool, proxy_url, victim_keys, wire, _STREAMS_PER_WORKER)
with owned.suspended(victim):
on_survivor: Final = _held_streams(stack, pool, proxy_url, survivor_keys, wire, 2 * _STREAMS_PER_WORKER)
with owned.replacing_worker(victim) as kill:
kill()
gate.set()
results: Final = tuple(future.result(timeout=60) for future in futures)
survivors: Final = tuple(
(key, result) for port, key, result in zip(ports, keys, results, strict=True) if port not in shares[victim]
)
assert survivors, shares
dropped: Final = tuple(future.result(timeout=60) for future in on_victim)
survivors: Final = tuple(
zip(survivor_keys, (future.result(timeout=60) for future in on_survivor), strict=True)
)
assert dropped == (None,) * _STREAMS_PER_WORKER, dropped
for _, result in survivors:
assert result is not None
_assert_one_stop_with(result, _GEMINI_USAGE)

View file

@ -379,48 +379,47 @@ def test_a_replacing_block_that_never_kills_its_worker_fails_and_leaves_it_servi
assert owned.worker_pids() == frozenset((worker,))
def _kill_of(owned: OwnedProxy, worker: int) -> WorkerKill:
return WorkerKill(psutil.Process(worker), psutil.Process(owned.process.pid), 5, threading.Event())
@contextmanager
def _unreaped_child() -> Iterator[psutil.Process]:
with subprocess.Popen((sys.executable, "-I", "-c", "import time; time.sleep(60)")) as child:
try:
yield psutil.Process(child.pid)
finally:
child.kill()
def test_a_worker_kill_returns_once_the_supervisor_reaped_the_victim_and_cannot_fire_twice(tmp_path: Path) -> None:
with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker):
kill: Final = _kill_of(owned, worker)
def test_a_worker_kill_returns_once_the_victim_is_dead_before_anything_reaps_it_and_cannot_fire_twice() -> None:
with _unreaped_child() as victim:
kill: Final = WorkerKill(victim, threading.Event())
kill()
assert not psutil.pid_exists(worker)
assert victim.status() == psutil.STATUS_ZOMBIE
assert kill.fired.is_set()
with pytest.raises(AssertionError, match="already killed"):
kill()
def test_a_frozen_victim_stays_stopped_until_the_kill(tmp_path: Path) -> None:
with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker):
kill: Final = _kill_of(owned, worker)
def test_a_frozen_victim_stays_stopped_until_the_kill() -> None:
with _unreaped_child() as victim:
kill: Final = WorkerKill(victim, threading.Event())
assert kill.freeze().status() == psutil.STATUS_STOPPED
kill()
assert not psutil.pid_exists(worker)
assert victim.status() == psutil.STATUS_ZOMBIE
def test_a_kill_under_a_suspended_supervisor_returns_with_the_victim_left_unreaped(tmp_path: Path) -> None:
with _supervised(tmp_path, WorkerHealthcheck.LENIENT) as (owned, worker):
kill: Final = _kill_of(owned, worker)
with owned.suspended(owned.process.pid):
kill()
assert psutil.Process(worker).status() == psutil.STATUS_ZOMBIE
eventually(lambda: psutil.pid_exists(worker), lambda exists: not exists, seconds=5)
def _kill_after_freezing_the_supervisor(owned: OwnedProxy, worker: int, supervisor: psutil.Process) -> None:
with owned.replacing_worker(worker) as kill:
supervisor.suspend()
kill()
def test_a_kill_fails_when_the_supervisor_never_reaps_the_victim(tmp_path: Path) -> None:
def test_a_replacing_block_fails_when_the_supervisor_never_reaps_the_killed_worker(tmp_path: Path) -> None:
with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker):
unreaping: Final = psutil.Process(owned.process.pid)
kill: Final = WorkerKill(psutil.Process(worker), psutil.Process(os.getpid()), 0.5, threading.Event())
unreaping.suspend()
try:
with pytest.raises(AssertionError, match="did not reap"):
kill()
_kill_after_freezing_the_supervisor(owned, worker, unreaping)
finally:
unreaping.resume()
assert not kill.fired.is_set()
def test_a_strict_healthcheck_proxy_refuses_suspension_before_stopping_anything(tmp_path: Path) -> None:

View file

@ -2,7 +2,7 @@ import asyncio
import json
import logging
import os
from collections.abc import Awaitable, Callable, Mapping
from collections.abc import Awaitable, Callable, Mapping, Sequence
from contextlib import ExitStack, asynccontextmanager
from types import SimpleNamespace
from typing import Final
@ -17,7 +17,7 @@ from fastapi import HTTPException, Request
import litellm
from litellm._uuid import uuid
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse
from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse, ProxyException
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO
from litellm.proxy.management_endpoints.types import CustomOpenID
@ -26,9 +26,11 @@ from litellm.proxy.management_endpoints.ui_sso import (
MicrosoftSSOHandler,
SSOAuthenticationHandler,
_complete_cli_sso_callback_session,
_insert_sso_user_or_find_concurrent_insert,
_setup_team_mappings,
_sync_user_role_from_jwt_role_map,
cli_sso_callback,
insert_sso_user,
normalize_email,
process_sso_jwt_access_token,
sso_lookup_user_ids,
@ -1425,6 +1427,82 @@ async def test_upsert_sso_user_no_role_in_sso_response():
assert "user_role" not in call_args.kwargs["data"]
class _UserTable:
def __init__(self, *rows: LiteLLM_UserTable) -> None:
self.rows: dict[str, LiteLLM_UserTable] = {row.user_id: row for row in rows}
async def find_first(self, where: dict[str, dict[str, str]]) -> LiteLLM_UserTable | None:
((field, condition),) = where.items()
def matches(stored: object) -> bool:
if condition.get("mode") == "insensitive":
return str(stored).lower() == condition["equals"].lower()
return stored == condition["equals"]
return next((row for row in self.rows.values() if matches(getattr(row, field))), None)
async def find_unique(self, where: dict[str, str]) -> LiteLLM_UserTable | None:
return self.rows.get(where["user_id"])
def _users_db(table: _UserTable) -> SimpleNamespace:
database: Final = SimpleNamespace(litellm_usertable=table)
return SimpleNamespace(db=database, writer_db=database)
async def _first_sign_in(table: _UserTable, **join: object) -> NewUserResponse | LiteLLM_UserTable:
prisma: Final = _users_db(table)
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
return await _insert_sso_user_or_find_concurrent_insert(
insert=insert_sso_user(
result_openid=CustomOpenID(id="sso-subject", email="person@example.com", team_ids=[]),
user_defined_values={
"user_id": "sso-subject",
"user_email": "person@example.com",
"user_role": None,
"max_budget": None,
"budget_duration": None,
"models": [],
},
),
user_id="sso-subject",
prisma_client=prisma,
**join,
)
@pytest.mark.asyncio
async def test_first_sign_in_signs_in_to_the_account_a_concurrent_first_sign_in_created():
winner: Final = LiteLLM_UserTable(user_id="sso-subject", user_email="person@example.com", user_role="internal_user")
signed_in: Final = await _first_sign_in(_UserTable(winner))
assert signed_in == winner
@pytest.mark.asyncio
async def test_first_sign_in_does_not_sign_in_to_another_account_holding_the_email():
other: Final = LiteLLM_UserTable(user_id="someone-else", user_email="Person@example.com", user_role="proxy_admin")
with pytest.raises(ProxyException, match=r"User with email Person@example\.com already exists"):
await _first_sign_in(_UserTable(other))
@pytest.mark.asyncio
async def test_first_sign_in_joins_default_teams_the_concurrent_first_sign_in_has_not_joined_yet(monkeypatch):
monkeypatch.setattr(litellm, "default_internal_user_params", {"teams": ["team-joined", "team-pending"]})
table: Final = _UserTable(
LiteLLM_UserTable(user_id="sso-subject", user_email="person@example.com", teams=["team-joined"])
)
async def join(user: LiteLLM_UserTable, teams: Sequence[object], prisma_client: object) -> None:
table.rows[user.user_id] = user.model_copy(update={"teams": [*user.teams, *map(str, teams)]})
signed_in: Final = await _first_sign_in(table, join_default_teams=join)
assert signed_in.teams == ["team-joined", "team-pending"]
def test_get_user_email_and_id_extracts_microsoft_role():
"""
Test that _get_user_email_and_id_from_result extracts user_role from Microsoft SSO.