fix(proxy): judge /key/health GCS status by the health check's own event

The flush result now carries the ids it uploaded and the ids it failed, so
/key/health reports unhealthy only when its own event failed to upload.
Another team's broken bucket or a leftover retry in the shared queue no
longer marks a working key unhealthy, and the flush walks every pending
batch so an event queued behind more than batch_size items is still judged

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/integrations/gcs_bucket/gcs_bucket.py

@@ -242,10 +242,14 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
         Send each log individually as separate GCS objects (legacy behavior).
         This is used when GCS_USE_BATCHED_LOGGING is disabled.
         """
-        failed_items: Final = tuple([item for item in items if not await self._send_single_log_item(item)])
+        outcomes: Final = tuple([(item, await self._send_single_log_item(item)) for item in items])
+        failed_items: Final = tuple(item for item, sent in outcomes if not sent)
         if failed_items:
             self._requeue(failed_items)
-        return GCSFlushResult(sent=len(items) - len(failed_items), failed=len(failed_items))
+        return GCSFlushResult(
+            sent_ids=tuple(item["payload"]["id"] for item, sent in outcomes if sent),
+            failed_ids=tuple(item["payload"]["id"] for item in failed_items),
+        )
 
     async def _send_single_log_item(self, item: GCSLogQueueItem) -> bool:
         """
@@ -288,8 +292,12 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
             if group_failed:
                 self._requeue(group_items)
         return GCSFlushResult(
-            sent=sum(group_sent for _, (group_sent, _) in results),
-            failed=sum(group_failed for _, (_, group_failed) in results),
+            sent_ids=tuple(
+                item["payload"]["id"] for group_items, (_, failed) in results if not failed for item in group_items
+            ),
+            failed_ids=tuple(
+                item["payload"]["id"] for group_items, (_, failed) in results if failed for item in group_items
+            ),
         )
 
     async def async_send_batch(self) -> None:
@@ -307,7 +315,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
         items_to_process: Final = self._drain_queue_batch()
 
         if not items_to_process:
