Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py: 34%

1""" 

2Agent Permission Handler for LiteLLM Proxy. 

3 

4Handles agent permission checking for keys and teams using object_permission_id. 

5Follows the same pattern as MCP permission handling. 

6""" 

7 

8import asyncio 

9from collections.abc import Awaitable, Callable, Sequence 

10from dataclasses import dataclass 

11from typing import Final, TypeAlias 

12 

13from litellm._logging import verbose_logger 

14from litellm.proxy._experimental.mcp_server.ui_session_utils import build_effective_auth_contexts 

15from litellm.proxy._types import ( 

16 UI_TEAM_ID, 

17 LiteLLM_ObjectPermissionTable, 

18 LiteLLM_TeamTable, 

19 LitellmUserRoles, 

20 UserAPIKeyAuth, 

21) 

22from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( 

23 CeilingResolver, 

24 resolve_agent_access_group_ceiling, 

25) 

26from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth 

27from litellm.repositories.table_repositories import AgentsRepository 

28from litellm.types.agents import AgentResponse 

29 

30 

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

32class UnrestrictedAgentAccess: 

33 """No agent grant exists on the key or its team, so every agent is reachable.""" 

34 

35 

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

37class RestrictedAgentAccess: 

38 """Only ``agent_ids`` are reachable. An empty set denies every agent.""" 

39 

40 agent_ids: frozenset[str] 

41 

42 

43AgentAccess: TypeAlias = UnrestrictedAgentAccess | RestrictedAgentAccess 

44 

45 

46def _to_stable_ids(agent_ids: frozenset[str]) -> frozenset[str]: 

47 from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry 

48 

49 return frozenset(global_agent_registry.stable_agent_id(agent_id) for agent_id in agent_ids) 

50 

51 

52def _restricted_ids(access: AgentAccess) -> frozenset[str] | None: 

53 if isinstance(access, UnrestrictedAgentAccess): 53 ↛ 55line 53 didn't jump to line 55 because the condition on line 53 was always true

54 return None 

55 return _to_stable_ids(access.agent_ids) 

56 

57 

58def _intersect_agent_access(key_access: AgentAccess, team_access: AgentAccess) -> AgentAccess: 

59 key_ids: Final = _restricted_ids(key_access) 

60 team_ids: Final = _restricted_ids(team_access) 

61 if key_ids is None: 61 ↛ 63line 61 didn't jump to line 63 because the condition on line 61 was always true

62 return UnrestrictedAgentAccess() if team_ids is None else RestrictedAgentAccess(team_ids) 

63 if team_ids is None: 

64 return RestrictedAgentAccess(key_ids) 

65 return RestrictedAgentAccess(key_ids & team_ids) 

66 

67 

68class AgentRequestHandler: 

69 """ 

70 Class to handle agent permission checking, including: 

71 1. Key-level agent permissions 

72 2. Team-level agent permissions 

73 3. Agent access group resolution 

74 

75 Follows the same inheritance logic as MCP: 

76 - If team has restrictions and key has restrictions: use intersection 

77 - If team has restrictions and key has none: inherit from team 

78 - If team has no restrictions: use key restrictions 

79 - If no restrictions: allow all agents 

80 """ 

81 

82 @staticmethod 

83 async def resolve_agent_access( 

84 user_api_key_auth: UserAPIKeyAuth | None = None, 

85 resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, 

86 ) -> AgentAccess: 

87 """Agents the key may reach: key and team grants, intersected with the agent's access group ceiling 

88 and, for an agent key acting on behalf of an invoking user, with that user's team grants.""" 

89 key_team_access: Final = await AgentRequestHandler._resolve_key_team_agent_access(user_api_key_auth) 

90 caller_access: Final = await AgentRequestHandler._agent_caller_access(user_api_key_auth) 

91 own_access: Final = _intersect_agent_access(key_team_access, caller_access) 

92 agent_ceiling: Final = await AgentRequestHandler._agent_access_group_ceiling(user_api_key_auth, resolve_ceiling) 

93 if agent_ceiling is None: 93 ↛ 95line 93 didn't jump to line 95 because the condition on line 93 was always true

94 return own_access 

95 if isinstance(own_access, UnrestrictedAgentAccess): 

96 return RestrictedAgentAccess(agent_ceiling) 

97 return RestrictedAgentAccess(own_access.agent_ids & agent_ceiling) 

98 

99 @staticmethod 

100 async def _agent_caller_access(user_api_key_auth: UserAPIKeyAuth | None) -> AgentAccess: 

101 caller_auth: Final = agent_caller_auth(user_api_key_auth) if user_api_key_auth else None 

