fix(datadog): drain cost-management queue + opt-in FinOps tag allowlist

litellm/integrations/datadog/datadog_cost_management.py

@@ -2,10 +2,17 @@
 import os
 import time
 from datetime import datetime
-from typing import Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, cast
 
 from litellm._logging import verbose_logger
 from litellm.integrations.custom_batch_logger import CustomBatchLogger
+from litellm.integrations.datadog.datadog_handler import (
+    get_datadog_env,
+    get_datadog_hostname,
+    get_datadog_pod_name,
+    get_datadog_service,
+)
+from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
     httpxSpecialProvider,
@@ -15,9 +22,30 @@
 )
 from litellm.types.utils import StandardLoggingPayload
 
+# Reserved tag keys whose values come from trusted sources (infra env, LiteLLM
+# core payload fields, or proxy-controlled auth metadata). User-supplied
+# request_tags / metadata cannot overwrite these, even when the key is
+# allowlisted via cost_tag_keys, because that would let an authenticated caller
+# spoof cost attribution (e.g. request_tags=["team:victim-team"]).
+_RESERVED_TAG_KEYS: frozenset = frozenset(
+    {
+        "env",
+        "service",
+        "host",
+        "pod_name",
+        "provider",
+        "model",
+        "model_id",
+        "team",
+        "user",
+        "model_group",
+    }
+)
+
 
 class DatadogCostManagementLogger(CustomBatchLogger):
-    def __init__(self, **kwargs):
+    def __init__(self, cost_tag_keys: Optional[List[str]] = None, **kwargs):
+        self.cost_tag_keys: List[str] = list(cost_tag_keys) if cost_tag_keys else []
         self.dd_api_key = os.getenv("DD_API_KEY")
         self.dd_app_key = os.getenv("DD_APP_KEY")
         self.dd_site = os.getenv("DD_SITE", "datadoghq.com")
@@ -68,20 +96,21 @@ async def async_send_batch(self):
         if not self.log_queue:
             return
 
-        try:
-            # Aggregate costs from the batch
-            aggregated_entries = self._aggregate_costs(self.log_queue)
+        batch_to_send = self.log_queue[:]
+        self.log_queue = []
 
+        try:
+            aggregated_entries = self._aggregate_costs(batch_to_send)
             if not aggregated_entries:
+                verbose_logger.debug(
+                    "Datadog Cost Management: batch produced no aggregable entries; "
+                    "dropping %d log(s) from queue.",
+                    len(batch_to_send),
+                )
                 return
-
-            # Send to Datadog
             await self._upload_to_datadog(aggregated_entries)
-
-            # Clear queue only on success (or if we decide to drop on failure)
-            # CustomBatchLogger clears queue in flush_queue, so we just process here
-
         except Exception as e:
+            self.log_queue = batch_to_send + self.log_queue
             verbose_logger.exception(
                 f"Datadog Cost Management: Error in async_send_batch: {str(e)}"
             )
