fix(access_group): propagate MCP/agents and preserve direct models on access group sync

When an Access Group is assigned to a Team or Key (via create/update or
retroactively via the access-group endpoints), MCP servers and agents defined
on the group were never written to the team/key's LiteLLM_ObjectPermissionTable
row. Similarly, removing an access group overwrote the entity's entire models
list, silently deleting directly-assigned models that had nothing to do with
the group.

Changes:
- Add `_upsert_mcp_agents_in_object_permission` helper: creates or merges a
  LiteLLM_ObjectPermissionTable row with MCP server IDs and agent IDs from
  the access group, returning the upserted `object_permission_id`.
- Add `_remove_mcp_agents_from_object_permission` helper: removes only the
  specified server/agent IDs from an existing object_permission row.
- Fix `_sync_add_access_group_to_teams`: call the upsert helper when the
  access group contains MCP servers or agents; link the resulting
  `object_permission_id` to the team when a new row is created.
- Fix `_sync_add_access_group_to_keys`: extract `access_mcp_server_ids` /
  `access_agent_ids` (previously ignored) and call the upsert helper.
- Fix `_sync_remove_access_group_from_teams` / `_sync_remove_access_group_from_keys`:
  - Accept optional `removed_access_group_record` so callers can supply the
    pre-update snapshot (avoids using stale post-update data in update_access_group).
  - Only remove models/MCP/agents that were *exclusively* contributed by the
    removed group; models shared with remaining groups or directly assigned to
    the entity are preserved.
- Fix `delete_access_group`: delegate team/key cleanup to `_sync_remove_*`
  helpers (previously bypassed them for the main affected set, leaving
  models/MCP/agents dirty).
- Update `update_access_group` to pass `existing` (pre-update snapshot) as
  `removed_access_group_record` so the removal delta is always computed from
  the correct resource lists.
- Update and expand unit tests: fix stale test assertions, add new tests for
  MCP/agent propagation on add, direct-model preservation and exclusive-MCP
  cleanup on remove, and helper-function tests.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
RoyVivat 2026-04-07 23:30:02 -07:00
parent 50f88c8642
commit 6aad101b96
No known key found for this signature in database
GPG key ID: 59743472EC86530E
4 changed files with 1466 additions and 95 deletions

View file