102 if caller_auth is None: 102 ↛ 104line 102 didn't jump to line 104 because the condition on line 102 was always true

103 return UnrestrictedAgentAccess() 

104 return await AgentRequestHandler._get_allowed_agents_for_team(caller_auth) 

105 

106 @staticmethod 

107 async def _resolve_key_team_agent_access( 

108 user_api_key_auth: UserAPIKeyAuth | None, 

109 ) -> AgentAccess: 

110 try: 

111 key_access: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) 

112 team_access: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) 

113 except Exception as e: 

114 verbose_logger.warning("Failed to get allowed agents: %s", e) 

115 return UnrestrictedAgentAccess() 

116 return _intersect_agent_access(key_access, team_access) 

117 

118 @staticmethod 

119 async def _agent_access_group_ceiling( 

120 user_api_key_auth: UserAPIKeyAuth | None, 

121 resolve_ceiling: CeilingResolver, 

122 ) -> frozenset[str] | None: 

123 if user_api_key_auth is None or not user_api_key_auth.agent_id: 123 ↛ 125line 123 didn't jump to line 125 because the condition on line 123 was always true

124 return None 

125 ceiling: Final = await resolve_ceiling(user_api_key_auth.agent_id) 

126 if ceiling is None: 

127 return None 

128 return _to_stable_ids(ceiling.agent_ids) 

129 

130 @staticmethod 

131 async def is_agent_allowed( 

132 agent_id: str, 

133 user_api_key_auth: UserAPIKeyAuth | None = None, 

134 resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, 

135 ) -> bool: 

136 """ 

137 Check if a specific agent is allowed for the given user/key. 

138 

139 Args: 

140 agent_id: The agent ID to check 

141 user_api_key_auth: User authentication info 

142 

143 Returns: 

144 bool: True if agent is allowed, False otherwise 

145 """ 

146 from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry 

147 

148 match await AgentRequestHandler.resolve_agent_access(user_api_key_auth, resolve_ceiling): 

149 case UnrestrictedAgentAccess(): 

150 return True 

151 case RestrictedAgentAccess(allowed_agent_ids): 

152 stable_id: Final = global_agent_registry.stable_agent_id(agent_id) 

153 return not global_agent_registry.ids_for_agent(stable_id).isdisjoint(allowed_agent_ids) 

154 

155 @staticmethod 

156 def _get_key_object_permission( 

157 user_api_key_auth: UserAPIKeyAuth | None = None, 

158 ) -> LiteLLM_ObjectPermissionTable | None: 

159 """ 

160 Get key object_permission - already loaded by get_key_object() in main auth flow. 

161 

162 Note: object_permission is automatically populated when the key is fetched via 

163 get_key_object() in litellm/proxy/auth/auth_checks.py 

164 """ 

165 if not user_api_key_auth: 165 ↛ 166line 165 didn't jump to line 166 because the condition on line 165 was never true

166 return None 

167 

168 return user_api_key_auth.object_permission 

169 

170 @staticmethod 

171 async def _get_team_object_permission( 

172 user_api_key_auth: UserAPIKeyAuth | None = None, 

173 ) -> LiteLLM_ObjectPermissionTable | None: 

174 """ 

175 Get team object_permission - automatically loaded by get_team_object() in main auth flow. 

176 

177 Note: object_permission is automatically populated when the team is fetched via 

178 get_team_object() in litellm/proxy/auth/auth_checks.py 

179 """ 

180 from litellm.proxy.auth.auth_checks import get_team_object 

181 from litellm.proxy.proxy_server import ( 

182 prisma_client, 

183 proxy_logging_obj, 

184 user_api_key_cache, 

185 ) 

186 

187 if not user_api_key_auth or not user_api_key_auth.team_id or not prisma_client: 

188 return None 

189 

190 # Get the team object (which has object_permission already loaded) 

191 team_obj: Final[LiteLLM_TeamTable | None] = await get_team_object( 

192 team_id=user_api_key_auth.team_id, 

193 prisma_client=prisma_client, 

194 user_api_key_cache=user_api_key_cache, 

195 parent_otel_span=user_api_key_auth.parent_otel_span, 

196 proxy_logging_obj=proxy_logging_obj, 

197 ) 

198 

199 if not team_obj: 

200 return None 

201 

202 return team_obj.object_permission 

203 

204 @staticmethod 

205 async def _get_allowed_agents_for_key( 

206 user_api_key_auth: UserAPIKeyAuth | None = None, 

207 ) -> AgentAccess: 