-            return GCSFlushResult(sent=0, failed=0)
+            return GCSFlushResult(sent_ids=(), failed_ids=())
 
         if self.use_batched_logging:
             return await self._send_grouped_batches(items_to_process)

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -157,6 +157,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import (
 from litellm.types.router import Deployment
 from litellm.types.utils import (
     BudgetConfig,
+    ModelResponse,
     PersonalUIKeyGenerationConfig,
     TeamUIKeyGenerationConfig,
 )
@@ -7140,15 +7141,20 @@ def _callback_entry_error(entry: Mapping[str, object]) -> str | None:
     return None
 
 
-async def flush_gcs_and_describe_failures(gcs_logger: CustomLogger | None) -> str | None:
+async def flush_gcs_and_describe_failures(gcs_logger: CustomLogger | None, health_check_event_id: str) -> str | None:
     from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
 
     if not isinstance(gcs_logger, GCSBucketLogger):
         return "gcs_bucket callback was selected but no GCS logger was initialized"
-    flush_result: Final = await gcs_logger.flush_queue_and_report()
-    if flush_result.failed == 0:
-        return None
-    return f"GCS upload failed for {flush_result.failed} event(s), {flush_result.sent} uploaded"
+    flush_rounds: Final = max(1, math.ceil(gcs_logger.log_queue.qsize() / gcs_logger.batch_size))
+    for _ in range(flush_rounds):
+        flush_result = await gcs_logger.flush_queue_and_report()
+        if health_check_event_id in flush_result.failed_ids:
+            return (
+                f"GCS upload failed for the /key/health event and {flush_result.failed - 1} other event(s), "
+                f"{flush_result.sent} uploaded"
+            )
+    return None
 
 
 async def test_key_logging(
@@ -7194,7 +7200,7 @@ async def test_key_logging(
             request=request,
         )
         data["mock_response"] = "test response"
-        await litellm.acompletion(**data)  # make mock completion call to trigger key based callbacks
+        health_check_response: Final = await litellm.acompletion(**data)
     except Exception as e:
         return LoggingCallbackStatus(
             callbacks=logging_callbacks,
@@ -7203,20 +7209,26 @@ async def test_key_logging(
         )
 
     await asyncio.sleep(2)  # wait for callbacks to run, callbacks use batching so wait for the flush event
+    callback_log_contents: Final = log_capture_string.getvalue()
 
+    health_check_event_id: Final = health_check_response.id if isinstance(health_check_response, ModelResponse) else ""
     gcs_failure: Final = (
-        await flush_gcs_and_describe_failures(get_custom_logger_compatible_class("gcs_bucket"))
+        await flush_gcs_and_describe_failures(get_custom_logger_compatible_class("gcs_bucket"), health_check_event_id)
         if "gcs_bucket" in logging_callbacks
         else None
     )
-
-    log_contents: Final = log_capture_string.getvalue()
     logger.removeHandler(ch)
-    if gcs_failure is not None or log_contents:
+    flush_log_contents: Final = (
+        log_capture_string.getvalue()[len(callback_log_contents) :] if gcs_failure is not None else ""
+    )
+    if gcs_failure is not None or callback_log_contents:
         return LoggingCallbackStatus(
             callbacks=logging_callbacks,
             status="unhealthy",
-            details=f"Logger exceptions triggered, system is unhealthy: {gcs_failure or ''} {log_contents}".strip(),
+            details=(
+                "Logger exceptions triggered, system is unhealthy: "
+                f"{gcs_failure or ''} {callback_log_contents}{flush_log_contents}"
+            ).strip(),
         )
     return LoggingCallbackStatus(
         callbacks=logging_callbacks,

litellm/types/integrations/gcs_bucket.py

@@ -38,5 +38,13 @@ class GCSLogQueueItem(TypedDict):
 
 @dataclass(frozen=True, slots=True)
 class GCSFlushResult:
-    sent: int
-    failed: int
+    sent_ids: tuple[str, ...]
+    failed_ids: tuple[str, ...]
+
+    @property
+    def sent(self) -> int:
+        return len(self.sent_ids)
+
+    @property
+    def failed(self) -> int:
+        return len(self.failed_ids)

tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket.py

@@ -61,14 +61,14 @@ async def test_failed_batch_stays_queued_and_is_retried_on_the_next_flush():
 
     failed_flush = await logger.flush_queue_and_report()
 
-    assert failed_flush == GCSFlushResult(sent=0, failed=2)
+    assert failed_flush == GCSFlushResult(sent_ids=(), failed_ids=("req-1", "req-2"))
     assert logger.log_queue.qsize() == 2
     assert logger.uploaded == []
 
     logger.failing_ids = frozenset()
     retried_flush = await logger.flush_queue_and_report()
 
-    assert retried_flush == GCSFlushResult(sent=2, failed=0)
+    assert retried_flush == GCSFlushResult(sent_ids=("req-1", "req-2"), failed_ids=())
     assert logger.log_queue.qsize() == 0
     assert logger.uploaded == [["req-1", "req-2"]]
 
@@ -83,7 +83,7 @@ async def test_individual_mode_requeues_only_the_failed_items():
 
     result = await logger.flush_queue_and_report()
 
-    assert result == GCSFlushResult(sent=1, failed=1)
+    assert result == GCSFlushResult(sent_ids=("req-ok",), failed_ids=("req-fail",))
     assert logger.uploaded == [["req-ok"]]
     assert logger.queued_ids() == ["req-fail"]
 
@@ -110,7 +110,7 @@ async def test_failed_batch_is_dropped_when_new_events_filled_the_queue_during_t
 
     result = await logger.flush_queue_and_report()
 
-    assert result == GCSFlushResult(sent=0, failed=2)
+    assert result == GCSFlushResult(sent_ids=(), failed_ids=("req-1", "req-2"))
     assert logger.queued_ids() == ["req-3", "req-4"]
 
 
@@ -118,4 +118,4 @@ async def test_failed_batch_is_dropped_when_new_events_filled_the_queue_during_t
 async def test_empty_queue_flush_reports_nothing_sent_or_failed():
     logger = _FakeUploadGCSLogger()
 
-    assert await logger.flush_queue_and_report() == GCSFlushResult(sent=0, failed=0)
+    assert await logger.flush_queue_and_report() == GCSFlushResult(sent_ids=(), failed_ids=())

tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py

@@ -18415,33 +18415,72 @@ async def test_key_health_rejects_key_logging_entries_without_a_callback_name():
     assert "callback_name is required" in exc.value.message
 
 
-def _gcs_logger_whose_flush_reports(sent: int, failed: int):
+def _fake_upload_gcs_logger(broken_bucket: str | None = None, batch_size: int = 2048, enqueue_error: str | None = None):
     from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
-    from litellm.types.integrations.gcs_bucket import GCSFlushResult
+    from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
+    from litellm.types.integrations.gcs_bucket import GCSLoggingConfig, GCSLogQueueItem
+    from litellm.types.utils import StandardLoggingPayload
+
+    class _FakeUploadGCSLogger(GCSBucketLogger):
+        """Skips GCP auth; an upload to `broken_bucket` raises, every other upload records its bucket"""
 
-    class _FixedFlushGCSLogger(GCSBucketLogger):
         def __init__(self) -> None:
             with patch("litellm.proxy.proxy_server.premium_user", True):  # test-quality-ok: GCS logging is premium-gated
-                super().__init__(bucket_name="test-bucket")
+                super().__init__(bucket_name="team-bucket")
+            self.batch_size = batch_size
+            self.uploaded_buckets: list[str] = []
 
-        async def flush_queue_and_report(self) -> GCSFlushResult:
-            return GCSFlushResult(sent=sent, failed=failed)
+        async def _enqueue(self, item: GCSLogQueueItem) -> None:
+            if enqueue_error is not None:
+                raise RuntimeError(enqueue_error)
+            await super()._enqueue(item)
 
-    return _FixedFlushGCSLogger()
+        async def enqueue(self, request_id: str, bucket_name: str) -> None:
+            payload: Final = StandardLoggingPayload(id=request_id)  # pyright: ignore[reportCallIssue]  # partial payload is enough for queueing
+            kwargs: Final = {"standard_callback_dynamic_params": {"gcs_bucket_name": bucket_name}}
+            await self._enqueue(GCSLogQueueItem(payload=payload, kwargs=kwargs, response_obj=None))
+
+        async def get_gcs_logging_config(self, kwargs: dict | None = None) -> GCSLoggingConfig:
+            dynamic_params: Final = (kwargs or {}).get("standard_callback_dynamic_params") or {}
+            return GCSLoggingConfig(
+                bucket_name=dynamic_params.get("gcs_bucket_name") or "team-bucket",
+                vertex_instance=None,
+                path_service_account=None,
+            )
+
+        async def construct_request_headers(
+            self, service_account_json: str | None, vertex_instance: VertexBase | None = None
+        ) -> dict[str, str]:
+            return {}
+
+        async def _log_json_data_on_gcs(
+            self, headers: dict[str, str], bucket_name: str, object_name: str, logging_payload: StandardLoggingPayload | str
+        ) -> None:
+            if bucket_name == broken_bucket:
+                raise RuntimeError("storage.googleapis.com returned 403")
+            self.uploaded_buckets.append(bucket_name)
+
+    return _FakeUploadGCSLogger()
 
 
 @pytest.mark.asyncio
-async def test_flush_gcs_reports_the_failed_upload_count_from_the_registered_logger():
+async def test_flush_gcs_reports_only_when_the_health_event_itself_failed_to_upload():
     from litellm.proxy.management_endpoints.key_management_endpoints import flush_gcs_and_describe_failures
 
+    logger = _fake_upload_gcs_logger(broken_bucket="team-bucket")
+    await logger.enqueue("req-1", "team-bucket")
+    await logger.enqueue("req-2", "team-bucket")
+    await logger.enqueue("req-3", "ok-bucket")
+    await logger.enqueue("health-event", "team-bucket")
+
     assert (
-        await flush_gcs_and_describe_failures(_gcs_logger_whose_flush_reports(sent=0, failed=3))
-        == "GCS upload failed for 3 event(s), 0 uploaded"
+        await flush_gcs_and_describe_failures(logger, "health-event")
+        == "GCS upload failed for the /key/health event and 2 other event(s), 1 uploaded"
     )
-    assert await flush_gcs_and_describe_failures(_gcs_logger_whose_flush_reports(sent=2, failed=0)) is None
+    assert await flush_gcs_and_describe_failures(logger, "health-event-already-uploaded") is None
 
 
-async def _key_logging_status_after_gcs_flush(sent: int, failed: int) -> LoggingCallbackStatus:
+async def _key_logging_status_with_gcs_logger(gcs_logger) -> LoggingCallbackStatus:
     from starlette.requests import Request as StarletteRequest
 
     from litellm.proxy.management_endpoints.key_management_endpoints import test_key_logging
@@ -18454,28 +18493,52 @@ async def _key_logging_status_after_gcs_flush(sent: int, failed: int) -> Logging
         patch("litellm.proxy.proxy_server.proxy_config", _default_team_gcs_proxy_config("team-gcs")),  # test-quality-ok: test_key_logging reads the module-level proxy config
         patch("litellm.proxy.proxy_server.premium_user", True),  # test-quality-ok: the mock completion's GCS success event is premium-gated
         patch(  # test-quality-ok: the mock completion and the flush both look the logger up in this process-wide registry
-            "litellm.litellm_core_utils.litellm_logging._in_memory_loggers",
-            [_gcs_logger_whose_flush_reports(sent=sent, failed=failed)],
+            "litellm.litellm_core_utils.litellm_logging._in_memory_loggers", [gcs_logger]
         ),
     ):
         return await test_key_logging(user_api_key_dict=caller, request=request, logging_callbacks=("gcs_bucket",))
 
 
 @pytest.mark.asyncio
-async def test_key_logging_marks_the_key_unhealthy_when_the_gcs_flush_leaves_events_undelivered():
-    status = await _key_logging_status_after_gcs_flush(sent=1, failed=3)
+async def test_key_logging_marks_the_key_unhealthy_when_its_own_gcs_upload_fails():
+    status = await _key_logging_status_with_gcs_logger(_fake_upload_gcs_logger(broken_bucket="team-bucket"))
 
     assert status["status"] == "unhealthy"
-    assert "GCS upload failed for 3 event(s), 1 uploaded" in (status["details"] or "")
+    assert "GCS upload failed for the /key/health event and 0 other event(s), 0 uploaded" in (status["details"] or "")
+    assert "storage.googleapis.com returned 403" in (status["details"] or "")
 
 
 @pytest.mark.asyncio
-async def test_key_logging_stays_healthy_when_the_gcs_flush_delivers_every_event():
-    status = await _key_logging_status_after_gcs_flush(sent=1, failed=0)
+async def test_key_logging_reports_the_callback_error_when_the_event_never_reaches_the_gcs_queue():
+    status = await _key_logging_status_with_gcs_logger(_fake_upload_gcs_logger(enqueue_error="queue closed"))
+
+    assert status["status"] == "unhealthy"
+    assert "GCS Bucket logging error: queue closed" in (status["details"] or "")
+    assert "GCS upload failed" not in (status["details"] or "")
+
+
+@pytest.mark.asyncio
+async def test_key_logging_stays_healthy_when_only_another_teams_queued_upload_fails():
+    logger = _fake_upload_gcs_logger(broken_bucket="other-team-bucket")
+    await logger.enqueue("req-other-team", "other-team-bucket")
+
+    status = await _key_logging_status_with_gcs_logger(logger)
 
     assert status["status"] == "healthy"
     assert status["callbacks"] == ("gcs_bucket",)
     assert "Manually check if logs were sent to gcs_bucket" in (status["details"] or "")
+    assert logger.uploaded_buckets == ["team-bucket"]
+
+
+@pytest.mark.asyncio
+async def test_key_logging_flushes_past_the_first_batch_to_reach_its_own_event():
+    logger = _fake_upload_gcs_logger(broken_bucket="team-bucket", batch_size=1)
+    await logger.enqueue("req-other-team", "other-team-bucket")
+
+    status = await _key_logging_status_with_gcs_logger(logger)
+
+    assert status["status"] == "unhealthy"
+    assert logger.uploaded_buckets == ["other-team-bucket"]
 
 
 @pytest.mark.asyncio
@@ -18483,6 +18546,6 @@ async def test_flush_gcs_names_a_missing_logger_when_the_callback_never_initiali
     from litellm.proxy.management_endpoints.key_management_endpoints import flush_gcs_and_describe_failures
 
     assert (
-        await flush_gcs_and_describe_failures(None)
+        await flush_gcs_and_describe_failures(None, "health-event")
         == "gcs_bucket callback was selected but no GCS logger was initialized"
     )