Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/management_helpers/object_permission_utils.py: 54%

1""" 

2Common utility functions for handling object permission updates across 

3organizations, teams, and keys. 

4""" 

5 

6import json 

7from collections.abc import Iterable, Mapping, Sequence 

8from collections.abc import Set as AbstractSet 

9from dataclasses import dataclass 

10from types import MappingProxyType 

11from typing import TYPE_CHECKING, Final, Optional 

12 

13from fastapi import HTTPException, status 

14from pydantic import TypeAdapter 

15 

16from litellm._logging import verbose_proxy_logger 

17from litellm._uuid import uuid 

18from litellm.litellm_core_utils.safe_json_dumps import safe_dumps 

19from litellm.proxy._types import ObjectPermissionDict, SpecialMCPServerName, SpecialMCPServerNames 

20from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast 

21from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, object_permission_cache_key 

22from litellm.proxy.utils import PrismaClient 

23from litellm.repositories.object_permission_repository import ObjectPermissionRepository 

24from litellm.repositories.table_repositories import MCPServerRepository 

25 

26if TYPE_CHECKING: 26 ↛ 27line 26 didn't jump to line 27 because the condition on line 26 was never true

27 from prisma import models as prisma_models 

28 

29 from litellm.proxy._types import ( 

30 LiteLLM_ObjectPermissionTable, 

31 LiteLLM_TeamTableCachedObj, 

32 ) 

33 

34 

35async def attach_object_permission_to_dict( 

36 data_dict: dict[str, object], 

37 prisma_client: PrismaClient, 

38) -> dict: 

39 """ 

40 Helper method to attach object_permission to a dictionary if object_permission_id is set. 

41 

42 This function: 

43 1. Checks if the dictionary has an object_permission_id 

44 2. If found, queries the database for the corresponding object permission 

45 3. Converts the object permission to a dictionary format 

46 4. Attaches it to the input dictionary under the 'object_permission' key 

47 

48 Args: 

49 data_dict: The dictionary to attach object_permission to 

50 prisma_client: The database client 

51 

52 Returns: 

53 Dict: The input dictionary with object_permission attached if found 

54 

55 Raises: 

56 ValueError: If prisma_client is None 

57 """ 

58 if prisma_client is None: 58 ↛ 59line 58 didn't jump to line 59 because the condition on line 58 was never true

59 raise ValueError("Prisma client not found") 

60 

61 object_permission_id: Final = data_dict.get("object_permission_id") 

62 if object_permission_id: 

63 object_permission = await ObjectPermissionRepository(prisma_client).table.find_unique( 

64 where={"object_permission_id": object_permission_id}, 

65 ) 

66 if object_permission: 66 ↛ 73line 66 didn't jump to line 73 because the condition on line 66 was always true

67 # Convert to dict if needed 

68 try: 

69 object_permission = object_permission.model_dump() 

70 except Exception: 

71 object_permission = object_permission.dict() # pyright: ignore[reportDeprecated] # pydantic v1 fallback 

72 data_dict["object_permission"] = object_permission 

73 return data_dict 

74 

75 

76@dataclass(frozen=True, slots=True) 

77class ObjectPermissionUpsert: 

78 object_permission_id: str 

79 record: dict[str, object] 

80 

81 

82async def prepare_object_permission_upsert( 

83 new_object_permission: Mapping[str, object], 

84 existing_object_permission_id: str | None, 

85 prisma_client: PrismaClient, 

86) -> ObjectPermissionUpsert: 

87 """ 

88 Read-and-merge half of an object permission upsert; performs no writes. 

89 

90 Merges the sent grants over the existing row (looked up by 

91 ``existing_object_permission_id``, or a fresh uuid when the entity has none) and 

92 returns the id plus the full record to upsert. The id is pinned inside the record 

93 because the column has ``@default(uuid())``, so a create without it would mint a 

94 different id than the one the caller links. ``mcp_tool_permissions`` is serialized 

95 to a JSON string to avoid GraphQL parsing issues (e.g. server IDs starting with 

96 "3e64" being interpreted as floats). 

97 

98 Keeping this separate from the write lets callers run the upsert inside the same 

99 transaction as the row that links ``object_permission_id``, so a rolled-back 

100 update cannot leave permission changes live. 

101 """ 

102 object_permission_id: Final = existing_object_permission_id or str(uuid.uuid4()) 

103 existing_object_permission: Final = await ObjectPermissionRepository(prisma_client).table.find_unique( 

104 where={"object_permission_id": object_permission_id}, 

105 ) 

106 existing_fields: Final[dict[str, object]] = ( 

107 existing_object_permission.model_dump(exclude_unset=True, exclude_none=True) 

108 if existing_object_permission is not None 

109 else {} 

110 ) 

