diff --git a/litellm/__init__.py b/litellm/__init__.py index bcac2f72f39..514e2519504 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -154,6 +154,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "gitlab", "cloudzero", "focus", + "mavvrik", "vantage", "posthog", "levo", diff --git a/litellm/constants.py b/litellm/constants.py index 57f55e6c177..714cfe9b114 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1480,6 +1480,7 @@ DB_SPEND_UPDATE_JOB_NAME = "db_spend_update_job" DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME = "db_daily_tag_spend_update_job" PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME = "prometheus_emit_budget_metrics" CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME = "cloudzero_export_usage_data" +MAVVRIK_FOCUS_EXPORT_JOB_NAME = "mavvrik_focus_export_usage_data" CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int( os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000) ) diff --git a/litellm/integrations/focus/destinations/__init__.py b/litellm/integrations/focus/destinations/__init__.py index e0cd90c1d61..21945c9b457 100644 --- a/litellm/integrations/focus/destinations/__init__.py +++ b/litellm/integrations/focus/destinations/__init__.py @@ -4,6 +4,7 @@ from .base import FocusDestination, FocusTimeWindow from .factory import FocusDestinationFactory from .gcs_destination import FocusGCSDestination from .s3_destination import FocusS3Destination +from .mavvrik_destination import FocusMavvrikDestination from .vantage_destination import FocusVantageDestination __all__ = [ @@ -12,5 +13,6 @@ __all__ = [ "FocusGCSDestination", "FocusTimeWindow", "FocusS3Destination", + "FocusMavvrikDestination", "FocusVantageDestination", ] diff --git a/litellm/integrations/focus/destinations/factory.py b/litellm/integrations/focus/destinations/factory.py index 7ce21d4040a..cd25a87729f 100644 --- a/litellm/integrations/focus/destinations/factory.py +++ b/litellm/integrations/focus/destinations/factory.py @@ -8,6 +8,7 @@ from typing import Any, Dict, Optional from .base import FocusDestination from .gcs_destination import FocusGCSDestination from .s3_destination import FocusS3Destination +from .mavvrik_destination import FocusMavvrikDestination from .vantage_destination import FocusVantageDestination @@ -32,6 +33,8 @@ class FocusDestinationFactory: return FocusVantageDestination(prefix=prefix, config=normalized_config) if provider_lower == "gcs": return FocusGCSDestination(prefix=prefix, config=normalized_config) + if provider_lower == "mavvrik": + return FocusMavvrikDestination(prefix=prefix, config=normalized_config) raise NotImplementedError( f"Provider '{provider}' not supported for Focus export" ) @@ -87,6 +90,15 @@ class FocusDestinationFactory: "FOCUS_GCS_BUCKET_NAME must be provided for GCS exports" ) return {k: v for k, v in resolved.items() if v is not None} + if provider == "mavvrik": + resolved = { + "api_key": overrides.get("api_key") or os.getenv("MAVVRIK_API_KEY"), + "api_endpoint": overrides.get("api_endpoint") + or os.getenv("MAVVRIK_API_ENDPOINT"), + "connection_id": overrides.get("connection_id") + or os.getenv("MAVVRIK_CONNECTION_ID"), + } + return {k: v for k, v in resolved.items() if v is not None} raise NotImplementedError( f"Provider '{provider}' not supported for Focus export configuration" ) diff --git a/litellm/integrations/focus/destinations/mavvrik_destination.py b/litellm/integrations/focus/destinations/mavvrik_destination.py new file mode 100644 index 00000000000..1e3c98b9a70 --- /dev/null +++ b/litellm/integrations/focus/destinations/mavvrik_destination.py @@ -0,0 +1,345 @@ +"""Mavvrik GCS destination for FOCUS export. + +Flow: + 1. GET /metrics/agent/ai/{connection_id}/upload-url → GCS signed URL + 2. PUT with CSV content +""" + +from __future__ import annotations + +import gzip +from typing import Any, Optional +from urllib.parse import urlparse + +from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, + httpxSpecialProvider, +) + +from .base import FocusDestination, FocusTimeWindow + +_MAVVRIK_ALLOWED_SUFFIXES = (".mavvrik.dev", ".mavvrik.ai", ".mavvrik.app") + +# GCS requires intermediate chunks to be a multiple of 256 KB. +# 8 MB gives a good balance between round-trips and memory pressure. +_GCS_CHUNK_SIZE = 8 * 1024 * 1024 # 8 MB + + +def _validate_api_endpoint(api_endpoint: str) -> None: + if not api_endpoint.startswith("https://"): + raise ValueError("MAVVRIK_API_ENDPOINT must be an HTTPS URL") + hostname = (urlparse(api_endpoint).hostname or "").lower() + if not any(hostname.endswith(suffix) for suffix in _MAVVRIK_ALLOWED_SUFFIXES): + raise ValueError( + "MAVVRIK_API_ENDPOINT host must be a Mavvrik domain " + "(e.g. https://api.mavvrik.dev/)" + ) + + +def _validate_gcs_url(url: str, label: str) -> None: + parsed = urlparse(url) + if parsed.scheme != "https": + raise ValueError( + f"Mavvrik FOCUS destination: {label} must be HTTPS, got scheme '{parsed.scheme}'" + ) + hostname = (parsed.hostname or "").lower() + if not ( + hostname == "storage.googleapis.com" + or hostname.endswith(".storage.googleapis.com") + ): + raise ValueError( + f"Mavvrik FOCUS destination: {label} must be a GCS endpoint " + f"(storage.googleapis.com), got '{hostname}'" + ) + + +class FocusMavvrikDestination(FocusDestination): + """Upload FOCUS CSV exports to Mavvrik via GCS signed URL.""" + + def __init__( + self, + *, + prefix: str, + config: Optional[dict[str, Any]] = None, + ) -> None: + config = config or {} + api_key = config.get("api_key") + api_endpoint = config.get("api_endpoint") + connection_id = config.get("connection_id") + + if not api_key: + raise ValueError( + "MAVVRIK_API_KEY must be provided for Mavvrik FOCUS destination " + "(set MAVVRIK_API_KEY env var or pass in destination_config)" + ) + if not api_endpoint: + raise ValueError( + "MAVVRIK_API_ENDPOINT must be provided for Mavvrik FOCUS destination " + "(set MAVVRIK_API_ENDPOINT env var or pass in destination_config)" + ) + if not connection_id: + raise ValueError( + "MAVVRIK_CONNECTION_ID must be provided for Mavvrik FOCUS destination " + "(set MAVVRIK_CONNECTION_ID env var or pass in destination_config)" + ) + + _validate_api_endpoint(api_endpoint) + + self.api_key = api_key + self.api_endpoint = api_endpoint.rstrip("/") + self.connection_id = connection_id + self.prefix = prefix + self._http: AsyncHTTPHandler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback + ) + self._registered = False + + @property + def _agent_url(self) -> str: + return f"{self.api_endpoint}/metrics/agent/ai/{self.connection_id}" + + @property + def _upload_url_endpoint(self) -> str: + return f"{self.api_endpoint}/metrics/agent/ai/{self.connection_id}/upload-url" + + @property + def _auth_headers(self) -> dict[str, str]: + return {"Content-Type": "application/json", "x-api-key": self.api_key} + + async def _ensure_registered(self) -> Optional[int]: + """POST agent endpoint to register/initialize the connector (once per instance). + + Returns metricsMarker from the Mavvrik response — the last date index + Mavvrik has successfully processed. Used by the logger to catch up any + dates that were missed due to previous export failures. + + Returns None if the connector was already registered (cached). + """ + if self._registered: + return None + resp = await self._http.client.request( + method="POST", + url=self._agent_url, + headers=self._auth_headers, + json={"name": self.connection_id}, + timeout=30.0, + ) + if resp.status_code == 410: + # Connector has been disconnected in Mavvrik — reset flag so next + # delivery attempt re-registers after it becomes active again. + self._registered = False + raise RuntimeError( + "Mavvrik FOCUS destination: connector is disconnected (410). " + "Re-enable the connection in the Mavvrik dashboard." + ) + if resp.status_code >= 400: + raise RuntimeError( + f"Mavvrik FOCUS destination: register failed " + f"({resp.status_code}): {resp.text[:200]}" + ) + self._registered = True + metrics_marker = resp.json().get("metricsMarker", 0) + verbose_logger.debug( + "Mavvrik FOCUS destination: connector registered (metricsMarker=%s)", + metrics_marker, + ) + return metrics_marker + + async def _get_signed_url(self, date_str: str) -> str: + """GET upload-url endpoint → GCS signed URL for the given date.""" + params = {"name": date_str, "type": "metrics", "datetime": date_str} + resp = await self._http.client.request( + method="GET", + url=self._upload_url_endpoint, + headers=self._auth_headers, + params=params, + timeout=30.0, + ) + if resp.status_code >= 400: + raise RuntimeError( + f"Mavvrik FOCUS destination: failed to get signed URL " + f"({resp.status_code}): {resp.text[:200]}" + ) + signed_url = resp.json().get("url") + if not signed_url: + raise RuntimeError( + f"Mavvrik FOCUS destination: response missing 'url' field: {resp.json()}" + ) + _validate_gcs_url(signed_url, "signed URL") + verbose_logger.debug( + "Mavvrik FOCUS destination: got signed URL for date %s", date_str + ) + return signed_url + + async def _upload_to_gcs(self, signed_url: str, content: bytes) -> None: + """Upload gzip-compressed CSV to GCS via chunked resumable upload. + + The full CSV is gzip-compressed first, then uploaded in _GCS_CHUNK_SIZE + chunks using the GCS resumable upload protocol. GCS assembles the chunks + server-side into a single complete object — the bucket receives one file + regardless of how many chunks were sent. + + Intermediate chunks: Content-Range: bytes X-Y/* → expect 308 + Final chunk: Content-Range: bytes X-Y/T → expect 200/201 + + This handles exports larger than available memory for a single PUT while + keeping the destination code self-contained (no changes to the FOCUS + pipeline upstream). + """ + gzip_bytes = gzip.compress(content) + total = len(gzip_bytes) + + # Step 1: initiate resumable upload session + metadata = b'{"contentEncoding":"gzip","contentDisposition":"attachment"}' + init_resp = await self._http.client.request( + method="POST", + url=signed_url, + headers={ + "Content-Type": "application/gzip", + "x-goog-resumable": "start", + }, + content=metadata, + timeout=30.0, + ) + if init_resp.status_code not in (200, 201): + raise RuntimeError( + f"Mavvrik FOCUS destination: GCS session init failed " + f"({init_resp.status_code}): {init_resp.text[:400]}" + ) + + session_uri = init_resp.headers.get("Location") + if not session_uri: + raise RuntimeError( + "Mavvrik FOCUS destination: GCS session init missing Location header" + ) + _validate_gcs_url(session_uri, "session URI") + + verbose_logger.debug( + "Mavvrik FOCUS destination: GCS session started, uploading %d gzip bytes " + "in %d chunk(s)", + total, + max(1, -(-total // _GCS_CHUNK_SIZE)), # ceiling division + ) + + # Step 2: upload in chunks; cancel session on any failure to avoid + # lingering GCS sessions (they stay open for ~1 week otherwise). + offset = 0 + try: + while offset < total: + chunk = gzip_bytes[offset : offset + _GCS_CHUNK_SIZE] + chunk_end = offset + len(chunk) - 1 + is_final = (offset + len(chunk)) >= total + content_range = ( + f"bytes {offset}-{chunk_end}/{total}" + if is_final + else f"bytes {offset}-{chunk_end}/*" + ) + expected_statuses = {200, 201} if is_final else {308} + + resp = await self._http.client.request( + method="PUT", + url=session_uri, + headers={ + "Content-Type": "application/gzip", + "Content-Range": content_range, + }, + content=chunk, + timeout=120.0, + ) + if resp.status_code not in expected_statuses: + raise RuntimeError( + f"Mavvrik FOCUS destination: GCS chunk upload failed " + f"(chunk offset={offset}, expected={expected_statuses}, " + f"got={resp.status_code}): {resp.text[:400]}" + ) + offset += len(chunk) + verbose_logger.debug( + "Mavvrik FOCUS destination: uploaded chunk offset=%d/%d", + offset, + total, + ) + except Exception: + # Cancel the open GCS session so it doesn't linger for up to 1 week. + try: + await self._http.client.request( + method="DELETE", url=session_uri, timeout=10.0 + ) + verbose_logger.debug( + "Mavvrik FOCUS destination: cancelled GCS session after error" + ) + except Exception: + pass + raise + + async def get_metrics_marker(self) -> Optional[int]: + """Register with Mavvrik and return the current metricsMarker. + + The metricsMarker is a Unix timestamp (seconds) representing the last + date Mavvrik has successfully ingested. Called on every scheduled run + so the logger can detect and catch up any dates missed due to previous + export failures. + + Always calls the Mavvrik register API — unlike deliver() which skips + registration once _registered is True, catch-up requires a fresh + marker value on every run. + """ + resp = await self._http.client.request( + method="POST", + url=self._agent_url, + headers=self._auth_headers, + json={"name": self.connection_id}, + timeout=30.0, + ) + if resp.status_code == 410: + self._registered = False + raise RuntimeError( + "Mavvrik FOCUS destination: connector is disconnected (410). " + "Re-enable the connection in the Mavvrik dashboard." + ) + if resp.status_code >= 400: + raise RuntimeError( + f"Mavvrik FOCUS destination: register failed " + f"({resp.status_code}): {resp.text[:200]}" + ) + self._registered = True + metrics_marker = resp.json().get("metricsMarker", 0) + verbose_logger.debug( + "Mavvrik FOCUS destination: got metricsMarker=%s", metrics_marker + ) + return metrics_marker + + async def deliver( + self, + *, + content: bytes, + time_window: FocusTimeWindow, + filename: str, + ) -> None: + """Upload FOCUS CSV to Mavvrik via GCS signed URL. + + Uses the start date of the time window as the object date key. + """ + if not content: + verbose_logger.debug( + "Mavvrik FOCUS destination: empty content, skipping upload" + ) + return + + date_str = time_window.start_time.strftime("%Y-%m-%d") + + verbose_logger.debug( + "Mavvrik FOCUS destination: uploading %d bytes for date=%s (%s)", + len(content), + date_str, + filename, + ) + + await self._ensure_registered() + signed_url = await self._get_signed_url(date_str) + await self._upload_to_gcs(signed_url, content) + + verbose_logger.debug( + "Mavvrik FOCUS destination: upload complete for date=%s", date_str + ) diff --git a/litellm/integrations/mavvrik_focus/__init__.py b/litellm/integrations/mavvrik_focus/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py new file mode 100644 index 00000000000..ddea7a3d604 --- /dev/null +++ b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py @@ -0,0 +1,272 @@ +"""MavvrikFocusLogger — FOCUS-based Mavvrik export logger. + +Usage in config.yaml: + litellm_settings: + callbacks: ["mavvrik"] + +Required env vars: + MAVVRIK_API_KEY + MAVVRIK_API_ENDPOINT + MAVVRIK_CONNECTION_ID + +Optional env vars: + MAVVRIK_FOCUS_MAX_ROWS — row cap per export window (default: 500000) + +Only daily frequency is supported. The Mavvrik ingestion protocol stores one +file per calendar date (metrics/YYYY-MM-DD). Hourly or interval exports would +overwrite each other within the same day, producing incomplete data. +""" + +from __future__ import annotations + +import os +from datetime import datetime, timedelta, timezone +from typing import TYPE_CHECKING, Any, List, Optional + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.constants import MAVVRIK_FOCUS_EXPORT_JOB_NAME +from litellm.integrations.focus.destinations.base import FocusTimeWindow +from litellm.integrations.focus.focus_logger import FocusLogger + +if TYPE_CHECKING: + from apscheduler.schedulers.asyncio import AsyncIOScheduler +else: + AsyncIOScheduler = Any + + +def _parse_metrics_marker( + marker: Optional[object], +) -> Optional[datetime]: + """Parse metricsMarker from Mavvrik register response into a UTC datetime. + + Handles both formats Mavvrik may return: + - Unix timestamp (int/float): e.g. 1749340800 + - ISO date string: e.g. "2026-06-09" or "2026-06-09T00:00:00Z" + + Returns None for falsy values (0, None, empty string) which indicate + no data has been ingested yet. + """ + if not marker: + return None + try: + if isinstance(marker, (int, float)): + return datetime.fromtimestamp(float(marker), tz=timezone.utc).replace( + hour=0, minute=0, second=0, microsecond=0 + ) + if isinstance(marker, str): + marker = marker.strip() + if not marker: + return None + # Try ISO date first (YYYY-MM-DD), then full ISO datetime + for fmt in ("%Y-%m-%d", "%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%dT%H:%M:%S"): + try: + return datetime.strptime(marker, fmt).replace(tzinfo=timezone.utc) + except ValueError: + continue + except Exception: + pass + verbose_proxy_logger.warning( + "Mavvrik FOCUS: could not parse metricsMarker %r — skipping catch-up", marker + ) + return None + + +class MavvrikFocusLogger(FocusLogger): + """FOCUS-based export logger that routes to the Mavvrik destination.""" + + def __init__(self, **kwargs: Any) -> None: + frequency = os.getenv("MAVVRIK_FOCUS_FREQUENCY", "daily").lower() + if frequency != "daily": + raise ValueError( + f"MAVVRIK_FOCUS_FREQUENCY='{frequency}' is not supported. " + "Only 'daily' is allowed — the Mavvrik ingestion protocol stores one " + "file per calendar date (metrics/YYYY-MM-DD). Hourly or interval " + "exports would overwrite each other within the same day." + ) + super().__init__( + provider="mavvrik", + export_format="csv", + frequency="daily", + prefix="mavvrik_focus_exports", + destination_config={ + "api_key": os.getenv("MAVVRIK_API_KEY"), + "api_endpoint": os.getenv("MAVVRIK_API_ENDPOINT"), + "connection_id": os.getenv("MAVVRIK_CONNECTION_ID"), + }, + **kwargs, + ) + raw = os.getenv("MAVVRIK_FOCUS_MAX_ROWS") + self._max_rows: Optional[int] = int(raw) if raw else 500_000 + + async def _export_window( + self, + *, + window: FocusTimeWindow, + limit: Optional[int], + ) -> None: + """Export with Mavvrik row cap applied when no explicit limit is passed.""" + effective_limit = limit if limit is not None else self._max_rows + engine = self._ensure_engine() + data = await engine._database.get_usage_data( + limit=effective_limit, + start_time_utc=window.start_time, + end_time_utc=window.end_time, + ) + if effective_limit is not None and len(data) >= effective_limit: + verbose_proxy_logger.warning( + "Mavvrik FOCUS export: row cap reached (%d rows). " + "Some data for window %s→%s may be excluded. " + "Increase MAVVRIK_FOCUS_MAX_ROWS to export all rows.", + effective_limit, + window.start_time.date(), + window.end_time.date(), + ) + if data.is_empty(): + verbose_proxy_logger.debug( + "Mavvrik FOCUS export: no usage data for window %s", window + ) + return + normalized = engine._transformer.transform(data) + if normalized.is_empty(): + return + payload = engine._serializer.serialize(normalized) + if not payload: + return + await engine._destination.deliver( + content=payload, + time_window=window, + filename=engine._build_filename(window), + ) + + # Maximum number of days to catch up in a single run. Prevents runaway + # loops if the connector was disabled for a long time, and avoids querying + # data that has likely been cleaned up from LiteLLM_DailyUserSpend. + _MAX_CATCHUP_DAYS = 7 + + async def _run_scheduled_export(self) -> None: + """Export today's window, catching up any dates Mavvrik has not yet received. + + On each run: + 1. Register with Mavvrik → get metricsMarker (last successfully ingested date) + 2. If metricsMarker is behind yesterday, catch up missed dates (capped at + _MAX_CATCHUP_DAYS to avoid runaway loops on long outages) + 3. Export yesterday (today's daily window) + + This ensures a failed export on day N is automatically retried on day N+1 + without any manual intervention. + """ + engine = self._ensure_engine() + from litellm.integrations.focus.destinations.mavvrik_destination import ( # noqa: PLC0415 + FocusMavvrikDestination, + ) + + destination = engine._destination + if not isinstance(destination, FocusMavvrikDestination): + await super()._run_scheduled_export() + return + + # Register and get the last date Mavvrik has processed. + # metricsMarker may be a Unix timestamp (int/float) or an ISO date string. + marker = await destination.get_metrics_marker() + + now = datetime.now(timezone.utc) + yesterday = now.replace(hour=0, minute=0, second=0, microsecond=0) - timedelta( + days=1 + ) + + last_ingested = _parse_metrics_marker(marker) + + # Catch up missed dates, capped at _MAX_CATCHUP_DAYS + if last_ingested and last_ingested < yesterday: + # Never go further back than _MAX_CATCHUP_DAYS from yesterday + earliest_catchup = yesterday - timedelta(days=self._MAX_CATCHUP_DAYS - 1) + catch_up_date = max(last_ingested + timedelta(days=1), earliest_catchup) + + if last_ingested + timedelta(days=1) < earliest_catchup: + verbose_proxy_logger.warning( + "Mavvrik FOCUS export: metricsMarker is more than %d days behind " + "(%s). Catching up from %s only; earlier data will not be re-exported.", + self._MAX_CATCHUP_DAYS, + last_ingested.date(), + catch_up_date.date(), + ) + + while catch_up_date < yesterday: + verbose_proxy_logger.info( + "Mavvrik FOCUS export: catching up missed date %s", + catch_up_date.date(), + ) + window = FocusTimeWindow( + start_time=catch_up_date, + end_time=catch_up_date + timedelta(days=1), + frequency="daily", + ) + await self._export_window(window=window, limit=None) + catch_up_date += timedelta(days=1) + + # Export yesterday's window (the normal daily run) + window = FocusTimeWindow( + start_time=yesterday, + end_time=yesterday + timedelta(days=1), + frequency="daily", + ) + await self._export_window(window=window, limit=None) + + async def initialize_mavvrik_focus_export_job(self) -> None: + """Scheduler entry point — uses Mavvrik-specific pod-lock key.""" + from litellm.proxy.proxy_server import proxy_logging_obj # noqa: PLC0415 + + pod_lock_manager = None + if proxy_logging_obj is not None: + writer = getattr(proxy_logging_obj, "db_spend_update_writer", None) + if writer is not None: + pod_lock_manager = getattr(writer, "pod_lock_manager", None) + + if pod_lock_manager and pod_lock_manager.redis_cache: + acquired = await pod_lock_manager.acquire_lock( + cronjob_id=MAVVRIK_FOCUS_EXPORT_JOB_NAME + ) + if not acquired: + verbose_proxy_logger.debug( + "Mavvrik FOCUS export: unable to acquire pod lock" + ) + return + try: + await self._run_scheduled_export() + finally: + await pod_lock_manager.release_lock( + cronjob_id=MAVVRIK_FOCUS_EXPORT_JOB_NAME + ) + else: + await self._run_scheduled_export() + + @staticmethod + async def init_mavvrik_focus_background_job( + scheduler: AsyncIOScheduler, + ) -> None: + """Register the Mavvrik FOCUS export job on the provided scheduler.""" + loggers: List[MavvrikFocusLogger] = [ + cb + for cb in litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=MavvrikFocusLogger + ) + if type(cb) is MavvrikFocusLogger + ] + if not loggers: + verbose_proxy_logger.debug( + "No MavvrikFocusLogger registered; skipping scheduler" + ) + return + + logger = loggers[0] + trigger_kwargs = logger._build_scheduler_trigger() + scheduler.add_job( # type: ignore[attr-defined] + logger.initialize_mavvrik_focus_export_job, + id=MAVVRIK_FOCUS_EXPORT_JOB_NAME, + replace_existing=True, + **trigger_kwargs, + ) + verbose_proxy_logger.info( + "mavvrik_focus: background export job scheduled (%s)", trigger_kwargs + ) diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index fd402b90d88..ff7f77cf130 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -25,6 +25,7 @@ from litellm.integrations.datadog.datadog_metrics import DatadogMetricsLogger from litellm.integrations.deepeval import DeepEvalLogger from litellm.integrations.dotprompt import DotpromptManager from litellm.integrations.focus.focus_logger import FocusLogger +from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import MavvrikFocusLogger from litellm.integrations.vantage.vantage_logger import VantageLogger from litellm.integrations.galileo import GalileoObserve from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger @@ -102,6 +103,7 @@ class CustomLoggerRegistry: "gitlab": GitLabPromptManager, "cloudzero": CloudZeroLogger, "focus": FocusLogger, + "mavvrik": MavvrikFocusLogger, "vantage": VantageLogger, "posthog": PostHogLogger, } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index dbfcf55d75d..a5bb1c2d56d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4124,6 +4124,17 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 focus_logger = FocusLogger() _in_memory_loggers.append(focus_logger) return focus_logger # type: ignore + elif logging_integration == "mavvrik": + from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import ( + MavvrikFocusLogger, + ) + + for callback in _in_memory_loggers: + if type(callback) is MavvrikFocusLogger: + return callback # type: ignore + mavvrik_focus_logger = MavvrikFocusLogger() + _in_memory_loggers.append(mavvrik_focus_logger) + return mavvrik_focus_logger # type: ignore elif logging_integration == "vantage": from litellm.integrations.vantage.vantage_logger import VantageLogger diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 67f26572c96..cc2f4298474 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7840,6 +7840,15 @@ class ProxyStartupEvent: ) await VantageLogger.init_vantage_background_job(scheduler=scheduler) + ######################################################## + # Mavvrik FOCUS Background Job + ######################################################## + from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import ( # noqa: PLC0415 + MavvrikFocusLogger, + ) + + await MavvrikFocusLogger.init_mavvrik_focus_background_job(scheduler=scheduler) + ######################################################## # Prometheus Background Job ######################################################## @@ -9504,9 +9513,7 @@ async def realtime_websocket_endpoint( if intent == "transcription": route_model = "gpt-realtime-whisper" else: - await websocket.close( - code=1008, reason="model query parameter is required" - ) + await websocket.close(code=1008, reason="model query parameter is required") return assert route_model is not None try: diff --git a/tests/test_litellm/integrations/focus/test_mavvrik_destination.py b/tests/test_litellm/integrations/focus/test_mavvrik_destination.py new file mode 100644 index 00000000000..797238ae238 --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_mavvrik_destination.py @@ -0,0 +1,751 @@ +"""Tests for FocusMavvrikDestination.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.integrations.focus.destinations.base import FocusTimeWindow +from litellm.integrations.focus.destinations.mavvrik_destination import ( + FocusMavvrikDestination, + _validate_api_endpoint, +) + +VALID_ENDPOINT = "https://api.mavvrik.ai/tenant123" + + +def _make_window() -> FocusTimeWindow: + return FocusTimeWindow( + start_time=datetime(2026, 1, 1, 0, 0, 0, tzinfo=timezone.utc), + end_time=datetime(2026, 1, 2, 0, 0, 0, tzinfo=timezone.utc), + frequency="daily", + ) + + +def _dest(**overrides) -> FocusMavvrikDestination: + config = { + "api_key": "test-key", + "api_endpoint": VALID_ENDPOINT, + "connection_id": "conn-123", + } + config.update(overrides) + return FocusMavvrikDestination(prefix="mavvrik_focus_exports", config=config) + + +def test_missing_api_key_raises(): + with pytest.raises(ValueError, match="MAVVRIK_API_KEY"): + FocusMavvrikDestination( + prefix="p", + config={"api_endpoint": VALID_ENDPOINT, "connection_id": "c"}, + ) + + +def test_missing_api_endpoint_raises(): + with pytest.raises(ValueError, match="MAVVRIK_API_ENDPOINT"): + FocusMavvrikDestination( + prefix="p", + config={"api_key": "k", "connection_id": "c"}, + ) + + +def test_missing_connection_id_raises(): + with pytest.raises(ValueError, match="MAVVRIK_CONNECTION_ID"): + FocusMavvrikDestination( + prefix="p", + config={"api_key": "k", "api_endpoint": VALID_ENDPOINT}, + ) + + +def test_non_https_endpoint_raises(): + with pytest.raises(ValueError, match="HTTPS"): + _validate_api_endpoint("http://api.mavvrik.ai/tenant") + + +def test_non_mavvrik_domain_raises(): + with pytest.raises(ValueError, match="Mavvrik domain"): + _validate_api_endpoint("https://evil.com/tenant") + + +def test_valid_mavvrik_domains_accepted(): + for domain in ( + "https://api.mavvrik.ai/tenant", + "https://api.mavvrik.dev/tenant", + "https://api.mavvrik.app/tenant", + ): + _validate_api_endpoint(domain) # must not raise + + +def test_initializes_with_not_registered(): + dest = _dest() + assert dest._registered is False + + +@pytest.mark.asyncio +async def test_deliver_skips_empty_content(): + dest = _dest() + await dest.deliver(content=b"", time_window=_make_window(), filename="usage.csv") + # _registered still False — _ensure_registered was never called + assert dest._registered is False + + +@pytest.mark.asyncio +async def test_large_content_uploads_in_multiple_chunks(): + """Content larger than _GCS_CHUNK_SIZE must be uploaded in multiple chunks. + + GCS assembles intermediate chunks (308) + final chunk (200) into one object. + The destination must send Content-Range headers for each chunk correctly. + """ + from litellm.integrations.focus.destinations.mavvrik_destination import ( + FocusMavvrikDestination, + _GCS_CHUNK_SIZE, + ) + + dest = FocusMavvrikDestination( + prefix="p", + config={"api_key": "k", "api_endpoint": VALID_ENDPOINT, "connection_id": "c"}, + ) + + register_resp = MagicMock() + register_resp.status_code = 200 + + signed_url_resp = MagicMock() + signed_url_resp.status_code = 200 + signed_url_resp.json.return_value = { + "url": "https://storage.googleapis.com/upload?sig=x" + } + + init_resp = MagicMock() + init_resp.status_code = 200 + init_resp.headers = {"Location": "https://storage.googleapis.com/session"} + + # First chunk → 308, second (final) chunk → 200 + chunk1_resp = MagicMock() + chunk1_resp.status_code = 308 + + chunk2_resp = MagicMock() + chunk2_resp.status_code = 200 + + mock_http = MagicMock() + mock_http.client = MagicMock() + mock_http.client.request = AsyncMock( + side_effect=[ + register_resp, + signed_url_resp, + init_resp, + chunk1_resp, + chunk2_resp, + ] + ) + dest._http = mock_http + + # Build content that when gzipped exceeds one chunk. + # Use incompressible random-ish bytes to ensure gzip doesn't shrink it below the chunk size. + import os as _os + + raw = b"col1,col2\n" + _os.urandom(_GCS_CHUNK_SIZE + 1024) + + await dest.deliver( + content=raw, + time_window=_make_window(), + filename="usage.csv", + ) + + # register + get_signed_url + init + 2 chunk PUTs = 5 calls + assert mock_http.client.request.call_count == 5 + + # Check Content-Range headers + put_calls = mock_http.client.request.call_args_list[3:] + assert "bytes" in put_calls[0].kwargs["headers"]["Content-Range"] + assert "/*" in put_calls[0].kwargs["headers"]["Content-Range"] # intermediate + assert "/*" not in put_calls[1].kwargs["headers"]["Content-Range"] # final + + +@pytest.mark.asyncio +async def test_deliver_calls_register_get_url_and_upload(): + dest = _dest() + + register_resp = MagicMock() + register_resp.status_code = 200 + + signed_url_resp = MagicMock() + signed_url_resp.status_code = 200 + signed_url_resp.json.return_value = {"url": "https://storage.googleapis.com/signed"} + + init_resp = MagicMock() + init_resp.status_code = 200 + init_resp.headers = {"Location": "https://storage.googleapis.com/session-uri"} + + upload_resp = MagicMock() + upload_resp.status_code = 200 + + mock_http = MagicMock() + mock_http.client = MagicMock() + # All 4 calls go through self._http.client.request: + # 1. register, 2. get_signed_url, 3. GCS session init POST, 4. GCS PUT + mock_http.client.request = AsyncMock( + side_effect=[register_resp, signed_url_resp, init_resp, upload_resp] + ) + dest._http = mock_http + + await dest.deliver( + content=b"header\nrow1\n", + time_window=_make_window(), + filename="usage.csv", + ) + + assert dest._registered is True + assert mock_http.client.request.call_count == 4 + # Verify Content-Range header was set on the PUT + put_call = mock_http.client.request.call_args_list[3] + assert "Content-Range" in put_call.kwargs["headers"] + + +@pytest.mark.asyncio +async def test_register_called_only_once_across_multiple_deliveries(): + dest = _dest() + + register_resp = MagicMock() + register_resp.status_code = 200 + + def _signed_url_resp(): + r = MagicMock() + r.status_code = 200 + r.json.return_value = {"url": "https://storage.googleapis.com/signed"} + return r + + init_resp = MagicMock() + init_resp.status_code = 200 + init_resp.headers = {"Location": "https://storage.googleapis.com/session-uri"} + + upload_resp = MagicMock() + upload_resp.status_code = 200 + + mock_http = MagicMock() + mock_http.client = MagicMock() + # First delivery: register, get_signed_url, GCS init, GCS PUT + # Second delivery: get_signed_url, GCS init, GCS PUT (register skipped) + mock_http.client.request = AsyncMock( + side_effect=[ + register_resp, + _signed_url_resp(), + init_resp, + upload_resp, + _signed_url_resp(), + init_resp, + upload_resp, + ] + ) + dest._http = mock_http + + window = _make_window() + await dest.deliver(content=b"header\nrow1\n", time_window=window, filename="1.csv") + await dest.deliver(content=b"header\nrow2\n", time_window=window, filename="2.csv") + + # 7 total: register(1) + [get_url+init+put](2) × 2 deliveries + assert mock_http.client.request.call_count == 7 + # First call was register + first_call = mock_http.client.request.call_args_list[0] + assert first_call.kwargs["method"] == "POST" + assert "/upload-url" not in first_call.kwargs["url"] + + +@pytest.mark.asyncio +async def test_deliver_raises_on_register_failure(): + dest = _dest() + + fail_resp = MagicMock() + fail_resp.status_code = 403 + fail_resp.text = "Forbidden" + + mock_http = MagicMock() + mock_http.client = MagicMock() + mock_http.client.request = AsyncMock(return_value=fail_resp) + dest._http = mock_http + + with pytest.raises(RuntimeError, match="register failed"): + await dest.deliver( + content=b"data", + time_window=_make_window(), + filename="usage.csv", + ) + + +@pytest.mark.asyncio +async def test_deliver_raises_on_signed_url_api_error(): + """_get_signed_url must raise RuntimeError when the API returns a 4xx.""" + dest = _dest() + + register_resp = MagicMock() + register_resp.status_code = 200 + + fail_resp = MagicMock() + fail_resp.status_code = 500 + fail_resp.text = "Internal Server Error" + + mock_http = MagicMock() + mock_http.client = MagicMock() + mock_http.client.request = AsyncMock(side_effect=[register_resp, fail_resp]) + dest._http = mock_http + + with pytest.raises(RuntimeError, match="failed to get signed URL"): + await dest.deliver( + content=b"data", + time_window=_make_window(), + filename="usage.csv", + ) + + +@pytest.mark.asyncio +async def test_deliver_raises_on_missing_signed_url(): + dest = _dest() + + register_resp = MagicMock() + register_resp.status_code = 200 + + bad_url_resp = MagicMock() + bad_url_resp.status_code = 200 + bad_url_resp.json.return_value = {} # no 'url' field + + mock_http = MagicMock() + mock_http.client = MagicMock() + mock_http.client.request = AsyncMock(side_effect=[register_resp, bad_url_resp]) + dest._http = mock_http + + with pytest.raises(RuntimeError, match="missing 'url' field"): + await dest.deliver( + content=b"data", + time_window=_make_window(), + filename="usage.csv", + ) + + +@pytest.mark.asyncio +async def test_deliver_raises_on_non_gcs_signed_url(): + """Signed URL pointing to a non-GCS host must be rejected before any upload.""" + from litellm.integrations.focus.destinations.mavvrik_destination import ( + _validate_gcs_url, + ) + + with pytest.raises(ValueError, match="GCS endpoint"): + _validate_gcs_url("https://evil.com/upload?token=abc", "signed URL") + + +@pytest.mark.asyncio +async def test_deliver_raises_on_non_gcs_session_uri(): + """Session URI from Location header pointing to a non-GCS host must be rejected.""" + dest = _dest() + + register_resp = MagicMock() + register_resp.status_code = 200 + + signed_url_resp = MagicMock() + signed_url_resp.status_code = 200 + # signed URL is valid GCS + signed_url_resp.json.return_value = { + "url": "https://storage.googleapis.com/upload?sig=abc" + } + + # Location header points to a non-GCS host + init_resp = MagicMock() + init_resp.status_code = 200 + init_resp.headers = {"Location": "https://evil.com/session-uri"} + + mock_http = MagicMock() + mock_http.client = MagicMock() + # register, get_signed_url, GCS session init (returns bad Location) + mock_http.client.request = AsyncMock( + side_effect=[register_resp, signed_url_resp, init_resp] + ) + dest._http = mock_http + + with pytest.raises(ValueError, match="GCS endpoint"): + await dest.deliver( + content=b"data", + time_window=_make_window(), + filename="usage.csv", + ) + + +def test_factory_creates_mavvrik_destination(monkeypatch): + monkeypatch.setenv("MAVVRIK_API_KEY", "k") + monkeypatch.setenv("MAVVRIK_API_ENDPOINT", VALID_ENDPOINT) + monkeypatch.setenv("MAVVRIK_CONNECTION_ID", "c") + + from litellm.integrations.focus.destinations.factory import FocusDestinationFactory + + dest = FocusDestinationFactory.create(provider="mavvrik", prefix="p") + + assert isinstance(dest, FocusMavvrikDestination) + assert dest.api_key == "k" + assert dest.connection_id == "c" + + +def test_only_daily_frequency_is_supported(): + """MavvrikFocusLogger must raise ValueError for non-daily frequencies.""" + import importlib + + for freq in ("hourly", "interval"): + + def _make(f=freq, monkeypatch=None): + import os + + old = os.environ.get("MAVVRIK_FOCUS_FREQUENCY") + os.environ["MAVVRIK_FOCUS_FREQUENCY"] = f + try: + from litellm.integrations.mavvrik_focus import mavvrik_focus_logger + + importlib.reload(mavvrik_focus_logger) + with pytest.raises(ValueError, match="Only 'daily' is allowed"): + mavvrik_focus_logger.MavvrikFocusLogger() + finally: + if old is None: + os.environ.pop("MAVVRIK_FOCUS_FREQUENCY", None) + else: + os.environ["MAVVRIK_FOCUS_FREQUENCY"] = old + + _make() + + +def test_max_rows_defaults_to_500k(): + """MAVVRIK_FOCUS_MAX_ROWS defaults to 500_000 when not set.""" + from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import ( + MavvrikFocusLogger, + ) + + logger = MavvrikFocusLogger() + assert logger._max_rows == 500_000 + + +def test_max_rows_reads_from_env(monkeypatch): + """MAVVRIK_FOCUS_MAX_ROWS env var is respected.""" + monkeypatch.setenv("MAVVRIK_FOCUS_MAX_ROWS", "100000") + + from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import ( + MavvrikFocusLogger, + ) + + logger = MavvrikFocusLogger() + assert logger._max_rows == 100_000 + + +@pytest.mark.asyncio +async def test_export_window_passes_max_rows_as_limit(monkeypatch): + """_export_window must pass _max_rows as limit to get_usage_data.""" + monkeypatch.setenv("MAVVRIK_FOCUS_MAX_ROWS", "1000") + + import polars as pl + from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import ( + MavvrikFocusLogger, + ) + from litellm.integrations.focus.destinations.base import FocusTimeWindow + from datetime import datetime, timezone + + logger = MavvrikFocusLogger() + assert logger._max_rows == 1000 + + # Mock the engine internals so _export_window runs through our new code path + db_mock = MagicMock() + db_mock.get_usage_data = AsyncMock(return_value=pl.DataFrame()) # empty → no upload + + engine_mock = MagicMock() + engine_mock._database = db_mock + logger._engine = engine_mock + + window = FocusTimeWindow( + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + end_time=datetime(2026, 1, 2, tzinfo=timezone.utc), + frequency="daily", + ) + await logger._export_window(window=window, limit=None) + + db_mock.get_usage_data.assert_called_once_with( + limit=1000, + start_time_utc=window.start_time, + end_time_utc=window.end_time, + ) + + +@pytest.mark.asyncio +async def test_run_scheduled_export_catches_up_missed_dates(): + """If metricsMarker is 2 days behind, _run_scheduled_export exports missed dates first.""" + import polars as pl + from datetime import datetime, timedelta, timezone + from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import ( + MavvrikFocusLogger, + ) + from litellm.integrations.focus.destinations.mavvrik_destination import ( + FocusMavvrikDestination, + ) + + logger = MavvrikFocusLogger() + + # metricsMarker = 3 days ago → 2 missed dates (day-2 and day-1) + today's run + now = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) + yesterday = now - timedelta(days=1) + two_days_ago = now - timedelta(days=2) + three_days_ago = now - timedelta(days=3) + + marker_ts = int(three_days_ago.timestamp()) + + # Mock destination + dest_mock = MagicMock(spec=FocusMavvrikDestination) + dest_mock.get_metrics_marker = AsyncMock(return_value=marker_ts) + + # Mock engine + db_mock = MagicMock() + db_mock.get_usage_data = AsyncMock(return_value=pl.DataFrame()) + engine_mock = MagicMock() + engine_mock._database = db_mock + engine_mock._destination = dest_mock + logger._engine = engine_mock + + await logger._run_scheduled_export() + + # Should have queried DB 3 times: day-2, day-1 (yesterday), and the normal yesterday window + # Actually: catch-up covers [three_days_ago+1 .. yesterday) = [two_days_ago, yesterday) + # = two_days_ago only (1 missed), then normal yesterday = 2 total calls + calls = db_mock.get_usage_data.call_args_list + assert len(calls) == 2 + # First call is the catch-up (two_days_ago) + assert calls[0].kwargs["start_time_utc"].date() == two_days_ago.date() + # Second call is yesterday's normal daily run + assert calls[1].kwargs["start_time_utc"].date() == yesterday.date() + + +@pytest.mark.asyncio +async def test_run_scheduled_export_no_catchup_when_marker_is_current(): + """If metricsMarker = yesterday, no catch-up needed — just export yesterday.""" + import polars as pl + from datetime import datetime, timedelta, timezone + from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import ( + MavvrikFocusLogger, + ) + from litellm.integrations.focus.destinations.mavvrik_destination import ( + FocusMavvrikDestination, + ) + + logger = MavvrikFocusLogger() + + now = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) + yesterday = now - timedelta(days=1) + marker_ts = int(yesterday.timestamp()) + + dest_mock = MagicMock(spec=FocusMavvrikDestination) + dest_mock.get_metrics_marker = AsyncMock(return_value=marker_ts) + + db_mock = MagicMock() + db_mock.get_usage_data = AsyncMock(return_value=pl.DataFrame()) + engine_mock = MagicMock() + engine_mock._database = db_mock + engine_mock._destination = dest_mock + logger._engine = engine_mock + + await logger._run_scheduled_export() + + # Only one call — yesterday's normal run, no catch-up + assert db_mock.get_usage_data.call_count == 1 + assert ( + db_mock.get_usage_data.call_args.kwargs["start_time_utc"].date() + == yesterday.date() + ) + + +@pytest.mark.asyncio +async def test_metrics_marker_always_calls_api(): + """get_metrics_marker must call the register API every time to get a fresh marker. + + This is the key difference from deliver() — catch-up requires the current + metricsMarker on every scheduled run, not just the first one. + """ + dest = _dest() + + register_resp = MagicMock() + register_resp.status_code = 200 + register_resp.json.return_value = { + "id": "litellm-conn-123", + "metricsMarker": 1749340800, + } + + mock_http = MagicMock() + mock_http.client = MagicMock() + mock_http.client.request = AsyncMock(return_value=register_resp) + dest._http = mock_http + + # First call + marker = await dest.get_metrics_marker() + assert marker == 1749340800 + assert dest._registered is True + + # Second call — must call API again to get fresh marker (not return None) + marker2 = await dest.get_metrics_marker() + assert marker2 == 1749340800 + assert mock_http.client.request.call_count == 2 # API called both times + + +def test_parse_metrics_marker_handles_unix_timestamp(): + from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import ( + _parse_metrics_marker, + ) + from datetime import datetime, timezone + + # Use a known date and compute its timestamp to avoid hardcoding + known_date = datetime(2026, 6, 9, 0, 0, 0, tzinfo=timezone.utc) + ts = int(known_date.timestamp()) + + result = _parse_metrics_marker(ts) + assert result is not None + assert result.date().isoformat() == "2026-06-09" + assert result.tzinfo == timezone.utc + + +def test_parse_metrics_marker_handles_iso_date_string(): + from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import ( + _parse_metrics_marker, + ) + + result = _parse_metrics_marker("2026-06-09") + assert result is not None + assert result.date().isoformat() == "2026-06-09" + + +def test_parse_metrics_marker_handles_iso_datetime_string(): + from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import ( + _parse_metrics_marker, + ) + + result = _parse_metrics_marker("2026-06-09T00:00:00Z") + assert result is not None + assert result.date().isoformat() == "2026-06-09" + + +def test_parse_metrics_marker_returns_none_for_zero(): + from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import ( + _parse_metrics_marker, + ) + + assert _parse_metrics_marker(0) is None + assert _parse_metrics_marker(None) is None + assert _parse_metrics_marker("") is None + + +def test_parse_metrics_marker_returns_none_for_garbage(): + from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import ( + _parse_metrics_marker, + ) + + # Should not raise — logs warning and returns None + assert _parse_metrics_marker("not-a-date") is None + + +@pytest.mark.asyncio +async def test_catchup_capped_at_max_catchup_days(): + """Catch-up must not go further back than _MAX_CATCHUP_DAYS.""" + import polars as pl + from datetime import datetime, timedelta, timezone + from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import ( + MavvrikFocusLogger, + ) + from litellm.integrations.focus.destinations.mavvrik_destination import ( + FocusMavvrikDestination, + ) + + logger = MavvrikFocusLogger() + max_days = MavvrikFocusLogger._MAX_CATCHUP_DAYS + + now = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) + yesterday = now - timedelta(days=1) + # Marker is 30 days ago — well beyond the cap + thirty_days_ago = now - timedelta(days=30) + marker_ts = int(thirty_days_ago.timestamp()) + + dest_mock = MagicMock(spec=FocusMavvrikDestination) + dest_mock.get_metrics_marker = AsyncMock(return_value=marker_ts) + + db_mock = MagicMock() + db_mock.get_usage_data = AsyncMock(return_value=pl.DataFrame()) + engine_mock = MagicMock() + engine_mock._database = db_mock + engine_mock._destination = dest_mock + logger._engine = engine_mock + + await logger._run_scheduled_export() + + # Should have queried at most _MAX_CATCHUP_DAYS times + # (max_days - 1 catch-up dates + 1 yesterday = max_days total) + assert db_mock.get_usage_data.call_count <= max_days + + # First catch-up date must not be earlier than (yesterday - max_days + 1) + earliest_allowed = yesterday - timedelta(days=max_days - 1) + first_call_start = db_mock.get_usage_data.call_args_list[0].kwargs["start_time_utc"] + assert first_call_start.date() >= earliest_allowed.date() + + +@pytest.mark.asyncio +async def test_register_resets_on_410(): + """_registered flag must be False after a 410 so next run re-registers.""" + dest = _dest() + + resp_410 = MagicMock() + resp_410.status_code = 410 + resp_410.text = "Gone" + + mock_http = MagicMock() + mock_http.client = MagicMock() + mock_http.client.request = AsyncMock(return_value=resp_410) + dest._http = mock_http + dest._registered = False # not yet registered — trigger the call + + with pytest.raises(RuntimeError, match="disconnected"): + await dest._ensure_registered() + + assert dest._registered is False + + +@pytest.mark.asyncio +async def test_gcs_session_cancelled_on_chunk_failure(): + """GCS session must be cancelled (DELETE) when a chunk PUT fails.""" + dest = _dest() + + register_resp = MagicMock() + register_resp.status_code = 200 + + signed_url_resp = MagicMock() + signed_url_resp.status_code = 200 + signed_url_resp.json.return_value = { + "url": "https://storage.googleapis.com/upload?sig=x" + } + + init_resp = MagicMock() + init_resp.status_code = 200 + init_resp.headers = {"Location": "https://storage.googleapis.com/session"} + + # Chunk PUT fails with 500 + fail_resp = MagicMock() + fail_resp.status_code = 500 + fail_resp.text = "Internal Server Error" + + # DELETE (session cancel) + delete_resp = MagicMock() + delete_resp.status_code = 200 + + mock_http = MagicMock() + mock_http.client = MagicMock() + mock_http.client.request = AsyncMock( + side_effect=[register_resp, signed_url_resp, init_resp, fail_resp, delete_resp] + ) + dest._http = mock_http + + with pytest.raises(RuntimeError, match="GCS chunk upload failed"): + await dest.deliver( + content=b"header\nrow1\n", + time_window=_make_window(), + filename="usage.csv", + ) + + # Verify DELETE was called to cancel the session + calls = mock_http.client.request.call_args_list + delete_call = calls[4] + assert delete_call.kwargs["method"] == "DELETE" + assert "storage.googleapis.com/session" in delete_call.kwargs["url"]