fix(proxy): tag the deprecation email cache calls with their service family

litellm/proxy/common_utils/model_deprecation_notifications.py

@@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Final, Protocol
 
 from pydantic import BaseModel, ConfigDict, Field, ValidationError
 
+from litellm._internal_context import with_service_target
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import EMAIL_MODEL_DEPRECATION_LOCK_ID
 from litellm.integrations.email_alerting import (
@@ -288,6 +289,9 @@ def render_model_deprecation_email(
     return subject, body + EMAIL_FOOTER
 
 
+_DEPRECATION_EMAIL_TARGET: Final = "model_deprecation_emails"
+
+
 _ENTERPRISE_BASE_EMAIL_LOGGER: Final = (
     "litellm_enterprise.enterprise_callbacks.send_emails.base_email",
     "BaseEmailLogger",
@@ -351,6 +355,7 @@ async def _stamp_pass(cache: DeprecationEmailCache) -> None:
     )
 
 
+@with_service_target(_DEPRECATION_EMAIL_TARGET)
 async def email_pass_done_today(cache: DeprecationEmailCache) -> bool:
     return (await cache.async_get_cache(key=SlackAlertingCacheKeys.deprecation_email_pass_key.value)) is not None
 
@@ -390,6 +395,7 @@ async def _send_team_notification(notification: TeamNotification, ctx: Deprecati
     return True
 
 
+@with_service_target(_DEPRECATION_EMAIL_TARGET)
 async def send_model_deprecation_emails(ctx: DeprecationEmailContext) -> int:
     """One pass: resolve affected teams, claim the day only when there is something to send, deliver, then stamp
 

tests/unit/proxy/common_utils/test_model_deprecation_notifications.py

@@ -10,6 +10,7 @@ from unittest.mock import MagicMock
 
 import pytest
 
+from litellm._internal_context import current_service_target
 from litellm.caching.caching import DualCache
 from litellm.constants import EMAIL_MODEL_DEPRECATION_LOCK_ID
 from litellm.models.team import LiteLLM_TeamTable
@@ -469,6 +470,20 @@ class _RecordingCache:
         self.ttls[key] = ttl
 
 
+class _TargetRecordingCache(_RecordingCache):
+    def __init__(self):
+        super().__init__()
+        self.targets = []
+
+    async def async_get_cache(self, *, key: str) -> object:
+        self.targets.append(current_service_target())
+        return await super().async_get_cache(key=key)
+
+    async def async_set_cache(self, *, key: str, value: float, ttl: float) -> None:
+        self.targets.append(current_service_target())
+        await super().async_set_cache(key=key, value=value, ttl=ttl)
+
+
 def _team_row(team_id: str, alias: str = "Team") -> Mapping[str, object]:
     return {
         "team_id": team_id,
@@ -707,6 +722,18 @@ async def test_should_report_whether_the_daily_pass_already_ran():
     assert await email_pass_done_today(cache) is True
 
 
+@pytest.mark.asyncio
+async def test_should_tag_every_cache_read_and_write_with_the_deprecation_email_family():
+    cache: Final = _TargetRecordingCache()
+    ctx: Final = _context(_router([DEAD_DEPLOYMENT]), _prisma(TWO_TEAMS[:1], TWO_ADMINS), _Deliverer(), cache)
+
+    assert await email_pass_done_today(cache) is False
+    assert await send_model_deprecation_emails(ctx) == 1
+    assert len(cache.targets) >= 3
+    assert set(cache.targets) == {"model_deprecation_emails"}
+    assert current_service_target() is None
+
+
 class TestMakeEmailDeliverer:
     @pytest.mark.asyncio
     async def test_should_use_the_configured_email_logger_once_per_recipient(self):