111 await reject_ambiguous_mcp_tool_permission_keys( 

112 new_mcp_tool_permissions=new_object_permission.get("mcp_tool_permissions"), 

113 existing_mcp_tool_permissions=existing_fields.get("mcp_tool_permissions"), 

114 prisma_client=prisma_client, 

115 ) 

116 merged: Final[dict[str, object]] = { 

117 **existing_fields, 

118 **new_object_permission, 

119 "object_permission_id": object_permission_id, 

120 } 

121 record: Final[dict[str, object]] = { 

122 **merged, 

123 **( 

124 {"mcp_tool_permissions": safe_dumps(merged["mcp_tool_permissions"])} 

125 if "mcp_tool_permissions" in merged 

126 else {} 

127 ), 

128 } 

129 return ObjectPermissionUpsert(object_permission_id=object_permission_id, record=record) 

130 

131 

132async def handle_update_object_permission_common( 

133 data_json: dict, 

134 existing_object_permission_id: str | None, 

135 prisma_client: PrismaClient | None, 

136) -> str | None: 

137 """ 

138 Common logic for handling object permission updates across organizations, teams, and keys. 

139 

140 This function: 

141 1. Extracts `object_permission` from data_json 

142 2. Looks up existing object permission if it exists 

143 3. Merges new permissions with existing ones 

144 4. Upserts to the LiteLLM_ObjectPermissionTable 

145 5. Returns the object_permission_id 

146 

147 Args: 

148 data_json: The data dictionary containing the object_permission to update 

149 existing_object_permission_id: The current object_permission_id from the entity (can be None) 

150 prisma_client: The database client 

151 

152 Returns: 

153 Optional[str]: The object_permission_id after the update/creation, or None if no object_permission to process 

154 

155 Raises: 

156 ValueError: If prisma_client is None 

157 """ 

158 if prisma_client is None: 158 ↛ 159line 158 didn't jump to line 159 because the condition on line 158 was never true

159 raise ValueError("Prisma client not found") 

160 

161 new_object_permission: dict | str | None = data_json.pop("object_permission", None) 

162 if new_object_permission is None: 162 ↛ 163line 162 didn't jump to line 163 because the condition on line 162 was never true

163 return None 

164 

165 if isinstance(new_object_permission, str): 165 ↛ 166line 165 didn't jump to line 166 because the condition on line 165 was never true

166 new_object_permission = json.loads(new_object_permission) 

167 

168 upsert: Final = await prepare_object_permission_upsert( 

169 new_object_permission=new_object_permission if isinstance(new_object_permission, dict) else {}, 

170 existing_object_permission_id=existing_object_permission_id, 

171 prisma_client=prisma_client, 

172 ) 

173 created_object_permission_row: Final = await ObjectPermissionRepository(prisma_client).table.upsert( 

174 where={"object_permission_id": upsert.object_permission_id}, 

175 data={ 

176 "create": upsert.record, 

177 "update": upsert.record, 

178 }, 

179 ) 

180 

181 verbose_proxy_logger.debug("created_object_permission_row: %s", created_object_permission_row) 

182 

183 return created_object_permission_row.object_permission_id 

184 

185 

186async def invalidate_cached_object_permissions( 

187 object_permission_ids: Iterable[object], 

188 user_api_key_cache: UserApiKeyCache, 

189) -> None: 

190 """Drop permission rows an entitlement change makes stale. 

191 

192 ``get_object_permission`` caches a row under its own id separate from the entity's cache entry, and an 

193 upsert keeps that id, so pass both the outgoing and incoming ids since a change can also mint a new row. 

194 """ 

195 cache_keys: Final = tuple( 

196 object_permission_cache_key(object_permission_id) 

197 for object_permission_id in dict.fromkeys(pid for pid in object_permission_ids if isinstance(pid, str)) 

198 ) 

199 await evict_and_broadcast(cache_keys, user_api_key_cache) 

200 

201 

202async def _set_object_permission( 

203 data_json: dict, 

204 prisma_client: PrismaClient | None, 

205): 

206 """ 

207 Creates the LiteLLM_ObjectPermissionTable record for the key/team. 

208 Handles permissions for vector stores and mcp servers. 

209 """ 

210 if prisma_client is None or "object_permission" not in data_json: 

211 return data_json 

212 

213 permission_data: Final = data_json["object_permission"] 

214 if not isinstance(permission_data, dict): 

215 data_json.pop("object_permission") 

216 return data_json 

217 

218 # Clean data: exclude None values and object_permission_id 

219 clean_data: Final[dict[str, object]] = { 

220 k: v for k, v in permission_data.items() if v is not None and k != "object_permission_id" 

221 } 

222 

223 await reject_ambiguous_mcp_tool_permission_keys( 

224 new_mcp_tool_permissions=clean_data.get("mcp_tool_permissions"), 

225 existing_mcp_tool_permissions=None, 

226 prisma_client=prisma_client, 

227 ) 

