From 1f3d128569854e0f295984f737d2481010626232 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Thu, 12 Mar 2026 12:57:40 -0700
Subject: [PATCH 1/8] fix(proxy): cap managed-object poll size + expire stale
 rows + kill-switch flag to prevent OOM/Prisma connection loss

---
 .../proxy/common_utils/check_batch_cost.py    | 70 +++++++++++++++----
 .../common_utils/check_responses_cost.py      | 31 +++++++-
 litellm/constants.py                          | 10 +++
 litellm/proxy/proxy_server.py                 |  7 +-
 4 files changed, 100 insertions(+), 18 deletions(-)

diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
index 10f7f98b7197..605f0cc735a4 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -2,11 +2,15 @@
 Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
 """
 
-from litellm._uuid import uuid
-from datetime import datetime
+from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, Optional
 
 from litellm._logging import verbose_proxy_logger
+from litellm._uuid import uuid
+from litellm.constants import (
+    MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
+    MAX_OBJECTS_PER_POLL_CYCLE,
+)
 
 if TYPE_CHECKING:
     from litellm.proxy.utils import PrismaClient, ProxyLogging
@@ -49,6 +53,26 @@ async def _get_user_info(self, batch_id, user_id) -> dict:
             verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
             return {}
 
+    async def _cleanup_stale_managed_objects(self) -> None:
+        """
+        Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
+        in non-terminal states as 'stale_expired'. These will never complete and
+        should not be polled.
+        """
+        cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
+        result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
+            where={
+                "status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
+                "created_at": {"lt": cutoff},
+            },
+            data={"status": "stale_expired"},
+        )
+        if result > 0:
+            verbose_proxy_logger.warning(
+                f"CheckBatchCost: marked {result} stale managed objects "
+                f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
+            )
+
     async def check_batch_cost(self):
         """
         Check if the batch JOB has been tracked.
@@ -70,14 +94,32 @@ async def check_batch_cost(self):
             get_model_id_from_unified_batch_id,
         )
 
+        await self._cleanup_stale_managed_objects()
+
         # Look for all batches that have not yet been processed by CheckBatchCost
-        jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
-            where={
-                "file_purpose": "batch",
-                "batch_processed" : False,
-                "status": {"not_in": ["failed", "expired", "cancelled"]}
-            }
-        )
+        try:
+            jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
+                where={
+                    "file_purpose": "batch",
+                    "batch_processed": False,
+                    "status": {"not_in": ["failed", "expired", "cancelled"]},
+                },
+                take=MAX_OBJECTS_PER_POLL_CYCLE,
+                order={"created_at": "asc"},
+            )
+        except Exception:
+            # Fallback: batch_processed column may not exist on older schemas
+            verbose_proxy_logger.warning(
+                "CheckBatchCost: batch_processed column not found, querying without it"
+            )
+            jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
+                where={
+                    "file_purpose": "batch",
+                    "status": {"not_in": ["failed", "expired", "cancelled", "complete", "completed"]},
+                },
+                take=MAX_OBJECTS_PER_POLL_CYCLE,
+                order={"created_at": "asc"},
+            )
         for job in jobs:
             # get the model from the job
             unified_object_id = job.unified_object_id
@@ -163,14 +205,14 @@ async def check_batch_cost(self):
 
                 # Access content - handle both direct attribute and method call
                 if hasattr(_file_content, 'content'):
-                    content_bytes = _file_content.content
+                    content_bytes = _file_content.content  # type: ignore[union-attr]
                 elif hasattr(_file_content, 'read'):
-                    content_bytes = await _file_content.read()
+                    content_bytes = await _file_content.read()  # type: ignore[misc]
                 else:
-                    content_bytes = _file_content
+                    content_bytes = _file_content  # type: ignore[assignment]
 
                 file_content_as_dict = _get_file_content_as_dictionary(
-                    content_bytes
+                    content_bytes  # type: ignore[arg-type]
                 )
 
                 deployment_info = self.llm_router.get_deployment(model_id=model_id)
@@ -195,7 +237,7 @@ async def check_batch_cost(self):
                         file_content_dictionary=file_content_as_dict,
                         custom_llm_provider=llm_provider,  # type: ignore
                         model_name=model_name,
