diff --git a/litellm/integrations/focus/destinations/mavvrik_destination.py b/litellm/integrations/focus/destinations/mavvrik_destination.py index 1e3c98b9a70..659f608a3e1 100644 --- a/litellm/integrations/focus/destinations/mavvrik_destination.py +++ b/litellm/integrations/focus/destinations/mavvrik_destination.py @@ -3,6 +3,7 @@ Flow: 1. GET /metrics/agent/ai/{connection_id}/upload-url → GCS signed URL 2. PUT with CSV content + 3. PATCH /metrics/agent/ai/{connection_id} → advance metricsMarker """ from __future__ import annotations @@ -127,8 +128,6 @@ class FocusMavvrikDestination(FocusDestination): 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). " @@ -273,14 +272,35 @@ class FocusMavvrikDestination(FocusDestination): pass raise + async def _update_metrics_marker(self, date_epoch: int) -> None: + """PATCH agent endpoint to advance metricsMarker after a successful upload.""" + resp = await self._http.client.request( + method="PATCH", + url=self._agent_url, + headers=self._auth_headers, + json={"metricsMarker": date_epoch}, + 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: + verbose_logger.warning( + "Mavvrik FOCUS destination: failed to update metricsMarker (%s): %s", + resp.status_code, + resp.text[:200], + ) + return + verbose_logger.debug( + "Mavvrik FOCUS destination: metricsMarker advanced to %s", date_epoch + ) + 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. @@ -328,6 +348,7 @@ class FocusMavvrikDestination(FocusDestination): return date_str = time_window.start_time.strftime("%Y-%m-%d") + date_epoch = int(time_window.start_time.timestamp()) verbose_logger.debug( "Mavvrik FOCUS destination: uploading %d bytes for date=%s (%s)", @@ -339,6 +360,7 @@ class FocusMavvrikDestination(FocusDestination): await self._ensure_registered() signed_url = await self._get_signed_url(date_str) await self._upload_to_gcs(signed_url, content) + await self._update_metrics_marker(date_epoch) verbose_logger.debug( "Mavvrik FOCUS destination: upload complete for date=%s", date_str diff --git a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py index 47d3e1da7bc..bbc9d1a6330 100644 --- a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py +++ b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py @@ -149,8 +149,8 @@ class MavvrikFocusLogger(FocusLogger): 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) + 2. If metricsMarker is behind yesterday (or 0/None for a fresh connector), + catch up missed dates (capped at _MAX_CATCHUP_DAYS) 3. Export yesterday (today's daily window) This ensures a failed export on day N is automatically retried on day N+1 @@ -177,13 +177,21 @@ class MavvrikFocusLogger(FocusLogger): 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) + # Catch up missed dates, capped at _MAX_CATCHUP_DAYS. + # last_ingested=None means metricsMarker=0 (fresh connector, never ingested) — + # treat the same as being _MAX_CATCHUP_DAYS behind so we export all available history. + earliest_catchup = yesterday - timedelta(days=self._MAX_CATCHUP_DAYS - 1) + if last_ingested is None or last_ingested < yesterday: + catch_up_date = ( + earliest_catchup + if last_ingested is None + else max(last_ingested + timedelta(days=1), earliest_catchup) + ) - if last_ingested + timedelta(days=1) < earliest_catchup: + if ( + last_ingested is not None + and 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.", @@ -197,18 +205,24 @@ class MavvrikFocusLogger(FocusLogger): "Mavvrik FOCUS export: catching up missed date %s", catch_up_date.date(), ) + # Use now as end_time for catch-up windows too — rows for old dates + # may have been flushed to DB well after their calendar day ended. + catch_up_end = min(catch_up_date + timedelta(days=1), now) window = FocusTimeWindow( start_time=catch_up_date, - end_time=catch_up_date + timedelta(days=1), + end_time=catch_up_end, frequency="daily", ) await self._export_window(window=window, limit=None) catch_up_date += timedelta(days=1) - # Export yesterday's window (the normal daily run) + # Export yesterday's window (the normal daily run). + # Use `now` as end_time so spend rows flushed after midnight are included. + # LiteLLM's DailyUserSpend rows for a given date keep getting updated_at + # bumped as the flush job runs; capping at midnight would miss those updates. window = FocusTimeWindow( start_time=yesterday, - end_time=yesterday + timedelta(days=1), + end_time=now, frequency="daily", ) await self._export_window(window=window, limit=None) @@ -253,6 +267,21 @@ class MavvrikFocusLogger(FocusLogger): ) if type(cb) is MavvrikFocusLogger ] + if not loggers and "mavvrik" in litellm.callbacks: + # The logger is registered as the string "mavvrik" but hasn't been + # instantiated yet (lazy init happens on first LLM call). Force it now + # so the scheduler can register the daily export job at startup. + from litellm.litellm_core_utils.litellm_logging import ( # noqa: PLC0415 + _init_custom_logger_compatible_class, + ) + + instance = _init_custom_logger_compatible_class( + logging_integration="mavvrik", + internal_usage_cache=None, + llm_router=None, + ) + if isinstance(instance, MavvrikFocusLogger): + loggers = [instance] if not loggers: verbose_proxy_logger.debug( "No MavvrikFocusLogger registered; skipping scheduler" diff --git a/tests/test_litellm/integrations/focus/test_mavvrik_destination.py b/tests/test_litellm/integrations/focus/test_mavvrik_destination.py index 797238ae238..3a23dc4ffb2 100644 --- a/tests/test_litellm/integrations/focus/test_mavvrik_destination.py +++ b/tests/test_litellm/integrations/focus/test_mavvrik_destination.py @@ -34,6 +34,12 @@ def _dest(**overrides) -> FocusMavvrikDestination: return FocusMavvrikDestination(prefix="mavvrik_focus_exports", config=config) +def _patch_resp(status: int = 204) -> MagicMock: + r = MagicMock() + r.status_code = status + return r + + def test_missing_api_key_raises(): with pytest.raises(ValueError, match="MAVVRIK_API_KEY"): FocusMavvrikDestination( @@ -127,6 +133,8 @@ async def test_large_content_uploads_in_multiple_chunks(): chunk2_resp = MagicMock() chunk2_resp.status_code = 200 + patch_resp = _patch_resp(204) + mock_http = MagicMock() mock_http.client = MagicMock() mock_http.client.request = AsyncMock( @@ -136,6 +144,7 @@ async def test_large_content_uploads_in_multiple_chunks(): init_resp, chunk1_resp, chunk2_resp, + patch_resp, ] ) dest._http = mock_http @@ -152,15 +161,20 @@ async def test_large_content_uploads_in_multiple_chunks(): filename="usage.csv", ) - # register + get_signed_url + init + 2 chunk PUTs = 5 calls - assert mock_http.client.request.call_count == 5 + # register + get_signed_url + init + 2 chunk PUTs + PATCH = 6 calls + assert mock_http.client.request.call_count == 6 - # Check Content-Range headers - put_calls = mock_http.client.request.call_args_list[3:] + # Check Content-Range headers on the chunk PUTs (calls 3 and 4) + put_calls = mock_http.client.request.call_args_list[3:5] 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 + # Verify the PATCH call advanced metricsMarker + patch_call = mock_http.client.request.call_args_list[5] + assert patch_call.kwargs["method"] == "PATCH" + assert "metricsMarker" in patch_call.kwargs["json"] + @pytest.mark.asyncio async def test_deliver_calls_register_get_url_and_upload(): @@ -180,12 +194,14 @@ async def test_deliver_calls_register_get_url_and_upload(): upload_resp = MagicMock() upload_resp.status_code = 200 + patch_resp = _patch_resp(204) + 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 + # All 5 calls go through self._http.client.request: + # 1. register, 2. get_signed_url, 3. GCS session init POST, 4. GCS PUT, 5. PATCH marker mock_http.client.request = AsyncMock( - side_effect=[register_resp, signed_url_resp, init_resp, upload_resp] + side_effect=[register_resp, signed_url_resp, init_resp, upload_resp, patch_resp] ) dest._http = mock_http @@ -196,10 +212,14 @@ async def test_deliver_calls_register_get_url_and_upload(): ) assert dest._registered is True - assert mock_http.client.request.call_count == 4 + assert mock_http.client.request.call_count == 5 # 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"] + # Verify PATCH was called last with metricsMarker + patch_call = mock_http.client.request.call_args_list[4] + assert patch_call.kwargs["method"] == "PATCH" + assert "metricsMarker" in patch_call.kwargs["json"] @pytest.mark.asyncio @@ -224,17 +244,19 @@ async def test_register_called_only_once_across_multiple_deliveries(): 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) + # First delivery: register, get_signed_url, GCS init, GCS PUT, PATCH + # Second delivery: get_signed_url, GCS init, GCS PUT, PATCH (register skipped) mock_http.client.request = AsyncMock( side_effect=[ register_resp, _signed_url_resp(), init_resp, upload_resp, + _patch_resp(204), _signed_url_resp(), init_resp, upload_resp, + _patch_resp(204), ] ) dest._http = mock_http @@ -243,8 +265,8 @@ async def test_register_called_only_once_across_multiple_deliveries(): 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 + # 9 total: register(1) + [get_url+init+put+patch](4) × 2 deliveries + assert mock_http.client.request.call_count == 9 # First call was register first_call = mock_http.client.request.call_args_list[0] assert first_call.kwargs["method"] == "POST" @@ -749,3 +771,43 @@ async def test_gcs_session_cancelled_on_chunk_failure(): delete_call = calls[4] assert delete_call.kwargs["method"] == "DELETE" assert "storage.googleapis.com/session" in delete_call.kwargs["url"] + + +@pytest.mark.asyncio +async def test_update_metrics_marker_warns_on_non_410_error(): + """_update_metrics_marker must log a warning on any >=400 (non-410) status but not raise.""" + dest = _dest() + + 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(return_value=fail_resp) + dest._http = mock_http + + # Must not raise — warning only + await dest._update_metrics_marker(1234567890) + assert mock_http.client.request.call_count == 1 + + +@pytest.mark.asyncio +async def test_update_metrics_marker_raises_on_410(): + """_update_metrics_marker must raise RuntimeError and reset _registered on 410.""" + dest = _dest() + dest._registered = True + + 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 + + with pytest.raises(RuntimeError, match="disconnected"): + await dest._update_metrics_marker(1234567890) + + assert dest._registered is False