228 

229 # Serialize mcp_tool_permissions to JSON string for GraphQL compatibility 

230 if "mcp_tool_permissions" in clean_data: 

231 clean_data["mcp_tool_permissions"] = safe_dumps(clean_data["mcp_tool_permissions"]) 

232 

233 created_permission: Final = await ObjectPermissionRepository(prisma_client).table.create(data=clean_data) 

234 

235 data_json["object_permission_id"] = created_permission.object_permission_id 

236 data_json.pop("object_permission") 

237 return data_json 

238 

239 

240def _dedupe_preserving_order(values: list[str]) -> list[str]: 

241 seen: Final[set[str]] = set() 

242 result: Final[list[str]] = [] 

243 for value in values: 

244 if value in seen: 

245 continue 

246 seen.add(value) 

247 result.append(value) 

248 return result 

249 

250 

251def _mcp_server_identifier_matches(server: object, identifier: str) -> bool: 

252 return identifier in { 

253 getattr(server, "server_id", None), 

254 getattr(server, "alias", None), 

255 getattr(server, "server_name", None), 

256 getattr(server, "name", None), 

257 } 

258 

259 

260async def _get_db_mcp_servers_by_identifiers( 

261 identifiers: AbstractSet[str], 

262 prisma_client: PrismaClient | None, 

263) -> "Sequence[prisma_models.LiteLLM_MCPServerTable]": 

264 if prisma_client is None or not identifiers: 264 ↛ 265line 264 didn't jump to line 265 because the condition on line 264 was never true

265 return [] 

266 

267 identifier_list: Final = list(identifiers) 

268 return await MCPServerRepository(prisma_client).table.find_many( 

269 where={ 

270 "OR": [ 

271 {"server_id": {"in": identifier_list}}, 

272 {"alias": {"in": identifier_list}}, 

273 {"server_name": {"in": identifier_list}}, 

274 ] 

275 } 

276 ) 

277 

278 

279async def _resolve_mcp_server_identifiers_to_ids( 

280 identifiers: AbstractSet[str], 

281 prisma_client: PrismaClient | None, 

282) -> dict[str, set[str]]: 

283 """ 

284 Resolve MCP permission entries written as server_id, alias, or server_name 

285 to canonical server IDs. 

286 

287 DB rows are authoritative when available; the in-memory registry is still 

288 consulted for config-file servers, which are not persisted in the MCP table. 

289 """ 

290 if not identifiers: 

291 return {} 

292 

293 from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( 

294 global_mcp_server_manager, 

295 ) 

296 

297 resolved: Final[dict[str, set[str]]] = {identifier: set() for identifier in identifiers} 

298 

299 for server in await _get_db_mcp_servers_by_identifiers( 

300 identifiers=identifiers, 

301 prisma_client=prisma_client, 

302 ): 

303 server_id = getattr(server, "server_id", None) 

304 if not server_id: 304 ↛ 305line 304 didn't jump to line 305 because the condition on line 304 was never true

305 continue 

306 for identifier in identifiers: 

307 if _mcp_server_identifier_matches(server, identifier): 

308 resolved[identifier].add(server_id) 

309 

310 for registry_key, server in global_mcp_server_manager.get_registry().items(): 

311 server_id = getattr(server, "server_id", None) or registry_key 

312 if not server_id: 312 ↛ 313line 312 didn't jump to line 313 because the condition on line 312 was never true

313 continue 

314 for identifier in identifiers: 

315 if identifier == registry_key or _mcp_server_identifier_matches(server, identifier): 315 ↛ 316line 315 didn't jump to line 316 because the condition on line 315 was never true

316 resolved[identifier].add(server_id) 

317 

318 return resolved 

319 

320 

321_MCP_TOOL_PERMISSIONS_ADAPTER: Final = TypeAdapter(dict[str, list[str] | None]) 

322 

323 

324def _mcp_tool_permission_entries(raw: object) -> Mapping[str, frozenset[str]]: 

325 parsed: Final[Mapping[str, Sequence[str] | None]] = ( 

326 _MCP_TOOL_PERMISSIONS_ADAPTER.validate_json(raw) 

327 if isinstance(raw, str) 

328 else _MCP_TOOL_PERMISSIONS_ADAPTER.validate_python(raw) 

329 if isinstance(raw, Mapping) 

330 else MappingProxyType({}) 

331 ) 

332 return MappingProxyType({identifier: frozenset(tools or ()) for identifier, tools in parsed.items()}) 

333 

334 

335async def reject_ambiguous_mcp_tool_permission_keys( 

336 new_mcp_tool_permissions: object, 

337 existing_mcp_tool_permissions: object, 

338 prisma_client: PrismaClient | None, 

339) -> None: 