208 """ 

209 Get allowed agents for a key. 

210 

211 1. First checks native key-level agent permissions (object_permission) 

212 2. Also includes agents from key's access_group_ids (unified access groups) 

213 

214 A key that declares agents or access groups is restricted even when those 

215 declarations resolve to nothing, so an emptied or deleted access group denies 

216 rather than opening the key up. Lookup failures still propagate to the caller, 

217 which keeps them fail-open. 

218 

219 Note: object_permission is already loaded by get_key_object() in main auth flow. 

220 """ 

221 if user_api_key_auth is None: 221 ↛ 222line 221 didn't jump to line 222 because the condition on line 221 was never true

222 return UnrestrictedAgentAccess() 

223 

224 try: 

225 # 1. Get agents from object_permission (native permissions) 

226 key_object_permission: Final = AgentRequestHandler._get_key_object_permission(user_api_key_auth) 

227 direct_agents: Final = tuple( 

228 key_object_permission.agents or () if key_object_permission is not None else () 

229 ) 

230 declared_access_groups: Final = tuple( 

231 key_object_permission.agent_access_groups or () if key_object_permission is not None else () 

232 ) 

233 # 2. Fallback: get agent IDs from key's access_group_ids (unified access groups) 

234 key_access_group_ids: Final = tuple(user_api_key_auth.access_group_ids or ()) 

235 

236 if not direct_agents and not declared_access_groups and not key_access_group_ids: 236 ↛ 239line 236 didn't jump to line 239 because the condition on line 236 was always true

237 return UnrestrictedAgentAccess() 

238 

239 access_group_agents: Final = ( 

240 tuple(await AgentRequestHandler._get_agents_from_access_groups(list(declared_access_groups))) 

241 if declared_access_groups 

242 else () 

243 ) 

244 unified_agents: Final = ( 

245 tuple(await AgentRequestHandler._get_unified_access_group_agents(list(key_access_group_ids))) 

246 if key_access_group_ids 

247 else () 

248 ) 

249 

250 return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents)) 

251 except Exception as e: 

252 verbose_logger.warning("Failed to get allowed agents for key: %s", e) 

253 return UnrestrictedAgentAccess() 

254 

255 @staticmethod 

256 async def _get_allowed_agents_for_team( 

257 user_api_key_auth: UserAPIKeyAuth | None = None, 

258 ) -> AgentAccess: 

259 """ 

260 Get allowed agents for a team. 

261 

262 1. First checks native team-level agent permissions (object_permission) 

263 2. Also includes agents from team's access_group_ids (unified access groups) 

264 

265 Fetches the team object once and reuses it for both permission sources. 

266 Declared-but-empty grants stay restricted; see `_get_allowed_agents_for_key`. 

267 """ 

268 if user_api_key_auth is None: 268 ↛ 269line 268 didn't jump to line 269 because the condition on line 268 was never true

269 return UnrestrictedAgentAccess() 

270 

271 if user_api_key_auth.team_id is None: 271 ↛ 274line 271 didn't jump to line 274 because the condition on line 271 was always true

272 return UnrestrictedAgentAccess() 

273 

274 try: 

275 from litellm.proxy.auth.auth_checks import get_team_object 

276 from litellm.proxy.proxy_server import ( 

277 prisma_client, 

278 proxy_logging_obj, 

279 user_api_key_cache, 

280 ) 

281 

282 if not prisma_client: 

283 return UnrestrictedAgentAccess() 

284 

285 # Fetch the team object once for both permission sources 

286 team_obj: Final = await get_team_object( 

287 team_id=user_api_key_auth.team_id, 

288 prisma_client=prisma_client, 

289 user_api_key_cache=user_api_key_cache, 

290 parent_otel_span=user_api_key_auth.parent_otel_span, 

291 proxy_logging_obj=proxy_logging_obj, 

292 ) 

293 

294 if team_obj is None: 

295 return UnrestrictedAgentAccess() 

296 

297 # 1. Get agents from object_permission (native permissions) 

298 object_permissions: Final = team_obj.object_permission 

299 direct_agents: Final = tuple(object_permissions.agents or () if object_permissions is not None else ()) 

300 declared_access_groups: Final = tuple( 

301 object_permissions.agent_access_groups or () if object_permissions is not None else () 

302 ) 

303 # 2. Also include agents from team's access_group_ids (unified access groups) 

304 team_access_group_ids: Final = tuple(team_obj.access_group_ids or ()) 

305 

306 if not direct_agents and not declared_access_groups and not team_access_group_ids: 

307 return UnrestrictedAgentAccess() 

308 

309 access_group_agents: Final = ( 

310 tuple(await AgentRequestHandler._get_agents_from_access_groups(list(declared_access_groups))) 

311 if declared_access_groups 

312 else () 

313 ) 