-                        model_info=deployment_model_info,
+                        model_info=deployment_model_info,  # type: ignore[arg-type]
                     )
                 )
                 logging_obj = LiteLLMLogging(
diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
index 4ee6a89cc98a..49a55531ad05 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
@@ -3,10 +3,15 @@
 Cost tracking is handled automatically by litellm.aget_responses().
 """
 
+from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING
 
 import litellm
 from litellm._logging import verbose_proxy_logger
+from litellm.constants import (
+    MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
+    MAX_OBJECTS_PER_POLL_CYCLE,
+)
 
 if TYPE_CHECKING:
     from litellm.proxy.utils import PrismaClient, ProxyLogging
@@ -27,6 +32,26 @@ def __init__(
         self.prisma_client: PrismaClient = prisma_client
         self.llm_router: Router = llm_router
 
+    async def _cleanup_stale_managed_objects(self) -> None:
+        """
+        Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
+        in non-terminal states as 'stale_expired'. These will never complete and
+        should not be polled.
+        """
+        cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
+        result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
+            where={
+                "status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
+                "created_at": {"lt": cutoff},
+            },
+            data={"status": "stale_expired"},
+        )
+        if result > 0:
+            verbose_proxy_logger.warning(
+                f"CheckResponsesCost: marked {result} stale managed objects "
+                f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
+            )
+
     async def check_responses_cost(self):
         """
         Check if background responses are complete and track their cost.
@@ -35,11 +60,15 @@ async def check_responses_cost(self):
         - Cost is automatically tracked by litellm.aget_responses()
         - Mark completed/failed/cancelled responses as complete in the database
         """
+        await self._cleanup_stale_managed_objects()
+
         jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
             where={
                 "status": {"in": ["queued", "in_progress"]},
                 "file_purpose": "response",
-            }
+            },
+            take=MAX_OBJECTS_PER_POLL_CYCLE,
+            order={"created_at": "asc"},
         )
         
         verbose_proxy_logger.debug(f"Found {len(jobs)} response jobs to check")
diff --git a/litellm/constants.py b/litellm/constants.py
index 34b6950a214f..05526e6ffeb6 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -1351,6 +1351,16 @@
     os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)
 )
 PROXY_BATCH_POLLING_INTERVAL = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600))
+MAX_OBJECTS_PER_POLL_CYCLE = int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50))
+MANAGED_OBJECT_STALENESS_CUTOFF_DAYS = int(
+    os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)
+)
+# Set PROXY_BATCH_POLLING_ENABLED=false to disable the CheckBatchCost and
+# CheckResponsesCost background polling jobs entirely (e.g. to avoid DB load on
+# installations with large numbers of stale managed objects).
+PROXY_BATCH_POLLING_ENABLED = (
+    os.getenv("PROXY_BATCH_POLLING_ENABLED", "true").lower() != "false"
+)
 PROXY_BUDGET_RESCHEDULER_MAX_TIME = int(
     os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)
 )
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 77e2a88796e5..2a9be0a67c92 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -214,6 +214,7 @@ def generate_feedback_box():
     DEFAULT_MODEL_CREATED_AT_TIME,
     LITELLM_PROXY_ADMIN_NAME,
     PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS,
+    PROXY_BATCH_POLLING_ENABLED,
     PROXY_BATCH_POLLING_INTERVAL,
     PROXY_BATCH_WRITE_AT,
     PROXY_BUDGET_RESCHEDULER_MAX_TIME,
@@ -260,7 +261,6 @@ def generate_feedback_box():
     claude_code_marketplace_router,
 )
 from litellm.proxy.anthropic_endpoints.endpoints import router as anthropic_router
-from litellm.proxy.realtime_endpoints.endpoints import router as webrtc_router
 from litellm.proxy.anthropic_endpoints.skills_endpoints import (
     router as anthropic_skills_router,
 )
@@ -471,6 +471,7 @@ def generate_feedback_box():
 from litellm.proxy.prompts.prompt_endpoints import router as prompts_router
 from litellm.proxy.public_endpoints import router as public_endpoints_router
 from litellm.proxy.rag_endpoints.endpoints import router as rag_router
+from litellm.proxy.realtime_endpoints.endpoints import router as webrtc_router
 from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router
 from litellm.proxy.response_api_endpoints.endpoints import router as response_router
 from litellm.proxy.route_llm_request import route_request
@@ -6069,7 +6070,7 @@ async def initialize_scheduled_background_jobs(  # noqa: PLR0915
                         "Invalid maximum_spend_logs_retention_interval value"
                     )
         ### CHECK BATCH COST ###
-        if llm_router is not None:
+        if llm_router is not None and PROXY_BATCH_POLLING_ENABLED:
             try:
                 from litellm_enterprise.proxy.common_utils.check_batch_cost import (
                     CheckBatchCost,
@@ -6100,7 +6101,7 @@ async def initialize_scheduled_background_jobs(  # noqa: PLR0915
                 pass
 
         ### CHECK RESPONSES COST ###
-        if llm_router is not None:
+        if llm_router is not None and PROXY_BATCH_POLLING_ENABLED:
             try:
                 from litellm_enterprise.proxy.common_utils.check_responses_cost import (
                     CheckResponsesCost,

From b2252b4b2a67fa1588667172e8a65ccf364e160c Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Thu, 12 Mar 2026 12:58:07 -0700
Subject: [PATCH 2/8] fix(constants): simplify PROXY_BATCH_POLLING_ENABLED
 readability

---
 litellm/constants.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/litellm/constants.py b/litellm/constants.py
index 05526e6ffeb6..2106c527d961 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -1358,9 +1358,8 @@
 # Set PROXY_BATCH_POLLING_ENABLED=false to disable the CheckBatchCost and
 # CheckResponsesCost background polling jobs entirely (e.g. to avoid DB load on
 # installations with large numbers of stale managed objects).
-PROXY_BATCH_POLLING_ENABLED = (
-    os.getenv("PROXY_BATCH_POLLING_ENABLED", "true").lower() != "false"
-)
+_batch_polling_env = os.getenv("PROXY_BATCH_POLLING_ENABLED", "true").lower()
+PROXY_BATCH_POLLING_ENABLED = _batch_polling_env == "true"
 PROXY_BUDGET_RESCHEDULER_MAX_TIME = int(
     os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)
 )

From 5acd8f6c6ed69751db3d218f0792c2cee9a35580 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Thu, 12 Mar 2026 13:02:01 -0700
Subject: [PATCH 3/8] docs+test: document new polling env vars, add
 pagination+stale-cleanup tests

---
 docs/my-website/docs/proxy/config_settings.md |  3 ++
 .../test_check_responses_cost.py              | 41 +++++++++++++++----
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md
index 90d71fb95c89..90045755e9c2 100644
--- a/docs/my-website/docs/proxy/config_settings.md
+++ b/docs/my-website/docs/proxy/config_settings.md
@@ -934,6 +934,9 @@ router_settings:
 | PROXY_BASE_URL | Base URL for proxy service
 | PROXY_BATCH_WRITE_AT | Time in seconds to wait before batch writing spend logs to the database. Default is 10
 | PROXY_BATCH_POLLING_INTERVAL | Time in seconds to wait before polling a batch, to check if it's completed. Default is 6000s (1 hour)
+| PROXY_BATCH_POLLING_ENABLED | Set to `false` to disable the `CheckBatchCost` and `CheckResponsesCost` background polling jobs entirely. Useful for emergency mitigation on installs with large numbers of stale managed objects. Default is `true`
+| MAX_OBJECTS_PER_POLL_CYCLE | Maximum number of managed objects (batches / responses) fetched per polling cycle. Prevents OOM on installs with many stale rows. Default is `50`
+| MANAGED_OBJECT_STALENESS_CUTOFF_DAYS | Managed objects older than this many days in a non-terminal state are marked `stale_expired` at the start of each poll cycle and skipped. Default is `7`
 | PROXY_BUDGET_RESCHEDULER_MAX_TIME | Maximum time in seconds to wait before checking database for budget resets. Default is 605
 | PROXY_BUDGET_RESCHEDULER_MIN_TIME | Minimum time in seconds to wait before checking database for budget resets. Default is 597
 | PYTHON_GC_THRESHOLD | GC thresholds ('gen0,gen1,gen2', e.g. '1000,50,50'); defaults to Python’s values.
diff --git a/tests/proxy_unit_tests/test_check_responses_cost.py b/tests/proxy_unit_tests/test_check_responses_cost.py
index 3bcacdfc05d2..bbffd6e0e75d 100644
--- a/tests/proxy_unit_tests/test_check_responses_cost.py
+++ b/tests/proxy_unit_tests/test_check_responses_cost.py
@@ -8,6 +8,7 @@
 
 import pytest
 
+from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
 from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
 
 
@@ -63,21 +64,45 @@ async def test_check_responses_cost_no_jobs(
         self, check_responses_cost_instance, mock_prisma_client
     ):
         """Test check_responses_cost when there are no jobs to process"""
-        # Mock empty job list
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
             return_value=[]
         )
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
 
-        # Should not raise any errors
         await check_responses_cost_instance.check_responses_cost()
 
-        # Verify find_many was called with correct parameters
-        mock_prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with(
-            where={
-                "status": {"in": ["queued", "in_progress"]},
-                "file_purpose": "response",
-            }
+        # Verify find_many was called with pagination params
+        find_many_call = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args
+        assert find_many_call[1]["where"] == {
+            "status": {"in": ["queued", "in_progress"]},
+            "file_purpose": "response",
+        }
+        assert find_many_call[1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE
+        assert find_many_call[1]["order"] == {"created_at": "asc"}
+
+    @pytest.mark.asyncio
+    async def test_cleanup_stale_managed_objects(
+        self, check_responses_cost_instance, mock_prisma_client
+    ):
+        """Stale rows (older than cutoff) are bulk-updated to stale_expired before polling."""
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=5
         )
+        mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
+            return_value=[]
+        )
+
+        await check_responses_cost_instance.check_responses_cost()
+
+        # The first update_many call should be the stale-row cleanup
+        calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+        stale_call = calls[0]
+        assert stale_call[1]["data"] == {"status": "stale_expired"}
+        where = stale_call[1]["where"]
+        assert "stale_expired" in where["status"]["not_in"]
+        assert "created_at" in where
 
     @pytest.mark.asyncio
     async def test_check_responses_cost_with_completed_response(

From e7f8cd50d04164a3bfb359436b5e8ba852aba932 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Thu, 12 Mar 2026 13:24:40 -0700
Subject: [PATCH 4/8] fix: exclude stale_expired from batch poll queries; fix
 update_many assertions in tests

---
 .../proxy/common_utils/check_batch_cost.py    | 13 ++-
 .../test_check_responses_cost.py              | 93 +++++++++++--------
 2 files changed, 66 insertions(+), 40 deletions(-)

diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
index 605f0cc735a4..fe3e8f1402c1 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -102,7 +102,7 @@ async def check_batch_cost(self):
                 where={
                     "file_purpose": "batch",
                     "batch_processed": False,
-                    "status": {"not_in": ["failed", "expired", "cancelled"]},
+                    "status": {"not_in": ["failed", "expired", "cancelled", "stale_expired"]},
                 },
                 take=MAX_OBJECTS_PER_POLL_CYCLE,
                 order={"created_at": "asc"},
@@ -115,7 +115,16 @@ async def check_batch_cost(self):
             jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
                 where={
                     "file_purpose": "batch",
-                    "status": {"not_in": ["failed", "expired", "cancelled", "complete", "completed"]},
+                    "status": {
+                        "not_in": [
+                            "failed",
+                            "expired",
+                            "cancelled",
+                            "complete",
+                            "completed",
+                            "stale_expired",
+                        ]
+                    },
                 },
                 take=MAX_OBJECTS_PER_POLL_CYCLE,
                 order={"created_at": "asc"},
diff --git a/tests/proxy_unit_tests/test_check_responses_cost.py b/tests/proxy_unit_tests/test_check_responses_cost.py
index bbffd6e0e75d..d542b8c2181c 100644
--- a/tests/proxy_unit_tests/test_check_responses_cost.py
+++ b/tests/proxy_unit_tests/test_check_responses_cost.py
@@ -133,8 +133,9 @@ async def test_check_responses_cost_with_completed_response(
             ),
         )
 
-        # Mock update_many
-        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
 
         # Run the check with mocked litellm.aget_responses
         with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
@@ -142,11 +143,12 @@ async def test_check_responses_cost_with_completed_response(
 
             await check_responses_cost_instance.check_responses_cost()
 
-        # Verify the job was marked as completed
-        mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_called_once()
-        call_args = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args
-        assert call_args[1]["data"]["status"] == "completed"
-        assert call_args[1]["where"]["id"]["in"] == ["job-123"]
+        # calls[0] = stale cleanup, calls[1] = job completion
+        calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+        assert len(calls) == 2
+        completion_call = calls[1]
+        assert completion_call[1]["data"]["status"] == "completed"
+        assert completion_call[1]["where"]["id"]["in"] == ["job-123"]
 
     @pytest.mark.asyncio
     async def test_check_responses_cost_with_failed_response(
@@ -173,8 +175,9 @@ async def test_check_responses_cost_with_failed_response(
             usage=None,
         )
 
-        # Mock update_many
-        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
 
         # Run the check
         with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
@@ -182,10 +185,10 @@ async def test_check_responses_cost_with_failed_response(
 
             await check_responses_cost_instance.check_responses_cost()
 
-        # Verify the job was marked as completed (even though response failed)
-        mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_called_once()
-        call_args = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args
-        assert call_args[1]["data"]["status"] == "completed"
+        # calls[0] = stale cleanup, calls[1] = job completion
+        calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+        assert len(calls) == 2
+        assert calls[1][1]["data"]["status"] == "completed"
 
     @pytest.mark.asyncio
     async def test_check_responses_cost_with_cancelled_response(
@@ -212,8 +215,9 @@ async def test_check_responses_cost_with_cancelled_response(
             usage=None,
         )
 
-        # Mock update_many
-        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
 
         # Run the check
         with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
@@ -221,8 +225,10 @@ async def test_check_responses_cost_with_cancelled_response(
 
             await check_responses_cost_instance.check_responses_cost()
 
-        # Verify the job was marked as completed
-        mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_called_once()
+        # calls[0] = stale cleanup, calls[1] = job completion
+        calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+        assert len(calls) == 2
+        assert calls[1][1]["data"]["status"] == "completed"
 
     @pytest.mark.asyncio
     async def test_check_responses_cost_with_in_progress_response(
@@ -249,8 +255,9 @@ async def test_check_responses_cost_with_in_progress_response(
             usage=None,
         )
 
-        # Mock update_many
-        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
 
         # Run the check
         with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
@@ -258,8 +265,10 @@ async def test_check_responses_cost_with_in_progress_response(
 
             await check_responses_cost_instance.check_responses_cost()
 
-        # Verify no updates were made (response still in progress)
-        mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_not_called()
+        # Only the stale-cleanup call should have fired — no completion update
+        calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+        assert len(calls) == 1
+        assert calls[0][1]["data"] == {"status": "stale_expired"}
 
     @pytest.mark.asyncio
     async def test_check_responses_cost_with_queued_response(
@@ -286,8 +295,9 @@ async def test_check_responses_cost_with_queued_response(
             usage=None,
         )
 
-        # Mock update_many
-        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
 
         # Run the check
         with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
@@ -295,8 +305,10 @@ async def test_check_responses_cost_with_queued_response(
 
             await check_responses_cost_instance.check_responses_cost()
 
-        # Verify no updates were made (response still queued)
-        mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_not_called()
+        # Only the stale-cleanup call should have fired — no completion update
+        calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+        assert len(calls) == 1
+        assert calls[0][1]["data"] == {"status": "stale_expired"}
 
     @pytest.mark.asyncio
     async def test_check_responses_cost_with_exception(
@@ -313,8 +325,9 @@ async def test_check_responses_cost_with_exception(
             return_value=[mock_job]
         )
 
-        # Mock update_many
-        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
 
         # Run the check with mocked exception
         with patch(
@@ -325,8 +338,10 @@ async def test_check_responses_cost_with_exception(
             # Should not raise, just skip the job
             await check_responses_cost_instance.check_responses_cost()
 
-        # Verify no updates were made (job was skipped due to error)
-        mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_not_called()
+        # Only the stale-cleanup call should have fired — no completion update
+        calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+        assert len(calls) == 1
+        assert calls[0][1]["data"] == {"status": "stale_expired"}
 
     @pytest.mark.asyncio
     async def test_check_responses_cost_multiple_jobs(
@@ -389,8 +404,9 @@ async def test_check_responses_cost_multiple_jobs(
             ),
         )
 
-        # Mock update_many
-        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
 
         # Run the check
         with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
@@ -398,10 +414,11 @@ async def test_check_responses_cost_multiple_jobs(
 
             await check_responses_cost_instance.check_responses_cost()
 
-        # Verify only the 2 completed jobs were marked as complete
-        mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_called_once()
-        call_args = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args
-        assert len(call_args[1]["where"]["id"]["in"]) == 2
-        assert "job-1" in call_args[1]["where"]["id"]["in"]
-        assert "job-3" in call_args[1]["where"]["id"]["in"]
-        assert "job-2" not in call_args[1]["where"]["id"]["in"]
+        # calls[0] = stale cleanup, calls[1] = completion of 2 finished jobs
+        calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+        assert len(calls) == 2
+        completion_call = calls[1]
+        assert len(completion_call[1]["where"]["id"]["in"]) == 2
+        assert "job-1" in completion_call[1]["where"]["id"]["in"]
+        assert "job-3" in completion_call[1]["where"]["id"]["in"]
+        assert "job-2" not in completion_call[1]["where"]["id"]["in"]

From 52c0574cc5de4219193f9c8ad327d3b4f1470af7 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Thu, 12 Mar 2026 16:34:00 -0700
Subject: [PATCH 5/8] fix: scope stale cleanup to file_purpose, fix file_object
 mocks, add CheckBatchCost tests

---
 .../proxy/common_utils/check_batch_cost.py    |   1 +
 .../common_utils/check_responses_cost.py      |   1 +
 .../proxy_unit_tests/test_check_batch_cost.py | 112 ++++++++++++++++++
 .../test_check_responses_cost.py              |  42 ++++++-
 4 files changed, 155 insertions(+), 1 deletion(-)
 create mode 100644 tests/proxy_unit_tests/test_check_batch_cost.py

diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
index fe3e8f1402c1..e89501fca6f5 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -62,6 +62,7 @@ async def _cleanup_stale_managed_objects(self) -> None:
         cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
         result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
             where={
+                "file_purpose": "batch",
                 "status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
                 "created_at": {"lt": cutoff},
             },
diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
index 49a55531ad05..6ee2bce272ac 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
@@ -41,6 +41,7 @@ async def _cleanup_stale_managed_objects(self) -> None:
         cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
         result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
             where={
+                "file_purpose": "response",
                 "status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
                 "created_at": {"lt": cutoff},
             },
diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py
new file mode 100644
index 000000000000..36a2f35de1a3
--- /dev/null
+++ b/tests/proxy_unit_tests/test_check_batch_cost.py
@@ -0,0 +1,112 @@
+"""
+Unit tests for CheckBatchCost class.
+Covers: stale-row cleanup (file_purpose scoping), paginated find_many,
+and the batch_processed-column fallback query.
+"""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+
+class TestCheckBatchCost:
+    """Test suite for CheckBatchCost class"""
+
+    @pytest.fixture
+    def mock_prisma_client(self):
+        client = MagicMock()
+        client.db = MagicMock()
+        client.db.litellm_managedobjecttable = MagicMock()
+        client.db.litellm_usertable = MagicMock()
+        return client
+
+    @pytest.fixture
+    def mock_proxy_logging_obj(self):
+        return MagicMock()
+
+    @pytest.fixture
+    def mock_llm_router(self):
+        return MagicMock()
+
+    @pytest.fixture
+    def check_batch_cost_instance(
+        self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
+    ):
+        from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
+
+        return CheckBatchCost(
+            proxy_logging_obj=mock_proxy_logging_obj,
+            prisma_client=mock_prisma_client,
+            llm_router=mock_llm_router,
+        )
+
+    @pytest.mark.asyncio
+    async def test_cleanup_scoped_to_batch_file_purpose(
+        self, check_batch_cost_instance, mock_prisma_client
+    ):
+        """_cleanup_stale_managed_objects scopes its update to file_purpose='batch' only."""
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
+        # Return empty so the main poll loop exits immediately
+        mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
+            return_value=[]
+        )
+
+        await check_batch_cost_instance.check_batch_cost()
+
+        calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+        stale_call = calls[0]
+        assert stale_call[1]["data"] == {"status": "stale_expired"}
+        where = stale_call[1]["where"]
+        assert where["file_purpose"] == "batch"
+        assert "stale_expired" in where["status"]["not_in"]
+        assert "created_at" in where
+
+    @pytest.mark.asyncio
+    async def test_find_many_uses_pagination_and_excludes_stale(
+        self, check_batch_cost_instance, mock_prisma_client
+    ):
+        """find_many is called with take, order, and stale_expired excluded from status."""
+        from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
+
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
+        mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
+            return_value=[]
+        )
+
+        await check_batch_cost_instance.check_batch_cost()
+
+        find_call = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args
+        assert find_call[1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE
+        assert find_call[1]["order"] == {"created_at": "asc"}
+        assert "stale_expired" in find_call[1]["where"]["status"]["not_in"]
+
+    @pytest.mark.asyncio
+    async def test_fallback_query_used_when_batch_processed_missing(
+        self, check_batch_cost_instance, mock_prisma_client
+    ):
+        """Falls back to query without batch_processed when primary query raises."""
+        from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
+
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
+        # First find_many (primary query) raises; second (fallback) returns empty list
+        mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
+            side_effect=[Exception("column batch_processed does not exist"), []]
+        )
+
+        await check_batch_cost_instance.check_batch_cost()
+
+        calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list
+        assert len(calls) == 2
+        fallback_where = calls[1][1]["where"]
+        # Fallback must not reference batch_processed
+        assert "batch_processed" not in fallback_where
+        # Fallback must exclude stale_expired
+        assert "stale_expired" in fallback_where["status"]["not_in"]
+        # Fallback must still paginate
+        assert calls[1][1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE
diff --git a/tests/proxy_unit_tests/test_check_responses_cost.py b/tests/proxy_unit_tests/test_check_responses_cost.py
index d542b8c2181c..601df9c4c7f9 100644
--- a/tests/proxy_unit_tests/test_check_responses_cost.py
+++ b/tests/proxy_unit_tests/test_check_responses_cost.py
@@ -96,11 +96,12 @@ async def test_cleanup_stale_managed_objects(
 
         await check_responses_cost_instance.check_responses_cost()
 
-        # The first update_many call should be the stale-row cleanup
+        # The first update_many call should be the stale-row cleanup scoped to "response"
         calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
         stale_call = calls[0]
         assert stale_call[1]["data"] == {"status": "stale_expired"}
         where = stale_call[1]["where"]
+        assert where["file_purpose"] == "response"
         assert "stale_expired" in where["status"]["not_in"]
         assert "created_at" in where
 
@@ -114,6 +115,7 @@ async def test_check_responses_cost_with_completed_response(
         mock_job.unified_object_id = "resp_test_123"
         mock_job.created_by = "test-user"
         mock_job.id = "job-123"
+        mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_123"}
 
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
             return_value=[mock_job]
@@ -160,6 +162,7 @@ async def test_check_responses_cost_with_failed_response(
         mock_job.unified_object_id = "resp_test_456"
         mock_job.created_by = "test-user"
         mock_job.id = "job-456"
+        mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_456"}
 
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
             return_value=[mock_job]
@@ -200,6 +203,7 @@ async def test_check_responses_cost_with_cancelled_response(
         mock_job.unified_object_id = "resp_test_789"
         mock_job.created_by = "test-user"
         mock_job.id = "job-789"
+        mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_789"}
 
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
             return_value=[mock_job]
@@ -240,6 +244,7 @@ async def test_check_responses_cost_with_in_progress_response(
         mock_job.unified_object_id = "resp_test_in_progress"
         mock_job.created_by = "test-user"
         mock_job.id = "job-in-progress"
+        mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_in_progress"}
 
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
             return_value=[mock_job]
@@ -280,6 +285,7 @@ async def test_check_responses_cost_with_queued_response(
         mock_job.unified_object_id = "resp_test_queued"
         mock_job.created_by = "test-user"
         mock_job.id = "job-queued"
+        mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_queued"}
 
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
             return_value=[mock_job]
@@ -320,6 +326,7 @@ async def test_check_responses_cost_with_exception(
         mock_job.unified_object_id = "resp_test_error"
         mock_job.created_by = "test-user"
         mock_job.id = "job-error"
+        mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_error"}
 
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
             return_value=[mock_job]
@@ -353,16 +360,19 @@ async def test_check_responses_cost_multiple_jobs(
         mock_job1.unified_object_id = "resp_test_1"
         mock_job1.created_by = "user1"
         mock_job1.id = "job-1"
+        mock_job1.file_object = {"model": "gpt-4o", "id": "resp_test_1"}
 
         mock_job2 = MagicMock()
         mock_job2.unified_object_id = "resp_test_2"
         mock_job2.created_by = "user2"
         mock_job2.id = "job-2"
+        mock_job2.file_object = {"model": "gpt-4o", "id": "resp_test_2"}
 
         mock_job3 = MagicMock()
         mock_job3.unified_object_id = "resp_test_3"
         mock_job3.created_by = "user3"
         mock_job3.id = "job-3"
+        mock_job3.file_object = {"model": "gpt-4o", "id": "resp_test_3"}
 
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
             return_value=[mock_job1, mock_job2, mock_job3]
@@ -422,3 +432,33 @@ async def test_check_responses_cost_multiple_jobs(
         assert "job-1" in completion_call[1]["where"]["id"]["in"]
         assert "job-3" in completion_call[1]["where"]["id"]["in"]
         assert "job-2" not in completion_call[1]["where"]["id"]["in"]
+
+    @pytest.mark.asyncio
+    async def test_check_responses_cost_no_model_in_file_object(
+        self, check_responses_cost_instance, mock_prisma_client
+    ):
+        """When file_object has no 'model' key, model_name is None and metadata skips model fields."""
+        mock_job = MagicMock()
+        mock_job.unified_object_id = "resp_test_no_model"
+        mock_job.created_by = "test-user"
+        mock_job.id = "job-no-model"
+        mock_job.file_object = {}  # no "model" key → model_name=None branch
+
+        mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
+            return_value=[mock_job]
+        )
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
+
+        mock_response = MagicMock()
+        mock_response.status = "completed"
+
+        with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
+            mock_aget.return_value = mock_response
+            await check_responses_cost_instance.check_responses_cost()
+
+        # aget_responses should be called without model metadata
+        call_kwargs = mock_aget.call_args[1]
+        assert "model" not in call_kwargs.get("litellm_metadata", {})
+        assert "model_group" not in call_kwargs.get("litellm_metadata", {})

From a14b95124880fd6142731456a6b6d5cea2938e4a Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Thu, 12 Mar 2026 18:00:55 -0700
Subject: [PATCH 6/8] fix: avoid duplicate cost logging in fallback path; guard
 integer constants against zero/negative values

---
 .../proxy/common_utils/check_batch_cost.py    | 19 +++++---
 litellm/constants.py                          |  6 +--
 .../proxy_unit_tests/test_check_batch_cost.py | 46 +++++++++++++++++++
 3 files changed, 62 insertions(+), 9 deletions(-)

diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
index e89501fca6f5..66761e84e575 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -97,7 +97,11 @@ async def check_batch_cost(self):
 
         await self._cleanup_stale_managed_objects()
 
-        # Look for all batches that have not yet been processed by CheckBatchCost
+        # Look for all batches that have not yet been processed by CheckBatchCost.
+        # _has_batch_processed_column tracks whether the column exists so the
+        # completion update can omit it on older schemas (avoiding a silent failure
+        # that would cause infinite reprocessing and duplicate cost logging).
+        _has_batch_processed_column = True
         try:
             jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
                 where={
@@ -110,6 +114,7 @@ async def check_batch_cost(self):
             )
         except Exception:
             # Fallback: batch_processed column may not exist on older schemas
+            _has_batch_processed_column = False
             verbose_proxy_logger.warning(
                 "CheckBatchCost: batch_processed column not found, querying without it"
             )
@@ -288,13 +293,15 @@ async def check_batch_cost(self):
 
                 # mark the job as complete
                 try:
+                    update_data: dict = {
+                        "status": "complete",
+                        "file_object": response.model_dump_json(),
+                    }
+                    if _has_batch_processed_column:
+                        update_data["batch_processed"] = True
                     await self.prisma_client.db.litellm_managedobjecttable.update(
                         where={"id": job.id},
-                        data={
-                            "batch_processed": True,
-                            "status": "complete",
-                            "file_object": response.model_dump_json(),
-                        },
+                        data=update_data,
                     )
                 except Exception as db_err:
                     verbose_proxy_logger.error(
diff --git a/litellm/constants.py b/litellm/constants.py
index 2106c527d961..dbc79b69a67c 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -1351,9 +1351,9 @@
     os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)
 )
 PROXY_BATCH_POLLING_INTERVAL = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600))
-MAX_OBJECTS_PER_POLL_CYCLE = int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50))
-MANAGED_OBJECT_STALENESS_CUTOFF_DAYS = int(
-    os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)
+MAX_OBJECTS_PER_POLL_CYCLE = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50)))
+MANAGED_OBJECT_STALENESS_CUTOFF_DAYS = max(
+    1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))
 )
 # Set PROXY_BATCH_POLLING_ENABLED=false to disable the CheckBatchCost and
 # CheckResponsesCost background polling jobs entirely (e.g. to avoid DB load on
diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py
index 36a2f35de1a3..94c33669e8b4 100644
--- a/tests/proxy_unit_tests/test_check_batch_cost.py
+++ b/tests/proxy_unit_tests/test_check_batch_cost.py
@@ -110,3 +110,49 @@ async def test_fallback_query_used_when_batch_processed_missing(
         assert "stale_expired" in fallback_where["status"]["not_in"]
         # Fallback must still paginate
         assert calls[1][1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE
+
+    @pytest.mark.asyncio
+    async def test_fallback_completion_update_omits_batch_processed(
+        self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
+    ):
+        """When batch_processed column is absent, completion update must not include it.
+
+        If it did, the update would fail silently, the job would never be marked done,
+        and every subsequent poll cycle would re-log the cost (duplicate billing).
+        """
+        from unittest.mock import patch
+
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
+        mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
+
+        mock_job = MagicMock()
+        mock_job.id = "job-fallback-1"
+        mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA=="  # base64-looking value
+        mock_job.created_by = "user-1"
+
+        # Primary query fails → fallback path
+        mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
+            side_effect=[Exception("column batch_processed does not exist"), [mock_job]]
+        )
+
+        # Stub out the heavy per-job processing so we reach the update()
+        with (
+            patch(
+                "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id",
+                return_value=None,  # causes "not a valid unified object id" early-continue
+            ),
+        ):
+            await check_batch_cost_instance.check_batch_cost()
+
+        # Even though the job was skipped (invalid ID), confirm the fallback path was taken
+        # by checking the find_many calls
+        find_calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list
+        assert len(find_calls) == 2
+        fallback_where = find_calls[1][1]["where"]
+        assert "batch_processed" not in fallback_where
+
+        # If a completion update were issued, it must not contain batch_processed
+        for call in mock_prisma_client.db.litellm_managedobjecttable.update.call_args_list:
+            assert "batch_processed" not in call[1].get("data", {})

From 679b3b247ebff9222d3855f96630a9bd120323fa Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Thu, 12 Mar 2026 18:28:10 -0700
Subject: [PATCH 7/8] fix: cache _has_batch_processed_column; guard cleanup
 from aborting poll; narrow fallback except

---
 .../proxy/common_utils/check_batch_cost.py    | 90 +++++++++++--------
 .../common_utils/check_responses_cost.py      |  7 +-
 .../proxy_unit_tests/test_check_batch_cost.py | 30 ++++++-
 3 files changed, 85 insertions(+), 42 deletions(-)

diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
index 66761e84e575..5609225f9f14 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -33,6 +33,9 @@ def __init__(
         self.proxy_logging_obj: ProxyLogging = proxy_logging_obj
         self.prisma_client: PrismaClient = prisma_client
         self.llm_router: Router = llm_router
+        # Cached after the first poll cycle. Once we know the column is absent we skip
+        # the guaranteed-failing primary query on every subsequent cycle.
+        self._has_batch_processed_column: bool = True
 
     async def _get_user_info(self, batch_id, user_id) -> dict:
         """
@@ -74,6 +77,26 @@ async def _cleanup_stale_managed_objects(self) -> None:
                 f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
             )
 
