fix(mavvrik): advance metricsMarker after upload + fix scheduler startup

Two bugs fixed:

1. deliver() never called PATCH /metrics/agent/ai/{connectionId} after a
   successful GCS upload, so metricsMarker stayed at 0 and every daily run
   re-exported the same dates in an infinite catch-up loop.
   Fix: add _update_metrics_marker(date_epoch) called at the end of deliver()
   after _upload_to_gcs() succeeds. A 4xx warns but does not raise (the GCS
   file is already committed). A 410 raises consistent with the rest of the
   destination.

2. init_mavvrik_focus_background_job runs at proxy startup before any LLM call
   has triggered lazy instantiation of MavvrikFocusLogger, so it found no
   logger instance and silently skipped registering the daily export job.
   Fix: if no instance is found but "mavvrik" is in litellm.callbacks, call
   _init_custom_logger_compatible_class to force instantiation before
   the APScheduler job is registered.
This commit is contained in:
Praveen Ghuge 2026-06-23 14:27:28 +05:30
parent 69b0dd2da0
commit b2c5d069bc
3 changed files with 116 additions and 19 deletions

View file

@ -3,6 +3,7 @@
Flow:
1. GET /metrics/agent/ai/{connection_id}/upload-url → GCS signed URL
2. PUT <signed_url> 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

View file

@ -253,6 +253,19 @@ 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"
)
if isinstance(instance, MavvrikFocusLogger):
loggers = [instance]
if not loggers:
verbose_proxy_logger.debug(
"No MavvrikFocusLogger registered; skipping scheduler"

View file

@ -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_4xx():
"""_update_metrics_marker must log a warning on 4xx 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