340 """ 

341 A name or alias shared by several MCP servers cannot key ``mcp_tool_permissions``: 

342 the read path unions the entry into every match, so no edit can narrow one of 

343 those servers without also changing the other. An exact server_id is never 

344 ambiguous, even when another server uses that string as its alias. Entries the 

345 row already stores with the same tool list are left alone, so unrelated edits 

346 to such an entity still succeed. 

347 

348 Raises HTTPException(400) naming the colliding servers. 

349 """ 

350 requested: Final = _mcp_tool_permission_entries(new_mcp_tool_permissions) 

351 stored: Final = _mcp_tool_permission_entries(existing_mcp_tool_permissions) 

352 resolved: Final = await _resolve_mcp_server_identifiers_to_ids( 

353 identifiers=frozenset(identifier for identifier, tools in requested.items() if stored.get(identifier) != tools), 

354 prisma_client=prisma_client, 

355 ) 

356 collisions: Final = "; ".join( 

357 f"'{identifier}' matches MCP servers {sorted(server_ids)}" 

358 for identifier, server_ids in sorted(resolved.items()) 

359 if identifier not in server_ids and len(server_ids) > 1 

360 ) 

361 if not collisions: 361 ↛ 363line 361 didn't jump to line 363 because the condition on line 361 was always true

362 return 

363 raise HTTPException( 

364 status_code=status.HTTP_400_BAD_REQUEST, 

365 detail={ # mutable-ok: HTTPException.detail has no immutable form; same shape as the sibling errors here 

366 "error": ( 

367 f"Ambiguous mcp_tool_permissions key: {collisions}. " 

368 "Key tool permissions by server_id when servers share a name or alias." 

369 ) 

370 }, 

371 ) 

372 

373 

374def _drop_stale_object_permission_mcp_servers( 

375 object_permission: ObjectPermissionDict, 

376 identifier_to_server_ids: dict[str, set[str]], 

377) -> None: 

378 mcp_servers: Final = object_permission.get("mcp_servers") 

379 if not isinstance(mcp_servers, list): 

380 return 

381 

382 # Persist original identifiers, never resolved ids: shared-DB multi-region 

383 # instances each expand a name/alias to their own local server id at read 

384 # time. Only entries resolving to nothing (deleted servers, typos) drop. 

385 kept_servers: Final = [ 

386 identifier 

387 for identifier in mcp_servers 

388 if identifier == SpecialMCPServerNames.no_mcp_servers.value or identifier_to_server_ids.get(identifier) 

389 ] 

390 object_permission["mcp_servers"] = _dedupe_preserving_order(kept_servers) 

391 

392 

393def _drop_stale_object_permission_mcp_tool_permissions( 

394 object_permission: ObjectPermissionDict, 

395 identifier_to_server_ids: dict[str, set[str]], 

396) -> None: 

397 mcp_tool_permissions: Final = object_permission.get("mcp_tool_permissions") 

398 if not isinstance(mcp_tool_permissions, dict): 

399 return 

400 

401 object_permission["mcp_tool_permissions"] = { 

402 identifier: _dedupe_preserving_order(tools if isinstance(tools, list) else []) 

403 for identifier, tools in mcp_tool_permissions.items() 

404 if identifier_to_server_ids.get(identifier) 

405 } 

406 

407 

408def _drop_stale_object_permission_mcp_identifiers( 

409 object_permission: ObjectPermissionDict | None, 

410 identifier_to_server_ids: dict[str, set[str]], 

411) -> None: 

412 if not object_permission or not isinstance(object_permission, dict): 

413 return 

414 

415 _drop_stale_object_permission_mcp_servers( 

416 object_permission=object_permission, 

417 identifier_to_server_ids=identifier_to_server_ids, 

418 ) 

419 _drop_stale_object_permission_mcp_tool_permissions( 

420 object_permission=object_permission, 

421 identifier_to_server_ids=identifier_to_server_ids, 

422 ) 

423 

424 

425def _flatten_resolved_mcp_server_ids( 

426 identifier_to_server_ids: dict[str, set[str]], 

427) -> set[str]: 

428 return {server_id for server_ids in identifier_to_server_ids.values() for server_id in server_ids} 

429 

430 

431async def _resolve_team_allowed_mcp_servers( 

432 team_object_permission: "LiteLLM_ObjectPermissionTable", 

433 prisma_client: PrismaClient | None = None, 

434) -> set[str]: 

435 """ 

436 Resolve the full set of MCP server IDs a team has access to. 

437 

438 Combines: 

439 - Direct mcp_servers list 

440 - Servers from mcp_access_groups 

441 - Server IDs referenced in mcp_tool_permissions keys 

442 """ 

443 from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( 

444 MCPRequestHandler, 

445 ) 

446 

447 direct_servers: Final[list[str]] = team_object_permission.mcp_servers or [] 