@@ -151,45 +180,81 @@ def _aggregate_costs(
         return list(aggregator.values())
 
     def _extract_tags(self, log: StandardLoggingPayload) -> Dict[str, str]:
-        from litellm.integrations.datadog.datadog_handler import (
-            get_datadog_env,
-            get_datadog_hostname,
-            get_datadog_pod_name,
-            get_datadog_service,
-        )
-
-        tags = {
+        tags: Dict[str, str] = {
             "env": get_datadog_env(),
             "service": get_datadog_service(),
             "host": get_datadog_hostname(),
             "pod_name": get_datadog_pod_name(),
         }
 
-        # Add metadata as tags
-        metadata = log.get("metadata", {})
-        if metadata:
-            # Add user info
-            # Add user info
-            if metadata.get("user_api_key_alias"):
-                tags["user"] = str(metadata["user_api_key_alias"])
-
-            # Add Team Tag
-            team_tag = (
-                metadata.get("user_api_key_team_alias")
-                or metadata.get("team_alias")  # type: ignore
-                or metadata.get("user_api_key_team_id")
-                or metadata.get("team_id")  # type: ignore
-            )
-
-            if team_tag:
-                tags["team"] = str(team_tag)
-            # model_group is not in StandardLoggingMetadata TypedDict, so we need to access it via dict.get()
-            model_group = metadata.get("model_group")  # type: ignore[misc]
-            if model_group:
-                tags["model_group"] = str(model_group)
+        # Always-on canonical FOCUS dimensions from top-level payload fields.
+        # Non-sensitive and required for Datadog Custom Costs per-model attribution.
+        self._add_tag(tags, "provider", log.get("custom_llm_provider"))
+        self._add_tag(tags, "model", log.get("model"))
+        self._add_tag(tags, "model_id", log.get("model_id"))
+
+        # cast because StandardLoggingMetadata is a TypedDict; we iterate it
+        # as a generic mapping below.
+        metadata: Dict[str, Any] = cast(Dict[str, Any], log.get("metadata") or {})
+
+        # Backwards-compat: team/user/model_group preserved regardless of allowlist.
+        if metadata.get("user_api_key_alias"):
+            tags["user"] = str(metadata["user_api_key_alias"])
+        team_tag = (
+            metadata.get("user_api_key_team_alias")
+            or metadata.get("team_alias")
+            or metadata.get("user_api_key_team_id")
+            or metadata.get("team_id")
+        )
+        if team_tag:
+            tags["team"] = str(team_tag)
+        if metadata.get("model_group"):
+            tags["model_group"] = str(metadata["model_group"])
+
+        # Allowlist-gated: request_tags (split on `:`) and arbitrary metadata.*.
+        # Reserved keys are hard-blocked here regardless of allowlist membership —
+        # see _RESERVED_TAG_KEYS for the rationale.
+        if self.cost_tag_keys:
+            allow = set(self.cost_tag_keys)
+            for rt in log.get("request_tags") or []:
+                if not isinstance(rt, str) or ":" not in rt:
+                    continue
+                k, _, v = rt.partition(":")
+                if k in allow and v:
+                    self._set_custom_tag(tags, k, v)
+            for k, v in metadata.items():
+                if k in allow and v is not None and not isinstance(v, (dict, list)):
+                    self._set_custom_tag(tags, k, str(v))
+            for nested_key in ("spend_logs_metadata", "requester_metadata"):
+                nested = metadata.get(nested_key)
+                if isinstance(nested, dict):
+                    for k, v in nested.items():
+                        if (
+                            k in allow
+                            and v is not None
+                            and not isinstance(v, (dict, list))
+                        ):
+                            self._set_custom_tag(tags, k, str(v))
 
         return tags
 
+    @staticmethod
+    def _set_custom_tag(tags: Dict[str, str], key: str, value: str) -> None:
+        if key in _RESERVED_TAG_KEYS:
+            verbose_logger.debug(
+                "Datadog Cost Management: dropping user-supplied tag %r=%r — "
+                "key is reserved for trusted cost attribution.",
+                key,
+                value,
+            )
+            return
+        tags[key] = value
+
+    @staticmethod
+    def _add_tag(tags: Dict[str, str], key: str, value: Any) -> None:
+        if value:
+            tags[key] = str(value)
+
     async def _upload_to_datadog(self, payload: List[Dict]):
         if not self.dd_api_key or not self.dd_app_key:
             return
@@ -201,8 +266,6 @@ async def _upload_to_datadog(self, payload: List[Dict]):
         }
 
         # The API endpoint expects a list of objects directly in the body (file content behavior)
-        from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
-
         data_json = safe_dumps(payload)
 
         response = await self.async_client.put(

litellm/proxy/common_utils/callback_utils.py

@@ -302,7 +302,15 @@ def initialize_callbacks_on_proxy(  # noqa: PLR0915
                     DatadogCostManagementLogger,
                 )
 
-                datadog_cost_management_obj = DatadogCostManagementLogger()
+                init_params = {}
+                if (
+                    "datadog_cost_management" in callback_specific_params
+                    and isinstance(
+                        callback_specific_params["datadog_cost_management"], dict
+                    )
+                ):
+                    init_params = callback_specific_params["datadog_cost_management"]
+                datadog_cost_management_obj = DatadogCostManagementLogger(**init_params)
                 imported_list.append(datadog_cost_management_obj)
             elif isinstance(callback, CustomLogger):
                 imported_list.append(callback)

litellm/types/integrations/datadog_cost_management.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, TypedDict
+from typing import Dict, List, Optional, TypedDict
 
 
 from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams
@@ -9,7 +9,7 @@ class DatadogCostManagementInitParams(StandardCustomLoggerInitParams):
     Init params for Datadog Cost Management
     """
 
-    datadog_cost_management_params: Optional[Dict] = None
+    cost_tag_keys: Optional[List[str]] = None
 
 
 class DatadogFOCUSCostEntry(TypedDict):