From 527bb6f99d1fcd236edc84eee03cebd8d8a19856 Mon Sep 17 00:00:00 2001 From: Lin Junrong Date: Mon, 17 Aug 2026 16:36:20 +0800 Subject: [PATCH] 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,