448 if SpecialMCPServerName.all_proxy_servers.value in direct_servers: 448 ↛ 449line 448 didn't jump to line 449 because the condition on line 448 was never true

449 return _get_all_mcp_server_ids() 

450 access_group_servers: Final[list[str]] = await MCPRequestHandler._get_mcp_servers_from_access_groups( 

451 team_object_permission.mcp_access_groups or [] 

452 ) 

453 raw_tool_perms = team_object_permission.mcp_tool_permissions or {} 

454 if isinstance(raw_tool_perms, str): 454 ↛ 455line 454 didn't jump to line 455 because the condition on line 454 was never true

455 raw_tool_perms = json.loads(raw_tool_perms) 

456 tool_perm_servers: Final[list[str]] = list(raw_tool_perms.keys()) 

457 raw_servers: Final = set(direct_servers + access_group_servers + tool_perm_servers) 

458 resolved_servers: Final = await _resolve_mcp_server_identifiers_to_ids( 

459 identifiers=raw_servers, 

460 prisma_client=prisma_client, 

461 ) 

462 unresolved_servers: Final = {server_id for server_id in raw_servers if not resolved_servers.get(server_id)} 

463 return _flatten_resolved_mcp_server_ids(resolved_servers) | unresolved_servers 

464 

465 

466def _get_allow_all_keys_server_ids() -> set[str]: 

467 """Return the set of MCP server IDs marked with allow_all_keys=True.""" 

468 from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( 

469 global_mcp_server_manager, 

470 ) 

471 

472 return set(global_mcp_server_manager.get_allow_all_keys_server_ids()) 

473 

474 

475def _get_all_mcp_server_ids() -> set[str]: 

476 """Return every MCP server id registered on the proxy (config + DB union).""" 

477 from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( 

478 global_mcp_server_manager, 

479 ) 

480 

481 return set(global_mcp_server_manager.get_registry().keys()) 

482 

483 

484async def _existing_object_permission_mcp_servers( 

485 object_permission_id: str | None, 

486 prisma_client: PrismaClient | None, 

487) -> list[str]: 

488 if not object_permission_id or prisma_client is None: 

489 return [] 

490 existing: Final = await ObjectPermissionRepository(prisma_client).table.find_unique( 

491 where={"object_permission_id": object_permission_id}, 

492 ) 

493 if existing is None: 

494 return [] 

495 return existing.mcp_servers or [] 

496 

497 

498async def enforce_all_proxy_mcp_servers_grant_is_admin_only( 

499 requested_mcp_servers: list[str] | None, 

500 existing_object_permission_id: str | None, 

501 is_proxy_admin: bool, 

502 prisma_client: PrismaClient | None, 

503) -> None: 

504 """ 

505 Only a proxy admin may newly grant the all-proxy MCP sentinel. 

506 

507 Scoping a team to every MCP server on the proxy is a proxy-wide authorization 

508 decision, so a caller who is not a proxy admin (e.g. a team admin managing their 

509 own team) cannot add ``all-proxy-mcpservers``. A sentinel a proxy admin already 

510 granted is left untouched, so unrelated edits to such a team still succeed. 

511 

512 Raises HTTPException(403) when a non-admin tries to add the sentinel. 

513 """ 

514 sentinel: Final = SpecialMCPServerName.all_proxy_servers.value 

515 if is_proxy_admin or sentinel not in (requested_mcp_servers or []): 515 ↛ 517line 515 didn't jump to line 517 because the condition on line 515 was always true

516 return 

517 existing_mcp_servers: Final = await _existing_object_permission_mcp_servers( 

518 object_permission_id=existing_object_permission_id, 

519 prisma_client=prisma_client, 

520 ) 

521 if sentinel in existing_mcp_servers: 

522 return 

523 raise HTTPException( 

524 status_code=status.HTTP_403_FORBIDDEN, 

525 detail={ 

526 "error": "Only a proxy admin can grant a team access to all proxy MCP servers ('all-proxy-mcpservers')." 

527 }, 

528 ) 

529 

530 

531async def _get_grandfathered_key_mcp_server_ids( 

532 existing_object_permission: Optional["LiteLLM_ObjectPermissionTable"], 

533 prisma_client: PrismaClient | None, 

534) -> frozenset[str]: 

535 """ 

536 Resolve the canonical MCP server IDs a key's stored object_permission already 

537 grants. Updates that keep or shrink those grants stay valid even when the 

538 team allowlist has since changed; sentinels are excluded so they cannot 

539 grandfather anything. 

540 """ 

541 if existing_object_permission is None or prisma_client is None: 

542 return frozenset() 

543 raw_tool_perms: Final = existing_object_permission.mcp_tool_permissions or {} 