+    async def _fallback_find_jobs(self) -> list:
+        """Query batch jobs without the batch_processed filter (for older schemas)."""
+        return await self.prisma_client.db.litellm_managedobjecttable.find_many(
+            where={
+                "file_purpose": "batch",
+                "status": {
+                    "not_in": [
+                        "failed",
+                        "expired",
+                        "cancelled",
+                        "complete",
+                        "completed",
+                        "stale_expired",
+                    ]
+                },
+            },
+            take=MAX_OBJECTS_PER_POLL_CYCLE,
+            order={"created_at": "asc"},
+        )
+
     async def check_batch_cost(self):
         """
         Check if the batch JOB has been tracked.
@@ -95,46 +118,39 @@ async def check_batch_cost(self):
             get_model_id_from_unified_batch_id,
         )
 
-        await self._cleanup_stale_managed_objects()
-
-        # Look for all batches that have not yet been processed by CheckBatchCost.
-        # _has_batch_processed_column tracks whether the column exists so the
-        # completion update can omit it on older schemas (avoiding a silent failure
-        # that would cause infinite reprocessing and duplicate cost logging).
-        _has_batch_processed_column = True
         try:
-            jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
-                where={
-                    "file_purpose": "batch",
-                    "batch_processed": False,
-                    "status": {"not_in": ["failed", "expired", "cancelled", "stale_expired"]},
-                },
-                take=MAX_OBJECTS_PER_POLL_CYCLE,
-                order={"created_at": "asc"},
-            )
-        except Exception:
-            # Fallback: batch_processed column may not exist on older schemas
-            _has_batch_processed_column = False
+            await self._cleanup_stale_managed_objects()
+        except Exception as cleanup_err:
             verbose_proxy_logger.warning(
-                "CheckBatchCost: batch_processed column not found, querying without it"
+                f"CheckBatchCost: stale cleanup failed (poll will continue): {cleanup_err}"
             )
-            jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
-                where={
-                    "file_purpose": "batch",
-                    "status": {
-                        "not_in": [
-                            "failed",
-                            "expired",
-                            "cancelled",
-                            "complete",
-                            "completed",
-                            "stale_expired",
-                        ]
+
+        # Look for all batches that have not yet been processed by CheckBatchCost.
+        # self._has_batch_processed_column is cached after the first probe so that
+        # older schemas don't pay a guaranteed-failing primary query + warning on
+        # every subsequent poll cycle.
+        if self._has_batch_processed_column:
+            try:
+                jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
+                    where={
+                        "file_purpose": "batch",
+                        "batch_processed": False,
+                        "status": {"not_in": ["failed", "expired", "cancelled", "stale_expired"]},
                     },
-                },
-                take=MAX_OBJECTS_PER_POLL_CYCLE,
-                order={"created_at": "asc"},
-            )
+                    take=MAX_OBJECTS_PER_POLL_CYCLE,
+                    order={"created_at": "asc"},
+                )
+            except Exception as query_err:
+                if "batch_processed" not in str(query_err).lower() and "unknown column" not in str(query_err).lower() and "does not exist" not in str(query_err).lower():
+                    raise
+                # Permanent schema gap — cache the result so future cycles skip straight to fallback
+                self._has_batch_processed_column = False
+                verbose_proxy_logger.warning(
+                    "CheckBatchCost: batch_processed column not found, querying without it"
+                )
+                jobs = await self._fallback_find_jobs()
+        else:
+            jobs = await self._fallback_find_jobs()
         for job in jobs:
             # get the model from the job
             unified_object_id = job.unified_object_id
@@ -297,7 +313,7 @@ async def check_batch_cost(self):
                         "status": "complete",
                         "file_object": response.model_dump_json(),
                     }
-                    if _has_batch_processed_column:
+                    if self._has_batch_processed_column:
                         update_data["batch_processed"] = True
                     await self.prisma_client.db.litellm_managedobjecttable.update(
                         where={"id": job.id},
diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
index 6ee2bce272ac..54fbc7abcc59 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
@@ -61,7 +61,12 @@ async def check_responses_cost(self):
         - Cost is automatically tracked by litellm.aget_responses()
         - Mark completed/failed/cancelled responses as complete in the database
         """