314 unified_agents: Final = ( 

315 tuple(await AgentRequestHandler._get_unified_access_group_agents(list(team_access_group_ids))) 

316 if team_access_group_ids 

317 else () 

318 ) 

319 

320 return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents)) 

321 except Exception as e: 

322 # litellm-dashboard is the default UI team and will never have agents; 

323 # skip noisy warnings for it. 

324 if user_api_key_auth.team_id != UI_TEAM_ID: 

325 verbose_logger.warning("Failed to get allowed agents for team: %s", e) 

326 return UnrestrictedAgentAccess() 

327 

328 @staticmethod 

329 def _get_config_agent_ids_for_access_groups(config_agents: list, access_groups: list[str]) -> set[str]: 

330 """ 

331 Helper to get agent_ids from config-loaded agents that match any of the given access groups. 

332 """ 

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

334 for agent in config_agents: 

335 agent_access_groups = getattr(agent, "agent_access_groups", None) 

336 if agent_access_groups: 

337 if any(group in agent_access_groups for group in access_groups): 

338 server_ids.add(agent.agent_id) 

339 return server_ids 

340 

341 @staticmethod 

342 async def _get_db_agent_ids_for_access_groups(prisma_client, access_groups: list[str]) -> set[str]: 

343 """ 

344 Helper to get agent_ids from DB agents that match any of the given access groups. 

345 

346 Query failures propagate so the caller can tell "this group is empty" (deny) 

347 apart from "the lookup failed" (fail-open). 

348 """ 

349 if not access_groups or prisma_client is None: 

350 return set() 

351 

352 agents: Final = await AgentsRepository(prisma_client).table.find_many( 

353 where={"agent_access_groups": {"hasSome": access_groups}} 

354 ) 

355 return {agent.agent_id for agent in agents} 

356 

357 @staticmethod 

358 async def _get_unified_access_group_agents(access_group_ids: list[str]) -> list[str]: 

359 """ 

360 Resolve unified access group ids to agent IDs. 

361 """ 

362 from litellm.proxy.auth.auth_checks import _get_agent_ids_from_access_groups 

363 

364 return await _get_agent_ids_from_access_groups(access_group_ids=access_group_ids) 

365 

366 @staticmethod 

367 async def _get_agents_from_access_groups( 

368 access_groups: list[str], 

369 ) -> list[str]: 

370 """ 

371 Resolve agent access groups to agent IDs by querying BOTH the agent table (DB) AND config-loaded agents. 

372 """ 

373 from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry 

374 from litellm.proxy.proxy_server import prisma_client 

375 

376 # Use the helper for config-loaded agents 

377 config_agent_ids: Final = AgentRequestHandler._get_config_agent_ids_for_access_groups( 

378 global_agent_registry.agent_list, access_groups 

379 ) 

380 

381 # Use the helper for DB agents 

382 db_agent_ids: Final = await AgentRequestHandler._get_db_agent_ids_for_access_groups( 

383 prisma_client, access_groups 

384 ) 

385 

386 return list(config_agent_ids | db_agent_ids) 

387 

388 @staticmethod 

389 async def get_agent_access_groups( 

390 user_api_key_auth: UserAPIKeyAuth | None = None, 

391 ) -> list[str]: 

392 """ 

393 Get list of agent access groups for the given user/key based on permissions. 

394 """ 

395 access_groups: list[str] = [] 

396 access_groups_for_key: Final = await AgentRequestHandler._get_agent_access_groups_for_key(user_api_key_auth) 

397 access_groups_for_team: Final = await AgentRequestHandler._get_agent_access_groups_for_team(user_api_key_auth) 

398 

399 # If team has access groups, then key must have a subset of the team's access groups 

400 if len(access_groups_for_team) > 0: 

401 for access_group in access_groups_for_key: 

402 if access_group in access_groups_for_team: 

403 access_groups.append(access_group) 

404 else: 

405 access_groups = access_groups_for_key 

406 

407 return list(set(access_groups)) 

408 

409 @staticmethod 

410 async def _get_agent_access_groups_for_key( 

411 user_api_key_auth: UserAPIKeyAuth | None = None, 

412 ) -> list[str]: 

413 """Get agent access groups for the key.""" 

414 from litellm.proxy.auth.auth_checks import get_object_permission 

415 from litellm.proxy.proxy_server import ( 

416 prisma_client, 

417 proxy_logging_obj, 

418 user_api_key_cache, 

419 ) 

420 

421 if user_api_key_auth is None: 

422 return [] 

423 

424 if user_api_key_auth.object_permission_id is None: 

425 return [] 

426 

427 if prisma_client is None: 

428 verbose_logger.debug("prisma_client is None") 

429 return [] 

430 

431 try: 

432 key_object_permission: Final = await get_object_permission( 

433 object_permission_id=user_api_key_auth.object_permission_id, 

434 prisma_client=prisma_client, 

435 user_api_key_cache=user_api_key_cache, 

436 parent_otel_span=user_api_key_auth.parent_otel_span, 

437 proxy_logging_obj=proxy_logging_obj, 

438 ) 

439 if key_object_permission is None: 

440 return [] 

441 

442 return key_object_permission.agent_access_groups or [] 

443 except Exception as e: 

444 verbose_logger.warning("Failed to get agent access groups for key: %s", e) 

445 return [] 

446 

447 @staticmethod 

448 async def _get_agent_access_groups_for_team( 

449 user_api_key_auth: UserAPIKeyAuth | None = None, 

450 ) -> list[str]: 

451 """Get agent access groups for the team.""" 

452 from litellm.proxy.auth.auth_checks import get_team_object 

453 from litellm.proxy.proxy_server import ( 

454 prisma_client, 

455 proxy_logging_obj, 

456 user_api_key_cache, 

457 ) 

458 

459 if user_api_key_auth is None: 

460 return [] 

461 

462 if user_api_key_auth.team_id is None: 

463 return [] 

464 

465 if prisma_client is None: 

466 verbose_logger.debug("prisma_client is None") 

467 return [] 

468 

469 try: 

470 team_obj: Final[LiteLLM_TeamTable | None] = await get_team_object( 

471 team_id=user_api_key_auth.team_id, 

472 prisma_client=prisma_client, 

473 user_api_key_cache=user_api_key_cache, 

474 parent_otel_span=user_api_key_auth.parent_otel_span, 

475 proxy_logging_obj=proxy_logging_obj, 

476 ) 

477 if team_obj is None: 

478 verbose_logger.debug("team_obj is None") 

479 return [] 

480 

481 object_permissions: Final = team_obj.object_permission 

482 if object_permissions is None: 

483 return [] 

484 

485 return object_permissions.agent_access_groups or [] 

486 except Exception as e: 

487 verbose_logger.warning("Failed to get agent access groups for team: %s", e) 

488 return [] 

489 

490 

491def _granted_ids(access: AgentAccess) -> frozenset[str]: 

492 match access: 

493 case UnrestrictedAgentAccess(): 

494 return frozenset() 

495 case RestrictedAgentAccess(agent_ids): 

496 return agent_ids 

497 

498 

499ResolveAgentAccess: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[AgentAccess]] 

500EffectiveAuthContexts: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[Sequence[UserAPIKeyAuth]]] 

501 

502 

503async def _granted_agent_ids( 

504 user_api_key_auth: UserAPIKeyAuth, 

505 resolve_access: ResolveAgentAccess, 

506 effective_contexts: EffectiveAuthContexts, 

507) -> frozenset[str]: 

508 """Union of the explicit grants reachable from the key, its team, or (for a dashboard session) 

509 the user's real teams and user row. No grant anywhere yields the empty set, unlike the 

510 open-by-default ``resolve_agent_access`` that guards direct access.""" 

511 accesses: Final = await asyncio.gather( 

512 *(resolve_access(auth_context) for auth_context in await effective_contexts(user_api_key_auth)) 

513 ) 

514 return frozenset().union(*(_granted_ids(access) for access in accesses)) 

515 

516 

517async def accessible_agents( 

518 user_api_key_auth: UserAPIKeyAuth, 

519 all_agents: tuple[AgentResponse, ...] | None = None, 

520 resolve_access: ResolveAgentAccess | None = None, 

521 effective_contexts: EffectiveAuthContexts = build_effective_auth_contexts, 

522) -> tuple[AgentResponse, ...]: 

523 """Every registry agent for proxy admins, else only the agents the caller was granted.""" 

524 from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry 

525 

526 agents: Final = global_agent_registry.get_agent_list() if all_agents is None else all_agents 

527 if user_api_key_auth.user_role in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN.value): 527 ↛ 529line 527 didn't jump to line 529 because the condition on line 527 was always true

528 return agents 

529 allowed_agent_ids: Final = await _granted_agent_ids( 

530 user_api_key_auth, 

531 AgentRequestHandler.resolve_agent_access if resolve_access is None else resolve_access, 

532 effective_contexts, 

533 ) 

534 return tuple(agent for agent in agents if agent.agent_id in allowed_agent_ids)