544 tool_perm_keys: Final[frozenset[str]] = frozenset( 

545 json.loads(raw_tool_perms).keys() if isinstance(raw_tool_perms, str) else raw_tool_perms.keys() 

546 ) 

547 identifiers: Final = (frozenset(existing_object_permission.mcp_servers or []) | tool_perm_keys) - { 

548 SpecialMCPServerNames.no_mcp_servers.value, 

549 SpecialMCPServerName.all_proxy_servers.value, 

550 } 

551 return frozenset( 

552 _flatten_resolved_mcp_server_ids( 

553 await _resolve_mcp_server_identifiers_to_ids( 

554 identifiers=set(identifiers), 

555 prisma_client=prisma_client, 

556 ) 

557 ) 

558 ) 

559 

560 

561async def _get_team_allowed_mcp_servers( 

562 team_obj: Optional["LiteLLM_TeamTableCachedObj"], 

563 prisma_client: PrismaClient | None = None, 

564) -> set[str]: 

565 """ 

566 Get the full set of MCP server IDs a team allows. 

567 

568 Combines servers granted via the team's object_permission with servers 

569 granted via the team's unified access groups (access_group_ids). If the 

570 team grants neither, returns empty set (meaning only allow_all_keys 

571 servers are permitted). 

572 """ 

573 if team_obj is None: 573 ↛ 574line 573 didn't jump to line 574 because the condition on line 573 was never true

574 return set() 

575 

576 from litellm.proxy.auth.auth_checks import ( 

577 _get_mcp_server_ids_from_access_groups, # pyright: ignore[reportPrivateUsage] # same resolver runtime MCP auth calls 

578 ) 

579 

580 access_group_servers: Final = await _get_mcp_server_ids_from_access_groups( 

581 access_group_ids=team_obj.access_group_ids or [], 

582 prisma_client=prisma_client, 

583 ) 

584 resolved_access_group_servers: Final = await _resolve_mcp_server_identifiers_to_ids( 

585 identifiers=set(access_group_servers), 

586 prisma_client=prisma_client, 

587 ) 

588 unified_servers: Final = _flatten_resolved_mcp_server_ids(resolved_access_group_servers) | { 

589 server for server in access_group_servers if not resolved_access_group_servers.get(server) 

590 } 

591 

592 team_object_permission: Final = team_obj.object_permission 

593 if team_object_permission is None: 

594 return unified_servers 

595 

596 return unified_servers | await _resolve_team_allowed_mcp_servers( 

597 team_object_permission=team_object_permission, 

598 prisma_client=prisma_client, 

599 ) 

600 

601 

602def _extract_requested_mcp_server_ids( 

603 object_permission: ObjectPermissionDict | None, 

604) -> set[str]: 

605 """ 

606 Extract all MCP server IDs referenced in a key's object_permission dict. 

607 

608 Includes: 

609 - mcp_servers list 

610 - Keys from mcp_tool_permissions 

611 """ 

612 if not object_permission or not isinstance(object_permission, dict): 

613 return set() 

614 

615 server_ids: Final[set[str]] = set() 

616 mcp_servers: Final = object_permission.get("mcp_servers") 

617 if isinstance(mcp_servers, list): 617 ↛ 621line 617 didn't jump to line 621 because the condition on line 617 was always true

618 server_ids.update(mcp_servers) 

619 server_ids.discard(SpecialMCPServerNames.no_mcp_servers.value) 

620 

621 mcp_tool_permissions: Final = object_permission.get("mcp_tool_permissions") 

622 if isinstance(mcp_tool_permissions, dict): 622 ↛ 625line 622 didn't jump to line 625 because the condition on line 622 was always true

623 server_ids.update(mcp_tool_permissions.keys()) 

624 

625 return server_ids 

626 

627 

628def _extract_requested_mcp_access_groups( 

629 object_permission: ObjectPermissionDict | None, 

630) -> set[str]: 

631 """Extract MCP access groups from a key's object_permission dict.""" 

632 if not object_permission or not isinstance(object_permission, dict): 

633 return set() 

634 

635 groups: Final = object_permission.get("mcp_access_groups") 

636 if isinstance(groups, list): 636 ↛ 638line 636 didn't jump to line 638 because the condition on line 636 was always true

637 return set(groups) 

638 return set() 

639 

640 

641def _extract_requested_mcp_toolsets( 

642 object_permission: ObjectPermissionDict | None, 

643) -> set[str]: 

644 """Extract MCP toolset IDs from a key's object_permission dict.""" 

645 if not object_permission or not isinstance(object_permission, dict): 

646 return set() 

647 

648 toolsets: Final = object_permission.get("mcp_toolsets") 

649 if isinstance(toolsets, list): 649 ↛ 651line 649 didn't jump to line 651 because the condition on line 649 was always true

650 return set(toolsets) 

651 return set() 

652 

653 