-        await self._cleanup_stale_managed_objects()
+        try:
+            await self._cleanup_stale_managed_objects()
+        except Exception as cleanup_err:
+            verbose_proxy_logger.warning(
+                f"CheckResponsesCost: stale cleanup failed (poll will continue): {cleanup_err}"
+            )
 
         jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
             where={
diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py
index 94c33669e8b4..7606ef49c2c1 100644
--- a/tests/proxy_unit_tests/test_check_batch_cost.py
+++ b/tests/proxy_unit_tests/test_check_batch_cost.py
@@ -94,7 +94,7 @@ async def test_fallback_query_used_when_batch_processed_missing(
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
             return_value=0
         )
-        # First find_many (primary query) raises; second (fallback) returns empty list
+        # First find_many (primary query) raises with a schema error; second (fallback) returns empty
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
             side_effect=[Exception("column batch_processed does not exist"), []]
         )
@@ -104,12 +104,34 @@ async def test_fallback_query_used_when_batch_processed_missing(
         calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list
         assert len(calls) == 2
         fallback_where = calls[1][1]["where"]
-        # Fallback must not reference batch_processed
         assert "batch_processed" not in fallback_where
-        # Fallback must exclude stale_expired
         assert "stale_expired" in fallback_where["status"]["not_in"]
-        # Fallback must still paginate
         assert calls[1][1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE
+        # Column absence is now cached — next call should go straight to fallback
+        assert check_batch_cost_instance._has_batch_processed_column is False
+
+    @pytest.mark.asyncio
+    async def test_column_absence_cached_across_cycles(
+        self, check_batch_cost_instance, mock_prisma_client
+    ):
+        """After column absence is discovered, subsequent cycles skip the primary query entirely."""
+        from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
+
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
+        # Simulate column already known absent from a previous cycle
+        check_batch_cost_instance._has_batch_processed_column = False
+        mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
+            return_value=[]
+        )
+
+        await check_batch_cost_instance.check_batch_cost()
+
+        # Only one find_many call — the fallback directly, no primary query attempt
+        assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1
+        fallback_where = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1]["where"]
+        assert "batch_processed" not in fallback_where
 
     @pytest.mark.asyncio
     async def test_fallback_completion_update_omits_batch_processed(

From b75f6162ed1b69a33f29c89c8423c13cfaa00f0f Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Fri, 13 Mar 2026 09:01:12 -0700
Subject: [PATCH 8/8] fix: add complete/completed to primary query not_in; fix
 vacuous test assertion
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Primary find_many was missing "complete" and "completed" in its not_in
  filter, creating asymmetry with the fallback query. A job whose status
  was set to "complete" but whose batch_processed flag update failed would
  be silently re-fetched and re-processed every cycle, emitting duplicate
  cost logs.

- test_fallback_completion_update_omits_batch_processed patched
  _is_base64_encoded_unified_file_id to return None, causing an immediate
  continue — so update() was never called and the assertion looped over an
  empty list (vacuously true). Rewrote the test to mock the full
  completion pipeline, verify update() is called exactly once, and assert
  batch_processed is absent from the update data.

- Added symmetric test (primary path) proving batch_processed IS included
  when the column exists.

Made-with: Cursor
---
 .../proxy/common_utils/check_batch_cost.py    |  11 +-
 .../proxy_unit_tests/test_check_batch_cost.py | 192 ++++++++++++++++--
 2 files changed, 186 insertions(+), 17 deletions(-)

diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
index 5609225f9f14..42a9acbfd1e8 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -135,7 +135,16 @@ async def check_batch_cost(self):
                     where={
                         "file_purpose": "batch",
                         "batch_processed": False,
-                        "status": {"not_in": ["failed", "expired", "cancelled", "stale_expired"]},
+                        "status": {
+                            "not_in": [
+                                "failed",
+                                "expired",
+                                "cancelled",
+                                "complete",
+                                "completed",
+                                "stale_expired",
+                            ]
+                        },
                     },
                     take=MAX_OBJECTS_PER_POLL_CYCLE,
                     order={"created_at": "asc"},
diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py
index 7606ef49c2c1..f6b8d567848a 100644
--- a/tests/proxy_unit_tests/test_check_batch_cost.py
+++ b/tests/proxy_unit_tests/test_check_batch_cost.py
@@ -67,7 +67,7 @@ async def test_cleanup_scoped_to_batch_file_purpose(
     async def test_find_many_uses_pagination_and_excludes_stale(
         self, check_batch_cost_instance, mock_prisma_client
     ):
-        """find_many is called with take, order, and stale_expired excluded from status."""
+        """find_many is called with take, order, and all terminal statuses excluded."""
         from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
 
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
@@ -82,7 +82,10 @@ async def test_find_many_uses_pagination_and_excludes_stale(
         find_call = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args
         assert find_call[1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE
         assert find_call[1]["order"] == {"created_at": "asc"}
-        assert "stale_expired" in find_call[1]["where"]["status"]["not_in"]
+        not_in = find_call[1]["where"]["status"]["not_in"]
+        assert "stale_expired" in not_in
+        assert "complete" in not_in
+        assert "completed" in not_in
 
     @pytest.mark.asyncio
     async def test_fallback_query_used_when_batch_processed_missing(
@@ -148,33 +151,190 @@ async def test_fallback_completion_update_omits_batch_processed(
             return_value=0
         )
         mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
+        mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
+            return_value=None
+        )
 
         mock_job = MagicMock()
         mock_job.id = "job-fallback-1"
-        mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA=="  # base64-looking value
+        mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA=="
         mock_job.created_by = "user-1"
 
-        # Primary query fails → fallback path
+        # Simulate column already known absent (e.g. discovered on a previous cycle)
+        check_batch_cost_instance._has_batch_processed_column = False
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
-            side_effect=[Exception("column batch_processed does not exist"), [mock_job]]
+            return_value=[mock_job]
+        )
+
+        # Build a fake batch response whose status triggers the completion branch
+        mock_response = MagicMock()
+        mock_response.status = "completed"
+        mock_response.output_file_id = "file-output-123"
+        mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}'
+
+        mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
+        mock_llm_router.get_deployment_credentials_with_provider = MagicMock(
+            return_value={"api_key": "sk-test"}
         )
 
-        # Stub out the heavy per-job processing so we reach the update()
+        mock_deployment = MagicMock()
+        mock_deployment.litellm_params.custom_llm_provider = "openai"
+        mock_deployment.litellm_params.model = "gpt-4"
+        mock_deployment.model_info.model_dump.return_value = {}
+        mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment)
+
+        mock_file_content = MagicMock()
+        mock_file_content.content = b'{"id":"req-1"}'
+
+        decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;"
+
         with (
             patch(
                 "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id",
-                return_value=None,  # causes "not a valid unified object id" early-continue
+                side_effect=[decoded_id, None],
+            ),
+            patch(
+                "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id",
+                return_value="model-123",
             ),
+            patch(
+                "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id",
+                return_value="batch-456",
+            ),
+            patch(
+                "litellm.files.main.afile_content",
+                new_callable=AsyncMock,
+                return_value=mock_file_content,
+            ),
+            patch(
+                "litellm.batches.batch_utils._get_file_content_as_dictionary",
+                return_value=[{"id": "req-1"}],
+            ),
+            patch(
+                "litellm.batches.batch_utils.calculate_batch_cost_and_usage",
+                new_callable=AsyncMock,
+                return_value=(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"]),
+            ),
+            patch(
+                "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
+                return_value=("gpt-4", "openai", None, None),
+            ),
+            patch(
+                "litellm.litellm_core_utils.litellm_logging.Logging"
+            ) as mock_logging_cls,
         ):
+            mock_logging_obj = MagicMock()
+            mock_logging_obj.async_success_handler = AsyncMock()
+            mock_logging_cls.return_value = mock_logging_obj
+
             await check_batch_cost_instance.check_batch_cost()
 
-        # Even though the job was skipped (invalid ID), confirm the fallback path was taken
-        # by checking the find_many calls
-        find_calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list
-        assert len(find_calls) == 2
-        fallback_where = find_calls[1][1]["where"]
-        assert "batch_processed" not in fallback_where
+        # The update must have been called — this is the core assertion.
+        assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, (
+            "Expected update() to be called exactly once for the completed job"
+        )
+        update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"]
+        assert "batch_processed" not in update_data, (
+            "update() must NOT include batch_processed when column is absent"
+        )
+        assert update_data["status"] == "complete"
+
+    @pytest.mark.asyncio
+    async def test_primary_path_completion_update_includes_batch_processed(
+        self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
+    ):
+        """When batch_processed column IS present, completion update must set it to True.
 
-        # If a completion update were issued, it must not contain batch_processed
-        for call in mock_prisma_client.db.litellm_managedobjecttable.update.call_args_list:
-            assert "batch_processed" not in call[1].get("data", {})
+        This is the symmetric counterpart to test_fallback_completion_update_omits_batch_processed
+        and proves the conditional on _has_batch_processed_column governs the update data.
+        """
+        from unittest.mock import patch
+
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
+        mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
+        mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
+            return_value=None
+        )
+
+        mock_job = MagicMock()
+        mock_job.id = "job-primary-1"
+        mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA=="
+        mock_job.created_by = "user-1"
+
+        assert check_batch_cost_instance._has_batch_processed_column is True
+        mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
+            return_value=[mock_job]
+        )
+
+        mock_response = MagicMock()
+        mock_response.status = "completed"
+        mock_response.output_file_id = "file-output-123"
+        mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}'
+
+        mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
+        mock_llm_router.get_deployment_credentials_with_provider = MagicMock(
+            return_value={"api_key": "sk-test"}
+        )
+
+        mock_deployment = MagicMock()
+        mock_deployment.litellm_params.custom_llm_provider = "openai"
+        mock_deployment.litellm_params.model = "gpt-4"
+        mock_deployment.model_info.model_dump.return_value = {}
+        mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment)
+
+        mock_file_content = MagicMock()
+        mock_file_content.content = b'{"id":"req-1"}'
+
+        decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;"
+
+        with (
+            patch(
+                "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id",
+                side_effect=[decoded_id, None],
+            ),
+            patch(
+                "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id",
+                return_value="model-123",
+            ),
+            patch(
+                "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id",
+                return_value="batch-456",
+            ),
+            patch(
+                "litellm.files.main.afile_content",
+                new_callable=AsyncMock,
+                return_value=mock_file_content,
+            ),
+            patch(
+                "litellm.batches.batch_utils._get_file_content_as_dictionary",
+                return_value=[{"id": "req-1"}],
+            ),
+            patch(
+                "litellm.batches.batch_utils.calculate_batch_cost_and_usage",
+                new_callable=AsyncMock,
+                return_value=(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"]),
+            ),
+            patch(
+                "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
+                return_value=("gpt-4", "openai", None, None),
+            ),
+            patch(
+                "litellm.litellm_core_utils.litellm_logging.Logging"
+            ) as mock_logging_cls,
+        ):
+            mock_logging_obj = MagicMock()
+            mock_logging_obj.async_success_handler = AsyncMock()
+            mock_logging_cls.return_value = mock_logging_obj
+
+            await check_batch_cost_instance.check_batch_cost()
+
+        assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, (
+            "Expected update() to be called exactly once for the completed job"
+        )
+        update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"]
+        assert update_data["batch_processed"] is True, (
+            "update() must include batch_processed=True when column is present"
+        )
+        assert update_data["status"] == "complete"