From 3a4cc872acca8cc74d399a713de0cb4536f7b5df Mon Sep 17 00:00:00 2001 From: Lin Junrong Date: Fri, 31 Jul 2026 23:34:43 +0800 Subject: [PATCH 1/2] fix: hold references to the background tasks in auth_checks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ten create_task() calls discard the returned task. The event loop only keeps a weak reference, so a task whose sole reference was the create_task() expression can be garbage collected before it runs. One of them is a database write, not telemetry: in get_user_object_from_db_or_cache, when a user is matched by email during SSO, the sso_user_id binding is persisted with asyncio.create_task( UserRepository(prisma_client).table.update( where={"user_id": response.user_id}, data={"sso_user_id": sso_user_id}, ) ) If that task is collected the login still succeeds, but the binding is never written — so the next login falls back to the email lookup again and the SSO identity never sticks. Silent, and it looks like a provider problem. The other nine are budget alerts (token, soft, max, organization, team, key). A collected alert simply never fires: no error, no notification, and the budget event is invisible. Applies the pattern already used in the repo (guardrail_hooks/akto/akto.py, management_helpers/audit_logs.py) — hold the task in a module-level set and drop it in a done callback. --- litellm/proxy/auth/auth_checks.py | 45 ++++++++++++++++++++++++------- 1 file changed, 35 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b1e444e55d6..f889b50ab06 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -254,6 +254,11 @@ db_cache_expiry: Final = DEFAULT_IN_MEMORY_TTL # refresh every 5s all_routes: Final = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value +# The event loop only keeps weak references to tasks, so a background task whose +# only reference was the create_task() call can be collected before it runs. Hold +# each one until it completes. +_background_tasks: set[asyncio.Task] = set() # mutable-ok: task registry + def _log_budget_lookup_failure(entity: str, error: Exception) -> None: """ @@ -1750,12 +1755,14 @@ async def _get_fuzzy_user_object( ) if response is not None and sso_user_id is not None: # update sso_user_id - asyncio.create_task( # background task to update user with sso id + _background_task = asyncio.create_task( # background task to update user with sso id _user_table(UserRepository(prisma_client)).update( where={"user_id": response.user_id}, data={"sso_user_id": sso_user_id}, ) ) + _background_tasks.add(_background_task) + _background_task.add_done_callback(_background_tasks.discard) return response @@ -3791,12 +3798,14 @@ async def _virtual_key_max_budget_check( key_alias=valid_token.key_alias, event_group=Litellm_EntityType.KEY, ) - asyncio.create_task( + _background_task = asyncio.create_task( proxy_logging_obj.budget_alerts( type="token_budget", user_info=call_info, ) ) + _background_tasks.add(_background_task) + _background_task.add_done_callback(_background_tasks.discard) #################################### # collect information for alerting # @@ -3895,12 +3904,14 @@ async def _virtual_key_soft_budget_check( event_group=Litellm_EntityType.KEY, ) - asyncio.create_task( + _background_task = asyncio.create_task( proxy_logging_obj.budget_alerts( type="soft_budget", user_info=call_info, ) ) + _background_tasks.add(_background_task) + _background_task.add_done_callback(_background_tasks.discard) def _parse_email_list(raw: str | Sequence[object] | None) -> list[str]: @@ -3986,12 +3997,14 @@ async def _virtual_key_max_budget_alert_check( event_group=Litellm_EntityType.KEY, max_budget_alert_emails=alert_email_config, ) - asyncio.create_task( + _background_task = asyncio.create_task( proxy_logging_obj.budget_alerts( type="max_budget_alert", user_info=call_info, ) ) + _background_tasks.add(_background_task) + _background_task.add_done_callback(_background_tasks.discard) else: # Old path: existing single 80% threshold — completely unchanged alert_threshold: Final = valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE @@ -4018,12 +4031,14 @@ async def _virtual_key_max_budget_alert_check( event_group=Litellm_EntityType.KEY, ) - asyncio.create_task( + _background_task = asyncio.create_task( proxy_logging_obj.budget_alerts( type="max_budget_alert", user_info=call_info, ) ) + _background_tasks.add(_background_task) + _background_task.add_done_callback(_background_tasks.discard) async def _check_team_member_budget( @@ -4182,12 +4197,14 @@ async def _team_max_budget_check( organization_id=valid_token.org_id, event_group=Litellm_EntityType.TEAM, ) - asyncio.create_task( + _background_task = asyncio.create_task( proxy_logging_obj.budget_alerts( type="team_budget", user_info=call_info, ) ) + _background_tasks.add(_background_task) + _background_task.add_done_callback(_background_tasks.discard) raise litellm.BudgetExceededError( current_cost=spend, @@ -4300,12 +4317,14 @@ async def _team_soft_budget_check( alert_emails=alert_emails, ) - asyncio.create_task( + _background_task = asyncio.create_task( proxy_logging_obj.budget_alerts( type="soft_budget", user_info=call_info, ) ) + _background_tasks.add(_background_task) + _background_task.add_done_callback(_background_tasks.discard) async def _project_max_budget_check( @@ -4344,12 +4363,14 @@ async def _project_max_budget_check( organization_id=valid_token.org_id, event_group=Litellm_EntityType.PROJECT, ) - asyncio.create_task( + _background_task = asyncio.create_task( proxy_logging_obj.budget_alerts( type="project_budget", user_info=call_info, ) ) + _background_tasks.add(_background_task) + _background_task.add_done_callback(_background_tasks.discard) raise litellm.BudgetExceededError( current_cost=project_object.spend, @@ -4396,12 +4417,14 @@ async def _project_soft_budget_check( organization_id=valid_token.org_id, event_group=Litellm_EntityType.PROJECT, ) - asyncio.create_task( + _background_task = asyncio.create_task( proxy_logging_obj.budget_alerts( type="soft_budget", user_info=call_info, ) ) + _background_tasks.add(_background_task) + _background_task.add_done_callback(_background_tasks.discard) def _project_cache_key(project_id: str) -> str: @@ -4564,12 +4587,14 @@ async def _organization_max_budget_check( organization_id=org_id, event_group=Litellm_EntityType.ORGANIZATION, ) - asyncio.create_task( + _background_task = asyncio.create_task( proxy_logging_obj.budget_alerts( type="organization_budget", user_info=call_info, ) ) + _background_tasks.add(_background_task) + _background_task.add_done_callback(_background_tasks.discard) raise litellm.BudgetExceededError( current_cost=org_spend, From 527bb6f99d1fcd236edc84eee03cebd8d8a19856 Mon Sep 17 00:00:00 2001 From: Lin Junrong Date: Mon, 17 Aug 2026 16:36:20 +0800 Subject: [PATCH 2/2] refactor(auth_checks): hold background tasks via a helper The per-site `_background_task = ...` temporaries each tripped LIT010 (assignment without a `Final` declaration), pushing the repo total to 16739 over the 16731 budget. Annotating them Final is not an option: `_virtual_key_max_budget_alert_check` binds the name in both branches of an if/else, and two Final declarations of one name in a scope is a redeclaration error. Collapsing the 2-line hold idiom into `_hold_background_task()` removes the temporaries entirely, so LIT010 for this file returns to its baseline 73 and the total type-discipline count lands at 301 -- four below the 305 the untouched file already carries. --- litellm/proxy/auth/auth_checks.py | 128 ++++++++++++++++-------------- 1 file changed, 67 insertions(+), 61 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index f889b50ab06..ab3e39ba84d 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -257,7 +257,13 @@ all_routes: Final = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management # The event loop only keeps weak references to tasks, so a background task whose # only reference was the create_task() call can be collected before it runs. Hold # each one until it completes. -_background_tasks: set[asyncio.Task] = set() # mutable-ok: task registry +_background_tasks: Final[set[asyncio.Task]] = set() # mutable-ok: task registry + + +def _hold_background_task(task: asyncio.Task) -> None: + """Keep a strong reference to a fire-and-forget task until it finishes.""" + _background_tasks.add(task) + task.add_done_callback(_background_tasks.discard) def _log_budget_lookup_failure(entity: str, error: Exception) -> None: @@ -1755,14 +1761,14 @@ async def _get_fuzzy_user_object( ) if response is not None and sso_user_id is not None: # update sso_user_id - _background_task = asyncio.create_task( # background task to update user with sso id - _user_table(UserRepository(prisma_client)).update( - where={"user_id": response.user_id}, - data={"sso_user_id": sso_user_id}, + _hold_background_task( + asyncio.create_task( # background task to update user with sso id + _user_table(UserRepository(prisma_client)).update( + where={"user_id": response.user_id}, + data={"sso_user_id": sso_user_id}, + ) ) ) - _background_tasks.add(_background_task) - _background_task.add_done_callback(_background_tasks.discard) return response @@ -3798,14 +3804,14 @@ async def _virtual_key_max_budget_check( key_alias=valid_token.key_alias, event_group=Litellm_EntityType.KEY, ) - _background_task = asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="token_budget", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="token_budget", + user_info=call_info, + ) ) ) - _background_tasks.add(_background_task) - _background_task.add_done_callback(_background_tasks.discard) #################################### # collect information for alerting # @@ -3904,14 +3910,14 @@ async def _virtual_key_soft_budget_check( event_group=Litellm_EntityType.KEY, ) - _background_task = asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="soft_budget", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="soft_budget", + user_info=call_info, + ) ) ) - _background_tasks.add(_background_task) - _background_task.add_done_callback(_background_tasks.discard) def _parse_email_list(raw: str | Sequence[object] | None) -> list[str]: @@ -3997,14 +4003,14 @@ async def _virtual_key_max_budget_alert_check( event_group=Litellm_EntityType.KEY, max_budget_alert_emails=alert_email_config, ) - _background_task = asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="max_budget_alert", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="max_budget_alert", + user_info=call_info, + ) ) ) - _background_tasks.add(_background_task) - _background_task.add_done_callback(_background_tasks.discard) else: # Old path: existing single 80% threshold — completely unchanged alert_threshold: Final = valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE @@ -4031,14 +4037,14 @@ async def _virtual_key_max_budget_alert_check( event_group=Litellm_EntityType.KEY, ) - _background_task = asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="max_budget_alert", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="max_budget_alert", + user_info=call_info, + ) ) ) - _background_tasks.add(_background_task) - _background_task.add_done_callback(_background_tasks.discard) async def _check_team_member_budget( @@ -4197,14 +4203,14 @@ async def _team_max_budget_check( organization_id=valid_token.org_id, event_group=Litellm_EntityType.TEAM, ) - _background_task = asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="team_budget", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="team_budget", + user_info=call_info, + ) ) ) - _background_tasks.add(_background_task) - _background_task.add_done_callback(_background_tasks.discard) raise litellm.BudgetExceededError( current_cost=spend, @@ -4317,14 +4323,14 @@ async def _team_soft_budget_check( alert_emails=alert_emails, ) - _background_task = asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="soft_budget", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="soft_budget", + user_info=call_info, + ) ) ) - _background_tasks.add(_background_task) - _background_task.add_done_callback(_background_tasks.discard) async def _project_max_budget_check( @@ -4363,14 +4369,14 @@ async def _project_max_budget_check( organization_id=valid_token.org_id, event_group=Litellm_EntityType.PROJECT, ) - _background_task = asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="project_budget", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="project_budget", + user_info=call_info, + ) ) ) - _background_tasks.add(_background_task) - _background_task.add_done_callback(_background_tasks.discard) raise litellm.BudgetExceededError( current_cost=project_object.spend, @@ -4417,14 +4423,14 @@ async def _project_soft_budget_check( organization_id=valid_token.org_id, event_group=Litellm_EntityType.PROJECT, ) - _background_task = asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="soft_budget", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="soft_budget", + user_info=call_info, + ) ) ) - _background_tasks.add(_background_task) - _background_task.add_done_callback(_background_tasks.discard) def _project_cache_key(project_id: str) -> str: @@ -4587,14 +4593,14 @@ async def _organization_max_budget_check( organization_id=org_id, event_group=Litellm_EntityType.ORGANIZATION, ) - _background_task = asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="organization_budget", - user_info=call_info, + _hold_background_task( + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="organization_budget", + user_info=call_info, + ) ) ) - _background_tasks.add(_background_task) - _background_task.add_done_callback(_background_tasks.discard) raise litellm.BudgetExceededError( current_cost=org_spend,