654async def validate_key_mcp_servers_against_team( 

655 object_permission: ObjectPermissionDict | None, 

656 team_obj: Optional["LiteLLM_TeamTableCachedObj"], 

657 prisma_client: PrismaClient | None = None, 

658 is_proxy_admin: bool = False, 

659 existing_key_object_permission: Optional["LiteLLM_ObjectPermissionTable"] = None, 

660) -> ObjectPermissionDict | None: 

661 """ 

662 Validate that MCP servers requested on a key are within the allowed scope. 

663 

664 When ``existing_key_object_permission`` is provided (key updates), servers 

665 the key already holds are grandfathered: keeping or removing them stays valid 

666 even if the team allowlist has since shrunk, while adding new servers outside 

667 the allowlist is still rejected. 

668 

669 Rules: 

670 - If key is in a team: key's mcp_servers must be a subset of 

671 (team's allowed servers + allow_all_keys servers), where the team's 

672 allowed servers include servers granted via the team's unified 

673 access groups 

674 - If key is NOT in a team and the caller is a proxy admin: any server or 

675 access group may be assigned. A proxy admin can already reach every MCP 

676 server, and runtime access is granted directly from the key's own 

677 object_permission, so the key is scoped to exactly what the admin selected 

678 - If key is NOT in a team and the caller is not a proxy admin: key's 

679 mcp_servers must only contain allow_all_keys servers 

680 - If team has no MCP config (no object_permission and no unified 

681 access groups): key can only use allow_all_keys servers 

682 

683 Raises HTTPException(403) if validation fails. 

684 """ 

685 teamless_admin_assignment: Final = team_obj is None and is_proxy_admin 

686 requested_servers: Final = _extract_requested_mcp_server_ids(object_permission) 

687 requested_access_groups: Final = _extract_requested_mcp_access_groups(object_permission) 

688 

689 requested_toolsets: Final = _extract_requested_mcp_toolsets(object_permission) 

690 

691 # Nothing to validate 

692 if not requested_servers and not requested_access_groups and not requested_toolsets: 692 ↛ 695line 692 didn't jump to line 695 because the condition on line 692 was always true

693 return object_permission 

694 

695 allow_all_keys_servers: Final = _get_allow_all_keys_server_ids() 

696 team_allowed_servers: Final = await _get_team_allowed_mcp_servers( 

697 team_obj=team_obj, 

698 prisma_client=prisma_client, 

699 ) 

700 

701 # Combined allowed set = team servers + allow_all_keys servers 

702 all_allowed_servers: Final = team_allowed_servers | allow_all_keys_servers 

703 

704 # Validate requested server IDs 

705 if requested_servers: 

706 # Normalize aliases/names before authorization. Only entries that do not 

707 # resolve to a server in the DB or config registry are treated as stale. 

708 identifier_to_server_ids: Final = await _resolve_mcp_server_identifiers_to_ids( 

709 identifiers=requested_servers, 

710 prisma_client=prisma_client, 

711 ) 

712 stale_identifiers: Final = { 

713 identifier for identifier in requested_servers if not identifier_to_server_ids.get(identifier) 

714 } 

715 if stale_identifiers: 

716 verbose_proxy_logger.warning( 

717 "validate_key_mcp_servers_against_team: ignoring stale MCP server identifiers (no longer in registry or DB): %s", 

718 sorted(stale_identifiers), 

719 ) 

720 _drop_stale_object_permission_mcp_identifiers( 

721 object_permission=object_permission, 

722 identifier_to_server_ids=identifier_to_server_ids, 

723 ) 

724 active_requested_servers: Final = _flatten_resolved_mcp_server_ids(identifier_to_server_ids) 

725 

726 allowed_servers = all_allowed_servers 

727 if teamless_admin_assignment: 

728 allowed_servers = all_allowed_servers | active_requested_servers 

729 

730 grandfathered_servers: Final = await _get_grandfathered_key_mcp_server_ids( 

731 existing_object_permission=existing_key_object_permission, 

732 prisma_client=prisma_client, 

733 ) 

734 disallowed_servers: Final = active_requested_servers - allowed_servers - grandfathered_servers 

735 if disallowed_servers: 

736 if team_obj is not None: 

737 team_id = team_obj.team_id 

738 detail = ( 

739 f"Key requests MCP servers not allowed by team '{team_id}': " 

740 f"{sorted(disallowed_servers)}. " 

741 f"Team allows: {sorted(team_allowed_servers)}. " 

742 f"Global (allow_all_keys) servers: {sorted(allow_all_keys_servers)}." 

743 ) 

744 else: 

745 detail = ( 

746 f"Key is not in a team. Only globally available (allow_all_keys) MCP servers " 

747 f"can be assigned: {sorted(allow_all_keys_servers)}. " 

748 f"Disallowed servers: {sorted(disallowed_servers)}." 

749 ) 