@ -1,8 +1,9 @@
from typing import List, Set
from typing import Dict, List, Optional, Set, Tuple
from fastapi import APIRouter, Depends, HTTPException, status
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.proxy._types import (
CommonProxyErrors,
LiteLLM_AccessGroupTable,
@ -38,6 +39,57 @@ def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None:
)
async def _merge_access_group_resources_into_data_json(
data_json: dict,
access_group_ids: List[str],
prisma_client,
) -> dict:
"""
Batch-fetch access groups and merge their models/mcp_servers/agents into data_json.
- Models are merged into data_json["models"]
- MCP server IDs are merged into data_json["object_permission"]["mcp_servers"]
- Agent IDs are merged into data_json["object_permission"]["agents"]
This ensures that resources defined on an access group are propagated directly
to any team or key that references those access groups.
"""
if not access_group_ids:
return data_json
records = await prisma_client.db.litellm_accessgrouptable.find_many(
where={"access_group_id": {"in": access_group_ids}}
)
ag_models: List[str] = list(
{m for r in records for m in (r.access_model_names or [])}
)
ag_mcp_servers: List[str] = list(
{s for r in records for s in (r.access_mcp_server_ids or [])}
)
ag_agents: List[str] = list(
{a for r in records for a in (r.access_agent_ids or [])}
)
if ag_models:
existing_models: List[str] = list(data_json.get("models") or [])
data_json["models"] = list(set(existing_models + ag_models))
if ag_mcp_servers or ag_agents:
obj_perm: Dict = data_json.get("object_permission") or {}
if not isinstance(obj_perm, dict):
obj_perm = {}
if ag_mcp_servers:
existing_mcp: List[str] = list(obj_perm.get("mcp_servers") or [])
obj_perm["mcp_servers"] = list(set(existing_mcp + ag_mcp_servers))
if ag_agents:
existing_agents: List[str] = list(obj_perm.get("agents") or [])
obj_perm["agents"] = list(set(existing_agents + ag_agents))
data_json["object_permission"] = obj_perm
return data_json
def _record_to_response(record) -> AccessGroupResponse:
return AccessGroupResponse(
access_group_id=record.access_group_id,
@ -94,75 +146,375 @@ async def _invalidate_cache_access_group(access_group_id: str) -> None:
)
# ---------------------------------------------------------------------------
# Object-permission helpers (called inside a Prisma transaction)
# ---------------------------------------------------------------------------
async def _upsert_mcp_agents_in_object_permission(
tx,
existing_op_id: Optional[str],
ag_mcp_servers: List[str],
ag_agents: List[str],
) -> Optional[str]:
"""
Upsert LiteLLM_ObjectPermissionTable to add MCP servers and agents.
Merges ``ag_mcp_servers`` / ``ag_agents`` into the existing record (if any),
creating a new record when ``existing_op_id`` is None.
Returns the ``object_permission_id`` of the upserted row, or ``None`` when
both lists are empty (nothing to do).
"""
if not ag_mcp_servers and not ag_agents:
return None
existing_mcp: List[str] = []
existing_agents: List[str] = []
existing_data: Dict = {}
if existing_op_id:
existing_op = await tx.litellm_objectpermissiontable.find_unique(
where={"object_permission_id": existing_op_id}
)
if existing_op is not None:
try:
existing_data = existing_op.model_dump(exclude_none=True)
except Exception:
existing_data = existing_op.dict(exclude_none=True)
existing_mcp = list(existing_data.get("mcp_servers") or [])
existing_agents = list(existing_data.get("agents") or [])
upsert_data: Dict = {
k: v for k, v in existing_data.items() if k != "object_permission_id"
}
if ag_mcp_servers:
upsert_data["mcp_servers"] = list(set(existing_mcp + ag_mcp_servers))
if ag_agents:
upsert_data["agents"] = list(set(existing_agents + ag_agents))
op_id_to_use: str = existing_op_id or str(uuid.uuid4())
created_row = await tx.litellm_objectpermissiontable.upsert(
where={"object_permission_id": op_id_to_use},
data={"create": upsert_data, "update": upsert_data},
)
return created_row.object_permission_id
async def _remove_mcp_agents_from_object_permission(
tx,
existing_op_id: Optional[str],
mcp_servers_to_remove: List[str],
agents_to_remove: List[str],
) -> None:
"""
Remove specific MCP server IDs and agent IDs from an existing
LiteLLM_ObjectPermissionTable row. No-ops when the record does not exist
or the removal sets are empty.
"""
if not existing_op_id or (not mcp_servers_to_remove and not agents_to_remove):
return
existing_op = await tx.litellm_objectpermissiontable.find_unique(
where={"object_permission_id": existing_op_id}
)
if existing_op is None:
return
op_update: Dict = {}
if mcp_servers_to_remove:
remove_set: Set[str] = set(mcp_servers_to_remove)
op_update["mcp_servers"] = [
s for s in (existing_op.mcp_servers or []) if s not in remove_set
]
if agents_to_remove:
remove_set = set(agents_to_remove)
op_update["agents"] = [
a for a in (existing_op.agents or []) if a not in remove_set
]
if op_update:
await tx.litellm_objectpermissiontable.update(
where={"object_permission_id": existing_op_id},
data=op_update,
)
# ---------------------------------------------------------------------------
# DB sync helpers (called inside a Prisma transaction)
# ---------------------------------------------------------------------------
async def _sync_add_access_group_to_teams(
tx, team_ids: List[str], access_group_id: str
tx,
team_ids: List[str],
access_group_id: str,
access_group_record=None,
) -> None:
"""Add access_group_id to each team's access_group_ids (idempotent)."""
"""Add access_group_id to each team's access_group_ids and merge the group's
models/mcp_servers/agents into the team's direct resource lists (idempotent).
access_group_record: the Prisma record for the access group being added (optional).
When provided, its resources are merged directly into the team rather than
making an extra DB round-trip.
"""
ag_models: List[str] = list(
getattr(access_group_record, "access_model_names", None) or []
)
ag_mcp_servers: List[str] = list(
getattr(access_group_record, "access_mcp_server_ids", None) or []
)
ag_agents: List[str] = list(
getattr(access_group_record, "access_agent_ids", None) or []
)
for team_id in team_ids:
team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id})
if team is not None and access_group_id not in (team.access_group_ids or []):
await tx.litellm_teamtable.update(
where={"team_id": team_id},
data={
"access_group_ids": list(team.access_group_ids or [])
+ [access_group_id]
},
if team is None or access_group_id in (team.access_group_ids or []):
continue
update_data: Dict = {
"access_group_ids": list(team.access_group_ids or []) + [access_group_id]
}
# Merge models from the access group into the team's model list
if ag_models:
merged_models = list(set(list(team.models or []) + ag_models))
update_data["models"] = merged_models
# Merge MCP servers and agents into the team's object_permission
if ag_mcp_servers or ag_agents:
existing_op_id: Optional[str] = getattr(
team, "object_permission_id", None
)
new_op_id = await _upsert_mcp_agents_in_object_permission(
tx,
existing_op_id=existing_op_id,
ag_mcp_servers=ag_mcp_servers,
ag_agents=ag_agents,
)
# Link the (possibly newly created) object_permission row to the team
if new_op_id is not None and new_op_id != existing_op_id:
update_data["object_permission_id"] = new_op_id
await tx.litellm_teamtable.update(
where={"team_id": team_id},
data=update_data,
)
async def _sync_remove_access_group_from_teams(
tx, team_ids: List[str], access_group_id: str
tx,
team_ids: List[str],
access_group_id: str,
removed_access_group_record=None,
) -> None:
"""Remove access_group_id from each team's access_group_ids (idempotent)."""
"""Remove access_group_id from each team's access_group_ids and clean up
models / object_permission resources that were exclusively contributed by
the removed access group (idempotent).
removed_access_group_record: the Prisma record for the access group being
removed (optional). Pass the *pre-update* snapshot when calling from
``update_access_group`` so that stale post-update data is not used to
compute the removal delta. When ``None`` the record is fetched from
the DB (safe for the delete path where the row still exists).
"""
# Resolve removed AG's resources once outside the per-team loop.
ag_record = removed_access_group_record
if ag_record is None:
ag_record = await tx.litellm_accessgrouptable.find_unique(
where={"access_group_id": access_group_id}
)
removed_ag_models: Set[str] = set(
getattr(ag_record, "access_model_names", None) or []
)
removed_ag_mcp: Set[str] = set(
getattr(ag_record, "access_mcp_server_ids", None) or []
)
removed_ag_agents: Set[str] = set(
getattr(ag_record, "access_agent_ids", None) or []
)
for team_id in team_ids:
team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id})
if team is not None and access_group_id in (team.access_group_ids or []):
await tx.litellm_teamtable.update(
where={"team_id": team_id},
data={
"access_group_ids": [
ag for ag in team.access_group_ids if ag != access_group_id
]
},
)
if team is None or access_group_id not in (team.access_group_ids or []):
continue
remaining_ag_ids = [
ag for ag in (team.access_group_ids or []) if ag != access_group_id
]
# Batch-fetch remaining AGs to compute what resources they still provide
remaining_records = await tx.litellm_accessgrouptable.find_many(
where={"access_group_id": {"in": remaining_ag_ids}}
)
models_in_remaining: Set[str] = {
m for r in remaining_records for m in (r.access_model_names or [])
}
mcp_in_remaining: Set[str] = {
s for r in remaining_records for s in (r.access_mcp_server_ids or [])
}
agents_in_remaining: Set[str] = {
a for r in remaining_records for a in (r.access_agent_ids or [])
}
# Only remove models/MCP/agents that were *exclusively* from the removed AG
models_to_remove = removed_ag_models - models_in_remaining
current_models = set(team.models or [])
updated_models = list(current_models - models_to_remove)
update_data: Dict = {
"access_group_ids": remaining_ag_ids,
"models": updated_models,
}
# Clean up MCP servers / agents in object_permission
mcp_to_remove = list(removed_ag_mcp - mcp_in_remaining)
agents_to_remove = list(removed_ag_agents - agents_in_remaining)
existing_op_id: Optional[str] = getattr(team, "object_permission_id", None)
await _remove_mcp_agents_from_object_permission(
tx,
existing_op_id=existing_op_id,
mcp_servers_to_remove=mcp_to_remove,
agents_to_remove=agents_to_remove,
)
await tx.litellm_teamtable.update(
where={"team_id": team_id},
data=update_data,
)
async def _sync_add_access_group_to_keys(
tx, key_tokens: List[str], access_group_id: str
tx, key_tokens: List[str], access_group_id: str, access_group_record=None
) -> None:
"""Add access_group_id to each key's access_group_ids (idempotent)."""
"""Add access_group_id to each key's access_group_ids and merge the group's
models/mcp_servers/agents into the key's direct resource lists (idempotent).
"""
ag_models: List[str] = list(
getattr(access_group_record, "access_model_names", None) or []
)
ag_mcp_servers: List[str] = list(
getattr(access_group_record, "access_mcp_server_ids", None) or []
)
ag_agents: List[str] = list(
getattr(access_group_record, "access_agent_ids", None) or []
)
for token in key_tokens:
key = await tx.litellm_verificationtoken.find_unique(where={"token": token})
if key is not None and access_group_id not in (key.access_group_ids or []):
await tx.litellm_verificationtoken.update(
where={"token": token},
data={
"access_group_ids": list(key.access_group_ids or [])
+ [access_group_id]
},
if key is None or access_group_id in (key.access_group_ids or []):
continue
update_data: Dict = {
"access_group_ids": list(key.access_group_ids or []) + [access_group_id]
}
if ag_models:
merged_models = list(set(list(key.models or []) + ag_models))
update_data["models"] = merged_models
# Merge MCP servers and agents into the key's object_permission
if ag_mcp_servers or ag_agents:
existing_op_id: Optional[str] = getattr(
key, "object_permission_id", None
)
new_op_id = await _upsert_mcp_agents_in_object_permission(
tx,
existing_op_id=existing_op_id,
ag_mcp_servers=ag_mcp_servers,
ag_agents=ag_agents,
)
# Link the (possibly newly created) object_permission row to the key
if new_op_id is not None and new_op_id != existing_op_id:
update_data["object_permission_id"] = new_op_id
await tx.litellm_verificationtoken.update(
where={"token": token},
data=update_data,
)
async def _sync_remove_access_group_from_keys(
tx, key_tokens: List[str], access_group_id: str
tx,
key_tokens: List[str],
access_group_id: str,
removed_access_group_record=None,
) -> None:
"""Remove access_group_id from each key's access_group_ids (idempotent)."""
"""Remove access_group_id from each key's access_group_ids and clean up
models / object_permission resources that were exclusively contributed by
the removed access group (idempotent).
removed_access_group_record: the Prisma record for the access group being
removed (optional). Pass the *pre-update* snapshot when calling from
``update_access_group`` so that stale post-update data is not used to
compute the removal delta. When ``None`` the record is fetched from
the DB (safe for the delete path where the row still exists).
"""
# Resolve removed AG's resources once outside the per-key loop.
ag_record = removed_access_group_record
if ag_record is None:
ag_record = await tx.litellm_accessgrouptable.find_unique(
where={"access_group_id": access_group_id}
)
removed_ag_models: Set[str] = set(
getattr(ag_record, "access_model_names", None) or []
)
removed_ag_mcp: Set[str] = set(
getattr(ag_record, "access_mcp_server_ids", None) or []
)
removed_ag_agents: Set[str] = set(
getattr(ag_record, "access_agent_ids", None) or []
)
for token in key_tokens:
key = await tx.litellm_verificationtoken.find_unique(where={"token": token})
if key is not None and access_group_id in (key.access_group_ids or []):
await tx.litellm_verificationtoken.update(
where={"token": token},
data={
"access_group_ids": [
ag for ag in key.access_group_ids if ag != access_group_id
]
},
)
if key is None or access_group_id not in (key.access_group_ids or []):
continue
remaining_ag_ids = [
ag for ag in (key.access_group_ids or []) if ag != access_group_id
]
# Batch-fetch remaining AGs to compute what resources they still provide
remaining_records = await tx.litellm_accessgrouptable.find_many(
where={"access_group_id": {"in": remaining_ag_ids}}
)
models_in_remaining: Set[str] = {
m for r in remaining_records for m in (r.access_model_names or [])
}
mcp_in_remaining: Set[str] = {
s for r in remaining_records for s in (r.access_mcp_server_ids or [])
}
agents_in_remaining: Set[str] = {
a for r in remaining_records for a in (r.access_agent_ids or [])
}
# Only remove models/MCP/agents that were *exclusively* from the removed AG
models_to_remove = removed_ag_models - models_in_remaining
current_models = set(key.models or [])
updated_models = list(current_models - models_to_remove)
# Clean up MCP servers / agents in object_permission
mcp_to_remove = list(removed_ag_mcp - mcp_in_remaining)
agents_to_remove = list(removed_ag_agents - agents_in_remaining)
existing_op_id: Optional[str] = getattr(key, "object_permission_id", None)
await _remove_mcp_agents_from_object_permission(
tx,
existing_op_id=existing_op_id,
mcp_servers_to_remove=mcp_to_remove,
agents_to_remove=agents_to_remove,
)
await tx.litellm_verificationtoken.update(
where={"token": token},
data={
"access_group_ids": remaining_ag_ids,
"models": updated_models,
},
)
# ---------------------------------------------------------------------------
@ -329,11 +681,18 @@ async def create_access_group(
)
# Sync team and key tables to reference the new access group
# Pass the created record so sync can merge models/MCP/agents directly.
await _sync_add_access_group_to_teams(
tx, data.assigned_team_ids or [], record.access_group_id
tx,
data.assigned_team_ids or [],
record.access_group_id,
access_group_record=record,
)
await _sync_add_access_group_to_keys(
tx, data.assigned_key_ids or [], record.access_group_id
tx,
data.assigned_key_ids or [],
record.access_group_id,
access_group_record=record,
)
except HTTPException:
raise
@ -481,13 +840,26 @@ async def update_access_group(
data=update_data,
)
await _sync_add_access_group_to_teams(tx, teams_to_add, access_group_id)
await _sync_remove_access_group_from_teams(
tx, teams_to_remove, access_group_id
await _sync_add_access_group_to_teams(
tx, teams_to_add, access_group_id, access_group_record=record
)
# Pass `existing` (pre-update snapshot) so remove logic uses the
# OLD resource lists when computing the removal delta, not the
# post-update ones.
await _sync_remove_access_group_from_teams(
tx,
teams_to_remove,
access_group_id,
removed_access_group_record=existing,
)
await _sync_add_access_group_to_keys(
tx, keys_to_add, access_group_id, access_group_record=record
)
await _sync_add_access_group_to_keys(tx, keys_to_add, access_group_id)
await _sync_remove_access_group_from_keys(
tx, keys_to_remove, access_group_id
tx,
keys_to_remove,
access_group_id,
removed_access_group_record=existing,
)
except HTTPException:
raise
@ -566,44 +938,21 @@ async def delete_access_group(
} | set(existing.assigned_key_ids or [])
affected_key_tokens = list(all_affected_key_tokens)
# Update teams returned by find_many directly — we already have their data.
for team in teams_with_group:
await tx.litellm_teamtable.update(
where={"team_id": team.team_id},
data={
"access_group_ids": [
ag
for ag in (team.access_group_ids or [])
if ag != access_group_id
]
},
)
# Use _sync_remove only for out-of-sync teams not found by the hasSome query.
out_of_sync_team_ids = set(existing.assigned_team_ids or []) - {
t.team_id for t in teams_with_group
}
# Use _sync_remove for ALL affected teams — it correctly handles
# model/MCP/agent cleanup and is idempotent. The `existing` record
# is passed as the pre-delete snapshot so that removal deltas are
# computed from the right resource lists before the row is deleted.
await _sync_remove_access_group_from_teams(
tx, list(out_of_sync_team_ids), access_group_id
tx,
affected_team_ids,
access_group_id,
removed_access_group_record=existing,
)
# Update keys returned by find_many directly — we already have their data.
for key in keys_with_group:
await tx.litellm_verificationtoken.update(
where={"token": key.token},
data={
"access_group_ids": [
ag
for ag in (key.access_group_ids or [])
if ag != access_group_id
]
},
)
# Use _sync_remove only for out-of-sync keys not found by the hasSome query.
out_of_sync_key_tokens = set(existing.assigned_key_ids or []) - {
k.token for k in keys_with_group
}
await _sync_remove_access_group_from_keys(
tx, list(out_of_sync_key_tokens), access_group_id
tx,
affected_key_tokens,
access_group_id,
removed_access_group_record=existing,
)
await tx.litellm_accessgrouptable.delete(

View file

@ -56,6 +56,9 @@ from litellm.proxy.management_endpoints.common_utils import (
_is_user_team_admin,
_set_object_metadata_field,
)
from litellm.proxy.management_endpoints.access_group_endpoints import (
_merge_access_group_resources_into_data_json,
)
from litellm.proxy.management_endpoints.model_management_endpoints import (
_add_model_to_db,
)
@ -672,6 +675,15 @@ async def _common_key_generation_helper( # noqa: PLR0915
data_json.pop("tags")
# Populate models/MCP servers/agents from key-level access groups (if provided).
# Only runs when access_group_ids is explicitly included in the request.
if data_json.get("access_group_ids"):
data_json = await _merge_access_group_resources_into_data_json(
data_json=data_json,
access_group_ids=data_json["access_group_ids"],
prisma_client=prisma_client,
)
# Validate MCP servers in object_permission are within team scope
await validate_key_mcp_servers_against_team(
object_permission=data_json.get("object_permission"),
@ -1559,6 +1571,20 @@ async def prepare_key_update_data(
non_default_values["budget_reset_at"] = key_reset_at
non_default_values["budget_duration"] = budget_duration
# Populate models/MCP servers/agents from key-level access groups when
# access_group_ids is explicitly provided in the update request.
if "access_group_ids" in non_default_values:
new_access_group_ids = non_default_values.get("access_group_ids") or []
if new_access_group_ids:
from litellm.proxy.proxy_server import prisma_client as _prisma_client
if _prisma_client is not None:
non_default_values = await _merge_access_group_resources_into_data_json(
data_json=non_default_values,
access_group_ids=new_access_group_ids,
prisma_client=_prisma_client,
)
if "object_permission" in non_default_values:
non_default_values = await _handle_update_object_permission(
data_json=non_default_values,

View file

@ -78,6 +78,9 @@ from litellm.proxy.management_endpoints.common_utils import (
_upsert_budget_and_membership,
_user_has_admin_view,
)
from litellm.proxy.management_endpoints.access_group_endpoints import (
_merge_access_group_resources_into_data_json,
)
from litellm.proxy.management_endpoints.tag_management_endpoints import (
get_daily_activity,
)
@ -925,6 +928,17 @@ async def new_team( # noqa: PLR0915
## Create Team Member Budget Table
data_json = data.json()
## Populate models/MCP servers/agents from access groups (if any)
## This ensures resources defined on access groups are reflected directly on the team,
## so they appear in model lists and MCP validation at key-generation time.
## Only runs when access_group_ids is explicitly provided in this request.
if data.access_group_ids:
data_json = await _merge_access_group_resources_into_data_json(
data_json=data_json,
access_group_ids=data.access_group_ids,
prisma_client=prisma_client,
)
## Handle Object Permission - MCP, Vector Stores etc.
data_json = await _set_object_permission(
data_json=data_json,
@ -1515,8 +1529,20 @@ async def update_team( # noqa: PLR0915
else:
TeamMemberBudgetHandler._clean_team_member_fields(updated_kv)
# Check object permission
if data.object_permission is not None:
# Populate models/MCP servers/agents from access groups when access_group_ids is
# being explicitly updated. Only runs when the caller provides access_group_ids.
if "access_group_ids" in updated_kv:
new_access_group_ids = updated_kv.get("access_group_ids") or []
if new_access_group_ids:
updated_kv = await _merge_access_group_resources_into_data_json(
data_json=updated_kv,
access_group_ids=new_access_group_ids,
prisma_client=prisma_client,
)
# Check object permission — fire when explicitly set OR when access group
# resolution added mcp_servers/agents to updated_kv["object_permission"].
if data.object_permission is not None or "object_permission" in updated_kv:
updated_kv = await handle_update_object_permission(
data_json=updated_kv,
existing_team_row=existing_team_row,