[April 6th] - Ishaan

[ishaan-berri]

…ATE (#25227)

• Add STALE_OBJECT_CLEANUP_BATCH_SIZE constant

Configurable batch limit (default 1000) for stale managed object cleanup, preventing unbounded UPDATE queries from hitting 300K+ rows at once.

• Batch-limit stale managed object cleanup with single bounded SQL query

Two fixes to _cleanup_stale_managed_objects:

1. Replace unbounded update_many with a single execute_raw using a subquery LIMIT, capping each poll cycle to STALE_OBJECT_CLEANUP_BATCH_SIZE rows. Zero rows loaded into Python memory — everything stays in Postgres. Uses the same PostgreSQL raw-SQL pattern as spend_log_cleanup.py (the proxy requires PostgreSQL per schema.prisma).

2. Extract _expire_stale_rows as a separate method for testability.

Keeps the file_purpose='response' filter to avoid incorrectly expiring long-running batch or fine-tune jobs that legitimately exceed the staleness cutoff.

Relevant issues

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have Added testing in the tests/test_litellm/ directory, Adding at least 1 test is a hard requirement - see details
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible, it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

CI (LiteLLM team)

CI status guideline:

• 50-55 passing tests: main is stable with minor issues.
• 45-49 passing tests: acceptable but needs attention
• <= 40 passing tests: unstable; be careful with your merges and assess the risk.

• Branch creation CI run
  Link:

• CI run for the last commit
  Link:

• Merge / cherry-pick CI run
  Links:

Type

🆕 New Feature
🐛 Bug Fix
🧹 Refactoring
📖 Documentation
🚄 Infrastructure
✅ Test

Changes

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
litellm	Ready	Preview, Comment	Apr 7, 2026 0:27am

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 3 committers have signed the CLA.

✅ otaviofbrito
❌ ishaan-berri
❌ github-actions[bot]
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[codspeed-hq]

Merging this PR will not alter performance

✅ 16 untouched benchmarks

---

_{Comparing litellm_ishaan_april6 (1c238b6) with main (d132b1b)}

[greptile-apps]

Greptile Summary

This PR adds three related features: (1) per-entity multi-window budget tracking (budget_limits) for both keys and teams, (2) per-team-member model scoping (allowed_models on the budget table, default_team_member_models on the team table), and (3) a batched stale managed-object cleanup for the responses API.

Several issues were flagged in the prior review round and remain open (merge conflict markers in test_bedrock_common_utils.py, budget_limits absent from all three schema.prisma files, raw SQL in _expire_stale_rows, tools/system dropped from the count_tokens call). This round surfaces one new finding:

• reset_budget_windows (P1): The new reset_budget_jobs.py method loads all keys and teams with budget_limits in a single unbounded find_many (no take/skip) and issues one update() per dirty record inside the loop. Both patterns violate CLAUDE.md's DB rules (bound large result sets with cursor pagination; batch writes rather than per-row calls)."

Confidence Score: 4/5

Not ready to merge — open P0/P1 findings from prior rounds (merge conflict breaks test file, budget_limits missing from Prisma schema breaks all multi-budget-window runtime paths) plus a new P1 unbounded-query/N+1 write pattern in reset_budget_windows.

Score of 4 is appropriate: there is one confirmed new P1 finding (unbounded find_many + N+1 updates in reset_budget_windows), and prior-round P1/P0 items (merge conflict markers in test file, budget_limits schema drift, raw SQL) are still open per file inspection. Multiple P1s keep the score at 4 rather than 5.

litellm/proxy/common_utils/reset_budget_job.py (new N+1 / unbounded query), tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py (merge conflict), litellm/proxy/schema.prisma + schema.prisma + litellm-proxy-extras/litellm_proxy_extras/schema.prisma (budget_limits column missing), enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py (raw SQL), litellm/proxy/proxy_server.py (tools/system dropped from count_tokens)

Important Files Changed

Filename	Overview
litellm/proxy/common_utils/reset_budget_job.py	Added reset_budget_windows() with unbounded find_many (no take/skip) and N+1 individual update calls inside a for-loop, violating CLAUDE.md DB rules for large result sets and batch writes
enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py	Extracted _expire_stale_rows using raw SQL execute_raw with batch LIMIT; raw SQL pattern flagged in prior review as violating CLAUDE.md
litellm/proxy/_types.py	Added BudgetLimitEntry, budget_limits on GenerateRequestBase/TeamBase/UpdateTeamRequest/LiteLLM_VerificationToken, allowed_models on LiteLLM_BudgetTable; all syntactically valid in current HEAD
litellm/proxy/auth/auth_checks.py	Added _check_team_member_model_access, _team_multi_budget_check, _virtual_key_multi_budget_check; all correctly follow the existing get_team_membership caching pattern
litellm/proxy/proxy_server.py	Added per-window spend counter increments for budget_limits keys/teams; tools/system params dropped from count_tokens call (flagged in prior review)
tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py	Unresolved git merge conflict markers at lines 37-43 prevent file from parsing as valid Python — all tests in module fail at collection time (flagged in prior review)
litellm/constants.py	Added STALE_OBJECT_CLEANUP_BATCH_SIZE constant with env-var override, max(1,...) guard, and sensible default of 1000; clean addition
litellm-proxy-extras/litellm_proxy_extras/schema.prisma	Added allowed_models to LiteLLM_BudgetTable and default_team_member_models to LiteLLM_TeamTable; budget_limits still absent from both LiteLLM_TeamTable and LiteLLM_VerificationToken (flagged in prior review)

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Incoming Request] --> B[user_api_key_auth]
    B --> C[common_checks]
    C --> D{team_object present?}
    D -- Yes --> E[_team_max_budget_check]
    D -- Yes --> F[_team_multi_budget_check]
    F --> F1[get_current_spend\nspend:team:ID:window:DURATION]
    F1 --> F2{spend >= max_budget?}
    F2 -- Yes --> ERR[BudgetExceededError]
    D -- Yes, with user_id --> G[_check_team_member_model_access]
    G --> G1[get_team_membership\ncache-first DB lookup]
    G1 --> G2{allowed_models non-empty?}
    G2 -- Yes --> G3[_can_object_call_model]
    G3 -- denied --> ERR
    C --> H{valid_token present?}
    H -- Yes --> I[_virtual_key_multi_budget_check]
    I --> I1[get_current_spend\nspend:key:TOKEN:window:DURATION]
    I1 --> I2{spend >= max_budget?}
    I2 -- Yes --> ERR
    C --> J[Request proceeds]

    K[Response complete] --> L[increment_spend_counters]
    L --> M[increment key spend counter]
    L --> N[increment per-window key counters\nspend:key:TOKEN:window:DURATION]
    L --> O[increment team spend counter]
    L --> P[increment per-window team counters\nspend:team:ID:window:DURATION]

    Q[ResetBudgetJob poll] --> R[reset_budget_windows]
    R --> R1[find_many keys with budget_limits\n⚠️ no take/skip limit]
    R1 --> R2[for each expired window\nreset Redis counter\nupdate DB one-by-one\n⚠️ N+1 writes]
    R --> R3[find_many teams with budget_limits\n⚠️ no take/skip limit]
    R3 --> R4[for each expired window\nreset Redis counter\nupdate DB one-by-one\n⚠️ N+1 writes]

Loading

_{Reviews (6): Last reviewed commit: "fix(tests): update upsert tests to refle..." | Re-trigger Greptile}

[greptile-apps]

Raw SQL bypasses ORM layer

The CLAUDE.md rule says: "Do not write raw SQL for proxy DB operations. Use Prisma model methods instead of execute_raw / query_raw". A Prisma-native implementation avoids hand-written SQL, keeps the code testable with simple mocks, and removes schema-drift risk — while still bounding the batch:

async def _expire_stale_rows(self, cutoff: datetime, batch_size: int) -> int:
    stale = await self.prisma_client.db.litellm_managedobjecttable.find_many(
        where={
            "file_purpose": "response",
            "status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
            "created_at": {"lt": cutoff},
        },
        order={"created_at": "asc"},
        take=batch_size,
        select={"id": True},
    )
    if not stale:
        return 0
    await self.prisma_client.db.litellm_managedobjecttable.update_many(
        where={"id": {"in": [r.id for r in stale]}},
        data={"status": "stale_expired"},
    )
    return len(stale)

This is two DB round-trips instead of one, but it stays in the ORM layer and matches how the rest of the proxy interacts with the DB. Note: spend_log_cleanup.py also uses execute_raw as a precedent for DELETE … WHERE … IN (SELECT … LIMIT n) — this is a style nudge rather than a hard blocker, but worth aligning with the stated convention.

Context Used: CLAUDE.md (source)

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

[greptile-apps]

No tests added for the new method

The PR's pre-submission checklist shows the test checkbox unchecked, and CLAUDE.md states: "Adding at least 1 test is a hard requirement". The _expire_stale_rows docstring explicitly calls out that it was "Isolated so it can be swapped / mocked in tests without touching the orchestration logic" — but no tests were added.

At minimum, a unit test in tests/test_litellm/ that mocks prisma_client.db.execute_raw should verify:

1. The batch cap (STALE_OBJECT_CLEANUP_BATCH_SIZE) is passed correctly.
2. The affected-row count is returned and triggers the warning log.
3. Zero rows → no warning emitted.

Without tests the batch-size guard and the refactored flow cannot be automatically regressed against.

[greptile-apps]

tools and system parameters silently dropped from token counting call

The refactor to add the httpx.HTTPStatusError handler also accidentally removed tools=tools and system=system from the count_tokens call. Both variables are captured from the request at lines 9064-9065 and are actively used by the Anthropic and Gemini token counters to include tool definitions and system prompts in the server-side token count.

With this change:

• Requests that include tools will receive an underestimated token count (tool definitions can be hundreds/thousands of tokens)
• Requests that include a system prompt will similarly receive an undercount

The BaseTokenCounter.count_tokens abstract signature explicitly accepts both params (tools: Optional[List[Dict[str, Any]]] = None, system: Optional[Any] = None), and AnthropicTokenCounter passes them straight through to anthropic_count_tokens_handler.handle_count_tokens_request.

Suggested change

	try:
	result = await provider_counter.count_tokens(
	model_to_use=model_to_use or "",
	messages=messages, # type: ignore
	contents=contents,
	deployment=deployment,
	request_model=request.model,
	)
	try:
	result = await provider_counter.count_tokens(
	model_to_use=model_to_use or "",
	messages=messages, # type: ignore
	contents=contents,
	deployment=deployment,
	request_model=request.model,
	tools=tools,
	system=system,
	)

[greptile-apps]

Unresolved git merge conflict markers

Lines 37–43 contain live conflict markers (<<<<<<< worktree-rustling-wishing-kite, =======, >>>>>>> main). Python cannot parse this file, so every test in test_bedrock_common_utils.py fails with a SyntaxError at import time.

The main-branch version is correct — the model under test is bedrock/us-gov.anthropic.claude-haiku-4-5-20251001-v1:0, so stripping the us-gov. cross-region prefix should yield anthropic.claude-haiku-4-5-20251001-v1:0. The worktree assertion (anthropic.claude-3-5-sonnet-20240620-v1:0) refers to a completely different model and would be wrong.

Resolve by keeping only the correct assertion:

Suggested change

	<<<<<<< worktree-rustling-wishing-kite
	assert base_model == "anthropic.claude-3-5-sonnet-20240620-v1:0"
	=======
	assert base_model == "anthropic.claude-haiku-4-5-20251001-v1:0"
	>>>>>>> main
	assert base_model == "anthropic.claude-haiku-4-5-20251001-v1:0"

[greptile-apps]

Syntax Error: missing closing ) on budget_limits field

The budget_limits field opens a parenthesis on line 1764 that is never closed before default_team_member_models begins on line 1766. The only ) on line 1768 is intended for default_team_member_models, so Python's implicit line continuation keeps the budget_limits expression open and sees an annotation statement (default_team_member_models: ...) inside a value expression — which raises SyntaxError: invalid syntax at import time. This makes the entire proxy unlaunchable.

Suggested change

	budget_limits: Optional[List[BudgetLimitEntry]] = (
	None # multiple concurrent budget windows
	default_team_member_models: Optional[List[str]] = (
	None # default allowed_models seeded onto new team members
	)
	budget_limits: Optional[List[BudgetLimitEntry]] = (
	None # multiple concurrent budget windows
	)
	default_team_member_models: Optional[List[str]] = (
	None # default allowed_models seeded onto new team members
	)

[greptile-apps]

budget_limits absent from all Prisma schema files — feature broken at runtime

The migration 20260401000000_add_budget_limits/migration.sql correctly adds budget_limits JSONB to LiteLLM_VerificationToken and LiteLLM_TeamTable in the database. However, none of the three Prisma schema files (schema.prisma, litellm/proxy/schema.prisma, litellm-proxy-extras/litellm_proxy_extras/schema.prisma) declare this field for those two models — confirmed by reading the current file, neither model block contains budget_limits.

Because Prisma generates its client from the schema (not from live DB introspection), every budget_limits reference in reset_budget_windows() will fail at runtime:

• find_many(where={"budget_limits": {"not": None}}) — Prisma rejects unknown filter fields with a validation error
• key.budget_limits / team.budget_limits attribute accesses — returned objects never carry this field (always None or AttributeError)
• update(data={"budget_limits": json.dumps(windows)}) — Prisma rejects unknown write fields

All of these failures are silently swallowed by the surrounding except Exception blocks, so budget windows will never reset and per-window spend counters will accumulate indefinitely. The # type: ignore[arg-type] / # type: ignore[attr-defined] annotations throughout confirm the author is aware the field is absent from the Prisma schema.

Fix: Add budget_limits Json? to LiteLLM_VerificationToken and LiteLLM_TeamTable in all three schema.prisma files and regenerate the Prisma client. Per CLAUDE.md, schema changes must be kept in sync across all copies.