fix: address Greptile review feedback on policy resolve endpoints

- Track unnamed keys/teams as separate counts instead of inflating
  affected_keys_count with duplicate "(unnamed key)" placeholders.
  Added unnamed_keys_count and unnamed_teams_count to response.
- Push alias pattern matching to DB via _build_alias_where() which
  converts exact patterns to Prisma "in" and suffix wildcards to
  "startsWith" filters.
- Gate sync_policies_from_db/sync_attachments_from_db behind
  force_sync query param (default false) to avoid 2 DB round-trips
  on every /policies/resolve request.
- Remove worktree-only conftest.py that cleared sys.modules at import
  time — no longer needed since code moved to main repo.
- Rename MAX_ESTIMATE_IMPACT_ROWS → MAX_POLICY_ESTIMATE_IMPACT_ROWS.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Ishaan Jaffer 2026-02-10 17:23:15 -08:00
parent 0eeddcc8be
commit e9ff805c26
3 changed files with 56 additions and 59 deletions

View file

@ -26,9 +26,6 @@ from litellm.types.proxy.policy_engine import (
router = APIRouter()
_UNNAMED_KEY_PLACEHOLDER = "(unnamed key)"
_UNNAMED_TEAM_PLACEHOLDER = "(unnamed team)"
def _build_alias_where(field: str, patterns: list) -> dict:
"""Build a Prisma ``where`` clause for alias patterns.
@ -78,11 +75,15 @@ def _get_tags_from_metadata(metadata: object, json_metadata: object = None) -> l
async def _find_affected_keys_by_tags(
prisma_client: object, tag_patterns: list
) -> list:
"""Find key aliases whose metadata.tags match any of the given patterns."""
) -> tuple:
"""Find key aliases whose metadata.tags match any of the given patterns.
Returns (named_aliases, unnamed_count).
"""
from litellm.proxy.auth.route_checks import RouteChecks
affected: list = []
unnamed_count = 0
keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore
where={}, order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS,
)
@ -96,17 +97,24 @@ async def _find_affected_keys_by_tags(
for tag in key_tags
for pat in tag_patterns
):
affected.append(key_alias or _UNNAMED_KEY_PLACEHOLDER)
return affected
if key_alias:
affected.append(key_alias)
else:
unnamed_count += 1
return affected, unnamed_count
async def _find_affected_teams_by_tags(
prisma_client: object, tag_patterns: list
) -> list:
"""Find team aliases whose metadata.tags match any of the given patterns."""
) -> tuple:
"""Find team aliases whose metadata.tags match any of the given patterns.
Returns (named_aliases, unnamed_count).
"""
from litellm.proxy.auth.route_checks import RouteChecks
affected: list = []
unnamed_count = 0
teams = await prisma_client.db.litellm_teamtable.find_many( # type: ignore
where={}, order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS,
)
@ -118,14 +126,20 @@ async def _find_affected_teams_by_tags(
for tag in team_tags
for pat in tag_patterns
):
affected.append(team_alias or _UNNAMED_TEAM_PLACEHOLDER)
return affected
if team_alias:
affected.append(team_alias)
else:
unnamed_count += 1
return affected, unnamed_count
async def _find_affected_by_team_patterns(
prisma_client: object, team_patterns: list, existing_teams: list, existing_keys: list
) -> tuple:
"""Find teams matching alias patterns and keys belonging to those teams."""
"""Find teams matching alias patterns and keys belonging to those teams.
Returns (new_teams, new_keys, unnamed_keys_count).
"""
from litellm.proxy.auth.route_checks import RouteChecks
new_teams: list = []
@ -146,17 +160,21 @@ async def _find_affected_by_team_patterns(
matched_team_ids.append(str(team.team_id))
new_keys: list = []
unnamed_keys_count = 0
if matched_team_ids:
keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore
where={"team_id": {"in": matched_team_ids}},
order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS,
)
for key in keys:
key_alias = key.key_alias or _UNNAMED_KEY_PLACEHOLDER
if key_alias not in existing_keys:
new_keys.append(key_alias)
key_alias = key.key_alias or ""
if key_alias:
if key_alias not in existing_keys:
new_keys.append(key_alias)
else:
unnamed_keys_count += 1
return new_teams, new_keys
return new_teams, new_keys, unnamed_keys_count
async def _find_affected_keys_by_alias(
@ -338,21 +356,28 @@ async def estimate_attachment_impact(
affected_keys: list = []
affected_teams: list = []
unnamed_keys = 0
unnamed_teams = 0
# Tag-based impact
tag_patterns = request.tags or []
if tag_patterns:
affected_keys = await _find_affected_keys_by_tags(prisma_client, tag_patterns)
affected_teams = await _find_affected_teams_by_tags(prisma_client, tag_patterns)
affected_keys, unnamed_keys = await _find_affected_keys_by_tags(
prisma_client, tag_patterns,
)
affected_teams, unnamed_teams = await _find_affected_teams_by_tags(
prisma_client, tag_patterns,
)
# Team-based impact (alias matching + keys belonging to those teams)
team_patterns = request.teams or []
if team_patterns:
new_teams, new_keys = await _find_affected_by_team_patterns(
new_teams, new_keys, new_unnamed = await _find_affected_by_team_patterns(
prisma_client, team_patterns, affected_teams, affected_keys,
)
affected_teams.extend(new_teams)
affected_keys.extend(new_keys)
unnamed_keys += new_unnamed
# Key-based impact (direct alias matching)
key_patterns = request.keys or []
@ -363,8 +388,10 @@ async def estimate_attachment_impact(
affected_keys.extend(new_keys)
return AttachmentImpactResponse(
affected_keys_count=len(affected_keys),
affected_teams_count=len(affected_teams),
affected_keys_count=len(affected_keys) + unnamed_keys,
affected_teams_count=len(affected_teams) + unnamed_teams,
unnamed_keys_count=unnamed_keys,
unnamed_teams_count=unnamed_teams,
sample_keys=affected_keys[:10],
sample_teams=affected_teams[:10],
)

View file

@ -343,10 +343,16 @@ class AttachmentImpactResponse(BaseModel):
"""Response for estimating the impact of a policy attachment."""
affected_keys_count: int = Field(
default=0, description="Number of keys that would be affected."
default=0, description="Number of keys that would be affected (named + unnamed)."
)
affected_teams_count: int = Field(
default=0, description="Number of teams that would be affected."
default=0, description="Number of teams that would be affected (named + unnamed)."
)
unnamed_keys_count: int = Field(
default=0, description="Number of affected keys without an alias."
)
unnamed_teams_count: int = Field(
default=0, description="Number of affected teams without an alias."
)
sample_keys: List[str] = Field(
default_factory=list,

View file

@ -1,36 +0,0 @@
"""
Local conftest for policy_engine tests.
The parent conftest (tests/test_litellm/conftest.py) inserts the main repo
root into sys.path and imports litellm from there. When running from a
worktree, this means all litellm imports resolve to the main repo's code
instead of the worktree's.
This conftest fixes the path and clears all cached litellm modules so
subsequent imports resolve from the worktree.
"""
import os
import sys
import pytest
# Fix sys.path: insert this worktree's root FIRST and remove the main repo root.
_this_dir = os.path.dirname(os.path.abspath(__file__))
_worktree_root = os.path.abspath(os.path.join(_this_dir, "..", "..", "..", ".."))
sys.path.insert(0, _worktree_root)
# Remove the main repo path that parent conftest inserted
_main_repo = os.path.abspath(os.path.join(_worktree_root, "..", ".."))
sys.path = [p for p in sys.path if os.path.abspath(p) != _main_repo]
# Clear ALL cached litellm modules so they're re-imported from the worktree
_to_remove = [key for key in sys.modules if key == "litellm" or key.startswith("litellm.")]
for key in _to_remove:
del sys.modules[key]
@pytest.fixture(scope="module", autouse=True)
def setup_and_teardown():
"""Override parent conftest - policy engine tests don't need litellm reload."""
yield