750 raise HTTPException( 

751 status_code=status.HTTP_403_FORBIDDEN, 

752 detail={"error": detail}, 

753 ) 

754 

755 # Validate requested access groups (must be subset of team's access groups) 

756 if requested_access_groups: 

757 team_access_groups: set[str] = set() 

758 if ( 

759 team_obj is not None 

760 and team_obj.object_permission is not None 

761 and team_obj.object_permission.mcp_access_groups 

762 ): 

763 team_access_groups = set(team_obj.object_permission.mcp_access_groups) 

764 

765 allowed_access_groups = team_access_groups 

766 if teamless_admin_assignment: 

767 allowed_access_groups = team_access_groups | requested_access_groups 

768 

769 disallowed_groups: Final = requested_access_groups - allowed_access_groups 

770 if disallowed_groups: 

771 if team_obj is not None: 

772 team_id = team_obj.team_id 

773 detail = ( 

774 f"Key requests MCP access groups not allowed by team '{team_id}': " 

775 f"{sorted(disallowed_groups)}. " 

776 f"Team allows: {sorted(team_access_groups)}." 

777 ) 

778 else: 

779 detail = ( 

780 f"Key is not in a team. MCP access groups cannot be assigned to " 

781 f"keys outside of a team. Disallowed groups: {sorted(disallowed_groups)}." 

782 ) 

783 raise HTTPException( 

784 status_code=status.HTTP_403_FORBIDDEN, 

785 detail={"error": detail}, 

786 ) 

787 

788 _validate_requested_toolsets( 

789 requested_toolsets=requested_toolsets, 

790 team_obj=team_obj, 

791 is_proxy_admin=is_proxy_admin, 

792 ) 

793 

794 return object_permission 

795 

796 

797def _validate_requested_toolsets( 

798 requested_toolsets: set[str], 

799 team_obj: Optional["LiteLLM_TeamTableCachedObj"], 

800 is_proxy_admin: bool, 

801) -> None: 

802 """ 

803 Validate mcp_toolsets requested on a key. 

804 

805 Non-admin callers cannot assign toolsets to a personal (no team) key. Team 

806 keys must request a subset of the team's own toolset allowlist. 

807 """ 

808 if not requested_toolsets: 

809 return 

810 if team_obj is None: 

811 if is_proxy_admin: 

812 return 

813 raise HTTPException( 

814 status_code=status.HTTP_403_FORBIDDEN, 

815 detail={ 

816 "error": ( 

817 "Key is not in a team. MCP toolsets cannot be assigned to " 

818 "personal keys by non-admin callers. Disallowed toolsets: " 

819 f"{sorted(requested_toolsets)}." 

820 ) 

821 }, 

822 ) 

823 team_op: Final = team_obj.object_permission 

824 team_mcp_toolsets: Final = team_op.mcp_toolsets if team_op is not None else None 

825 if not team_mcp_toolsets: 

826 return 

827 disallowed_toolsets: Final = requested_toolsets - set(team_mcp_toolsets) 

828 if not disallowed_toolsets: 

829 return 

830 raise HTTPException( 

831 status_code=status.HTTP_403_FORBIDDEN, 

832 detail={ 

833 "error": ( 

834 f"Key requests MCP toolsets not allowed by team '{team_obj.team_id}': " 

835 f"{sorted(disallowed_toolsets)}. " 

836 f"Team allows: {sorted(team_mcp_toolsets)}." 

837 ) 

838 }, 

839 ) 

840 

841 

842def _extract_requested_vector_stores( 

843 object_permission: ObjectPermissionDict | None, 

844) -> set[str]: 

845 """Return vector_store IDs from a key's object_permission dict.""" 

846 if not object_permission or not isinstance(object_permission, dict): 

847 return set() 

848 raw: Final = object_permission.get("vector_stores") 

849 if isinstance(raw, list): 849 ↛ 851line 849 didn't jump to line 851 because the condition on line 849 was always true

850 return {str(x) for x in raw if x} 

851 return set() 

852 

853 

854async def validate_key_vector_stores_against_team( 

855 object_permission: ObjectPermissionDict | None, 

856 team_obj: Optional["LiteLLM_TeamTableCachedObj"], 

857 is_proxy_admin: bool = False, 

858) -> None: 

859 """ 

860 Reject vector_stores requested on a personal (no team) key by a non-admin 

861 caller. Vector store access is granted at use-time from the key's 

862 object_permission.vector_stores list, so the assignment is the authorization 

863 boundary. Team keys and proxy admins are unaffected. 

864 """ 

865 requested: Final = _extract_requested_vector_stores(object_permission) 

866 if not requested: 866 ↛ 868line 866 didn't jump to line 868 because the condition on line 866 was always true

867 return 

868 if team_obj is not None or is_proxy_admin: