Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/agent_endpoints/agent_registry.py: 78%

1import asyncio 

2import hashlib 

3import json 

4from collections.abc import Callable, Iterator, Mapping, Sequence 

5from datetime import datetime, timezone 

6from types import MappingProxyType 

7from typing import TYPE_CHECKING, Final, NamedTuple, Protocol, TypedDict 

8 

9from pydantic import TypeAdapter, ValidationError 

10from typing_extensions import ReadOnly 

11 

12import litellm 

13from litellm.constants import REDACTED_BY_LITELM_STRING 

14from litellm.litellm_core_utils.safe_json_dumps import safe_dumps 

15from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker 

16from litellm.proxy.agent_endpoints.kill_switch import restore_kill_switch 

17from litellm.proxy.management_helpers.object_permission_utils import ( 

18 handle_update_object_permission_common, 

19) 

20from litellm.proxy.utils import PrismaClient 

21from litellm.repositories.prisma_protocols import TableActions 

22from litellm.repositories.table_repositories import AgentsRepository, ObjectPermissionRepository 

23from litellm.types.agents import AgentConfig, AgentKillSwitchConfig, AgentResponse, PatchAgentRequest 

24 

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

26 from prisma import models as prisma_models 

27 

28 

29class AgentObjectPermissionRecord(Protocol): 

30 def model_dump(self) -> dict[str, object]: ... 30 ↛ exitline 30 didn't return from function 'model_dump' because

31 

32 def dict(self) -> dict[str, object]: ... 32 ↛ exitline 32 didn't return from function 'dict' because

33 

34 

35class AgentIdWhere(TypedDict): 

36 agent_id: ReadOnly[str] 

37 

38 

39class AgentRecordDump(TypedDict): 

40 agent_id: str 

41 agent_name: str 

42 litellm_params: dict[str, object] | None 

43 agent_card_params: dict[str, object] 

44 static_headers: dict[str, str] | None 

45 extra_headers: list[str] | None 

46 kill_switch: ReadOnly[AgentKillSwitchConfig | None] 

47 access_group_ids: ReadOnly[Sequence[str] | None] 

48 object_permission: dict[str, object] | None 

49 spend: float 

50 tpm_limit: int | None 

51 rpm_limit: int | None 

52 session_tpm_limit: int | None 

53 session_rpm_limit: int | None 

54 created_at: datetime 

55 updated_at: datetime 

56 created_by: str | None 

57 updated_by: str | None 

58 

59 

60class AgentRecord(Protocol): 

61 @property 

62 def agent_id(self) -> str: ... 62 ↛ exitline 62 didn't return from function 'agent_id' because

63 

64 @property 

65 def agent_name(self) -> str: ... 65 ↛ exitline 65 didn't return from function 'agent_name' because

66 

67 @property 

68 def litellm_params(self) -> Mapping[str, object] | None: ... 68 ↛ exitline 68 didn't return from function 'litellm_params' because

69 

70 @property 

71 def object_permission_id(self) -> str | None: ... 71 ↛ exitline 71 didn't return from function 'object_permission_id' because

72 

73 @property 

74 def object_permission(self) -> AgentObjectPermissionRecord | None: ... 74 ↛ exitline 74 didn't return from function 'object_permission' because

75 

76 @property 

77 def access_group_ids(self) -> Sequence[str] | None: ... 77 ↛ exitline 77 didn't return from function 'access_group_ids' because

78 

79 @property 

80 def kill_switch(self) -> Mapping[str, object] | None: ... 80 ↛ exitline 80 didn't return from function 'kill_switch' because

81 

82 @property 

83 def spend(self) -> float: ... 83 ↛ exitline 83 didn't return from function 'spend' because

84 

85 def model_dump(self) -> AgentRecordDump: ... 85 ↛ exitline 85 didn't return from function 'model_dump' because

86 

87 def __iter__(self) -> Iterator[tuple[str, object]]: ... 87 ↛ exitline 87 didn't return from function '__iter__' because

88 

89 

90class AgentTableClient(Protocol): 

91 async def create( 91 ↛ exitline 91 didn't return from function 'create' because

92 self, 

93 data: Mapping[str, object], 

94 include: Mapping[str, object] | None = None, 

95 ) -> AgentRecord: ... 

96 

97 async def find_unique( 97 ↛ exitline 97 didn't return from function 'find_unique' because

98 self, 

99 where: Mapping[str, object], 

100 include: Mapping[str, object] | None = None, 

101 ) -> AgentRecord | None: ... 

102 

103 async def find_many( 103 ↛ exitline 103 didn't return from function 'find_many' because

104 self, 

105 where: Mapping[str, object] | None = None, 

106 order: Mapping[str, str] | None = None, 

107 include: Mapping[str, object] | None = None, 

108 ) -> Sequence[AgentRecord]: ... 

109 

110 async def update( 110 ↛ exitline 110 didn't return from function 'update' because

111 self, 

112 data: Mapping[str, object], 

113 where: Mapping[str, object], 

114 include: Mapping[str, object] | None = None, 

115 ) -> AgentRecord | None: ... 

116 

117 async def delete( 117 ↛ exitline 117 didn't return from function 'delete' because

118 self, 

119 where: Mapping[str, object], 

120 include: Mapping[str, object] | None = None, 

121 ) -> AgentRecord | None: ... 

122 

123 

124def agents_table(prisma_client: PrismaClient) -> AgentTableClient: 

125 table: Final[AgentTableClient] = AgentsRepository(prisma_client).table # pyright: ignore[reportAssignmentType] # prisma rows type model_dump() as dict[str, Any] 

126 return table 

127 

128 

129def object_permission_table( 

130 prisma_client: PrismaClient, 

131) -> "TableActions[prisma_models.LiteLLM_ObjectPermissionTable]": 

132 table: Final[TableActions[prisma_models.LiteLLM_ObjectPermissionTable]] = ObjectPermissionRepository( 

133 prisma_client 

134 ).table 

135 return table 

136 

137 

138def _dump_agent_params(raw: Mapping[str, object]) -> dict[str, object]: 

139 model_dump: Final[Callable[[], dict[str, object]] | None] = getattr(raw, "model_dump", None) 

140 if model_dump is not None: 140 ↛ 141line 140 didn't jump to line 141 because the condition on line 140 was never true

141 return model_dump() 

142 return dict(raw) if raw else {} 

143 

144 

145_AGENT_PARAMS_MASKER: Final = SensitiveDataMasker() 

146_REDACT_AGENT_PARAMS_MAX_DEPTH: Final = 10 

147_AGENT_PARAMS_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) 

148_AGENT_PARAMS_SEQUENCE_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...]) 

149_EMPTY_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) 

150 

151 

152def redact_sensitive_agent_litellm_params(litellm_params: object, _depth: int = 0) -> object: 

153 """ 

154 Replace credential-bearing values in an agent's litellm_params with 

155 ``REDACTED_BY_LITELM_STRING`` while preserving non-secret keys (``model``, 

156 ``is_public``, rate-limit config). Used so list/get/create/update 

157 responses never echo a stored provider credential back to the caller. 

158 

159 Handles a plain dict, a JSON-serialized string (some callers hold the 

160 in-memory registry's params that way), and ``None`` at the top level; 

161 anything else is passed through. Recursion depth is bounded to match the 

162 convention documented in ``tests/code_coverage_tests/recursive_detector.py``. 

163 """ 

164 if litellm_params is None: 164 ↛ 165line 164 didn't jump to line 165 because the condition on line 164 was never true

165 return None 

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

167 if _depth >= _REDACT_AGENT_PARAMS_MAX_DEPTH: 

168 return REDACTED_BY_LITELM_STRING 

169 try: 

170 parsed_params: Final = _AGENT_PARAMS_ADAPTER.validate_json(litellm_params) 

171 except ValidationError: 

172 return REDACTED_BY_LITELM_STRING 

173 return json.dumps(_redact_agent_params_tree(parsed_params, _depth + 1)) 

174 return _redact_agent_params_tree(litellm_params, _depth) 

175 

176 

177def _redact_agent_params_tree(value: object, _depth: int) -> object: 

178 """Structural recursion over an already-parsed litellm_params value: a 

179 dict redacts sensitive keys and recurses into the rest, a list redacts 

180 each element (so a secret nested inside a list of provider configs is 

181 still caught), and anything else -- including a plain string leaf, which 

182 must never be re-interpreted as a JSON blob -- passes through unchanged. 

183 """ 

184 if _depth >= _REDACT_AGENT_PARAMS_MAX_DEPTH: 184 ↛ 185line 184 didn't jump to line 185 because the condition on line 184 was never true

185 return REDACTED_BY_LITELM_STRING 

186 if isinstance(value, list): 

187 typed_items: Final = _AGENT_PARAMS_SEQUENCE_ADAPTER.validate_python(value) 

188 return tuple(_redact_agent_params_tree(item, _depth + 1) for item in typed_items) 

189 if not isinstance(value, dict): 

190 return value 

191 typed_params: Final = _AGENT_PARAMS_ADAPTER.validate_python(value) 

192 return { 

193 key: ( 

194 REDACTED_BY_LITELM_STRING 

195 if _AGENT_PARAMS_MASKER.is_sensitive_key(key) 

196 else _redact_agent_params_tree(nested_value, _depth + 1) 

197 ) 

198 for key, nested_value in typed_params.items() 

199 } 

200 

201 

202def parse_agent_litellm_params(value: object) -> Mapping[str, object]: 

203 """Normalize a stored litellm_params column to a read-only mapping. 

204 

205 The prisma Json column comes back as either an already-parsed dict or a 

206 JSON string depending on the read path, so handle both rather than 

207 assuming one. Only ever read from (merge-source lookups), never mutated 

208 or re-serialized directly, so a read-only view is enough here. 

209 """ 

210 if isinstance(value, str): 210 ↛ 211line 210 didn't jump to line 211 because the condition on line 210 was never true

211 try: 

212 return _AGENT_PARAMS_ADAPTER.validate_json(value) 

213 except ValidationError: 

214 return _EMPTY_LITELLM_PARAMS 

215 if isinstance(value, Mapping): 215 ↛ 220line 215 didn't jump to line 220 because the condition on line 215 was always true

216 try: 

217 return _AGENT_PARAMS_ADAPTER.validate_python(value) 

218 except ValidationError: 

219 return _EMPTY_LITELLM_PARAMS 

220 return _EMPTY_LITELLM_PARAMS 

221 

222 

223_KILL_SWITCH_ADAPTER: Final[TypeAdapter[AgentKillSwitchConfig | None]] = TypeAdapter(AgentKillSwitchConfig | None) 

224 

225 

226def parse_agent_kill_switch(value: object) -> AgentKillSwitchConfig | None: 

227 if value is None: 

228 return None 

229 try: 

230 if isinstance(value, str): 230 ↛ 231line 230 didn't jump to line 231 because the condition on line 230 was never true

231 return _KILL_SWITCH_ADAPTER.validate_json(value) 

232 return _KILL_SWITCH_ADAPTER.validate_python(value) 

233 except ValidationError: 

234 return None 

235 

236 

237def serialize_agent_kill_switch(incoming: object, existing: object) -> str: 

238 """prisma-client-py drops ``None`` from update data, so a cleared kill switch is stored as the JSON literal 

239 ``null`` (read back as ``None``), the same convention ``memory_endpoints`` uses for ``Json?`` columns.""" 

240 restored: Final = restore_kill_switch( 

241 _KILL_SWITCH_ADAPTER.validate_python(incoming), parse_agent_kill_switch(existing) 

242 ) 

243 return safe_dumps(restored.model_dump() if restored is not None else None) 

244 

245 

246_MISSING_AGENT_PARAM: Final = object() 

247_RESTORE_AGENT_PARAMS_MAX_DEPTH: Final = 10 

248 

249 

250def _restore_redacted_nested_value(incoming_value: object, existing_value: object, _depth: int) -> object: 

251 """Recurse into a non-sensitively-named dict/list value so a secret 

252 nested underneath it (e.g. inside a list of per-provider configs) is 

253 still restored, not just top-level keys. Mirrors the shapes 

254 ``redact_sensitive_agent_litellm_params`` recurses into on read, so 

255 restore and redact stay symmetric. 

256 

257 List elements are paired with the existing list by position: with no 

258 stable per-element identity in an arbitrary ``dict[str, object]`` schema, 

259 index is the same correspondence every other part of this restore (and 

260 the endpoints' existing full-replace-on-PUT semantics) already assumes. 

261 This correctly preserves a masked secret across an ordinary edit of that 

262 same entry's other fields; it does not protect against a caller who both 

263 reorders/resizes the list AND echoes back a masked marker in the same 

264 request, which is a known, narrow limitation (see LIT-6736 PR discussion) 

265 rather than a cross-entry credential leak in the common case. 

266 

267 A value collapsed to the flat marker by the read side's depth cap is 

268 recovered wholesale from ``existing_value`` (rather than the marker 

269 string itself getting persisted) whenever ``existing_value`` isn't 

270 already that same flat marker. Depth-bounded like its read-side 

271 counterpart; a value at the cap is returned unchanged rather than 

272 corrupted. 

273 """ 

274 if incoming_value == REDACTED_BY_LITELM_STRING and existing_value != REDACTED_BY_LITELM_STRING: 274 ↛ 275line 274 didn't jump to line 275 because the condition on line 274 was never true

275 return existing_value 

276 if _depth >= _RESTORE_AGENT_PARAMS_MAX_DEPTH: 276 ↛ 277line 276 didn't jump to line 277 because the condition on line 276 was never true

277 return incoming_value 

278 if isinstance(incoming_value, Mapping): 

279 typed_incoming_map: Final = _AGENT_PARAMS_ADAPTER.validate_python(incoming_value) 

280 existing_map: Final = ( 

281 _AGENT_PARAMS_ADAPTER.validate_python(existing_value) 

282 if isinstance(existing_value, Mapping) 

283 else _EMPTY_LITELLM_PARAMS 

284 ) 

285 return _restore_redacted_litellm_params(typed_incoming_map, existing_map, _depth + 1) 

286 if isinstance(incoming_value, (list, tuple)): 

287 typed_incoming_seq: Final = _AGENT_PARAMS_SEQUENCE_ADAPTER.validate_python(incoming_value) 

288 existing_seq: Final = ( 

289 _AGENT_PARAMS_SEQUENCE_ADAPTER.validate_python(existing_value) 

290 if isinstance(existing_value, (list, tuple)) 

291 else () 

292 ) 

293 return tuple( 

294 _restore_redacted_nested_value( 

295 item, 

296 existing_seq[index] if index < len(existing_seq) else None, 

297 _depth + 1, 

298 ) 

299 for index, item in enumerate(typed_incoming_seq) 

300 ) 

301 return incoming_value 

302 

303 

304def _resolved_agent_param_value( 

305 key: str, 

306 incoming: Mapping[str, object], 

307 existing: Mapping[str, object], 

308 _depth: int, 

309) -> object: 

310 """The value ``key`` should end up with in a restored litellm_params, or 

311 ``_MISSING_AGENT_PARAM`` when it should be dropped entirely.""" 

312 if key in incoming: 312 ↛ 317line 312 didn't jump to line 317 because the condition on line 312 was always true

313 value: Final = incoming[key] 

314 if _AGENT_PARAMS_MASKER.is_sensitive_key(key): 314 ↛ 315line 314 didn't jump to line 315 because the condition on line 314 was never true

315 return existing.get(key, _MISSING_AGENT_PARAM) if value == REDACTED_BY_LITELM_STRING else value 

316 return _restore_redacted_nested_value(value, existing.get(key), _depth) 

317 if _AGENT_PARAMS_MASKER.is_sensitive_key(key): 

318 return existing.get(key, _MISSING_AGENT_PARAM) 

319 return _MISSING_AGENT_PARAM 

320 

321 

322def _patched_access_group_ids(agent: PatchAgentRequest) -> Mapping[str, object]: 

323 if "access_group_ids" not in agent: 

324 return MappingProxyType({}) 

325 return MappingProxyType({"access_group_ids": tuple(dict.fromkeys(agent.get("access_group_ids") or ()))}) 

326 

327 

328def _patched_kill_switch(agent: PatchAgentRequest, existing: object) -> Mapping[str, object]: 

329 if "kill_switch" not in agent: 

330 return MappingProxyType({}) 

331 return MappingProxyType({"kill_switch": serialize_agent_kill_switch(agent.get("kill_switch"), existing)}) 

332 

333 

334def _restore_redacted_litellm_params( 

335 incoming: Mapping[str, object], 

336 existing: Mapping[str, object], 

337 _depth: int = 0, 

338) -> dict[str, object]: 

339 """Restore the real credential behind any litellm_params value the caller 

340 echoed back as ``REDACTED_BY_LITELM_STRING``, and behind any sensitive key 

341 omitted entirely, so an edit to an unrelated field never overwrites (or 

342 silently drops) a stored provider credential -- the UI never has to 

343 read-and-resend a secret to keep it. Recurses into nested dicts and lists 

344 so a secret nested under a non-sensitively-named key is restored too. 

345 

346 A sensitive key given a real (non-marker) value, including an explicit 

347 empty string, is treated as a deliberate update -- that's how a caller 

348 clears a credential. Non-sensitive keys always take the incoming value 

349 (recursed into), matching the endpoints' existing full-replace-on-PUT / 

350 merge-on-PATCH semantics for everything that isn't a secret. 

351 """ 

352 all_keys: Final = frozenset(incoming) | frozenset(existing) 

353 return { 

354 key: value 

355 for key in all_keys 

356 if (value := _resolved_agent_param_value(key, incoming, existing, _depth)) is not _MISSING_AGENT_PARAM 

357 } 

358 

359 

360class GrantMigrationResult(NamedTuple): 

361 rewritten: int 

362 missed: int 

363 

364 

365class AgentRegistry: 

366 def __init__(self): 

367 self.agent_list: list[AgentResponse] = [] 

368 self.config_agents: tuple[AgentConfig, ...] = () 

369 self.config_agent_legacy_ids: Mapping[str, str] = MappingProxyType({}) 

370 

371 def reset_agent_list(self): 

372 self.agent_list = [] 

373 

374 def register_agent(self, agent_config: AgentResponse): 

375 self.agent_list.append(agent_config) 

376 

377 def deregister_agent(self, agent_name: str): 

378 self.agent_list = [agent for agent in self.agent_list if agent.agent_name != agent_name] 

379 

380 def get_agent_list(self, agent_names: Sequence[str] | None = None) -> tuple[AgentResponse, ...]: 

381 if agent_names is not None: 381 ↛ 382line 381 didn't jump to line 382 because the condition on line 381 was never true

382 return tuple(agent for agent in self.agent_list if agent.agent_name in agent_names) 

383 return tuple(self.agent_list) 

384 

385 def get_public_agent_list(self) -> tuple[AgentResponse, ...]: 

386 public_agent_groups: Final = litellm.public_agent_groups 

387 if public_agent_groups is None: 

388 return () 

389 return tuple( 

390 agent for agent in self.agent_list if not self.ids_for_agent(agent.agent_id).isdisjoint(public_agent_groups) 

391 ) 

392 

393 def _create_agent_id(self, agent_config: AgentConfig) -> str: 

394 return hashlib.sha256(agent_config["agent_name"].encode()).hexdigest() 

395 

396 def _create_legacy_agent_id(self, agent_config: AgentConfig) -> str: 

397 return hashlib.sha256(json.dumps(agent_config, sort_keys=True).encode()).hexdigest() 

398 

399 def ids_for_agent(self, agent_id: str) -> frozenset[str]: 

400 return frozenset( 

401 {agent_id, *(legacy for legacy, stable in self.config_agent_legacy_ids.items() if stable == agent_id)} 

402 ) 

403 

404 def stable_agent_id(self, agent_id: str) -> str: 

405 return self.config_agent_legacy_ids.get(agent_id, agent_id) 

406 

407 def load_agents_from_config(self, agent_config: Sequence[AgentConfig] | None = None): 

408 """ 

409 Register the agents declared in config.yaml and remember them for later rebuilds. 

410 

411 A config entry is skipped when its ``agent_name`` is already registered, so a 

412 database record always wins over a config entry that reuses its name and the 

413 registry never holds two agents under one name. Enforcing that here rather than 

414 in the caller keeps the guarantee independent of the order the two sources load 

415 in. Passing ``None`` leaves the remembered agents untouched; passing an empty 

416 sequence clears them. 

417 """ 

418 if agent_config is None: 418 ↛ 419line 418 didn't jump to line 419 because the condition on line 418 was never true

419 return 

420 

421 for agent_config_item in agent_config: 421 ↛ 422line 421 didn't jump to line 422 because the loop on line 421 never started

422 if not isinstance(agent_config_item, dict): 

423 raise ValueError("agent_config must be a list of dictionaries") 

424 

425 self.config_agents = tuple(agent_config) 

426 self.config_agent_legacy_ids = MappingProxyType( 

427 { 

428 self._create_legacy_agent_id(agent_config_item): self._create_agent_id(agent_config_item) 

429 for agent_config_item in agent_config 

430 if agent_config_item.get("agent_name") and agent_config_item.get("agent_card_params") 

431 } 

432 ) 

433 

434 for agent_config_item in agent_config: 434 ↛ 435line 434 didn't jump to line 435 because the loop on line 434 never started

435 agent_name = agent_config_item.get("agent_name") 

436 agent_card_params = agent_config_item.get("agent_card_params") 

437 if not all([agent_name, agent_card_params]): 

438 continue 

439 

440 if any(agent.agent_name == agent_name for agent in self.agent_list): 

441 continue 

442 

443 # create a stable hash id for config item 

444 config_hash = self._create_agent_id(agent_config_item) 

445 

446 self.register_agent(agent_config=AgentResponse(agent_id=config_hash, **agent_config_item)) 

447 

448 def load_agents_from_db_and_config( 

449 self, 

450 agent_config: Sequence[AgentConfig] | None = None, 

451 db_agents: Sequence[Mapping[str, object]] | None = None, 

452 ): 

453 """ 

454 Rebuild the registry from the DB rows plus the agents declared in config.yaml. 

455 

456 ``agent_config`` defaults to the agents remembered by the last 

457 ``load_agents_from_config`` call, so a periodic DB reload does not drop 

458 config-defined agents. 

459 

460 The DB rows are registered first so that a config entry reusing one of their 

461 names is dropped by ``load_agents_from_config``, mirroring how config-declared 

462 MCP servers are unioned under the database registry. Name lookups and 

463 deregistration both address a single agent, so the registry must never hold two 

464 under one name. 

465 """ 

466 self.reset_agent_list() 

467 

468 if db_agents: 

469 for db_agent in db_agents: 

470 if not isinstance(db_agent, dict): 470 ↛ 471line 470 didn't jump to line 471 because the condition on line 470 was never true

471 raise ValueError("db_agents must be a list of dictionaries") 

472 

473 self.register_agent(agent_config=AgentResponse.model_validate(db_agent)) 

474 

475 self.load_agents_from_config(agent_config if agent_config is not None else self.config_agents) 

476 return self.agent_list 

477 

478 async def migrate_legacy_grant_ids( 

479 self, table: "TableActions[prisma_models.LiteLLM_ObjectPermissionTable]" 

480 ) -> GrantMigrationResult: 

481 """ 

482 Rewrite object_permission.agents rows holding a legacy full-entry hash to the 

483 stable name-derived id. 

484 

485 Only the running proxy can do this: the legacy hash is computed from the 

486 resolved config entry (secrets included), so no SQL migration can know it. 

487 Persisting the stable id here is what keeps a grant alive across a later 

488 secret rotation, which re-mints the legacy hash and would otherwise orphan 

489 the stored value. Idempotent; runs of it after the first find no rows. 

490 

491 Each write is a compare-and-swap against the agents array read above, so a 

492 grant edited concurrently is left untouched; the runtime alias keeps covering 

493 it and the next boot retries the rewrite. 

494 """ 

495 legacy_ids: Final = tuple(legacy for legacy, stable in self.config_agent_legacy_ids.items() if legacy != stable) 

496 if not legacy_ids: 496 ↛ 498line 496 didn't jump to line 498 because the condition on line 496 was always true

497 return GrantMigrationResult(rewritten=0, missed=0) 

498 rows: Final = await table.find_many(where={"agents": {"has_some": legacy_ids}}) 

499 updates: Final = tuple( 

500 ( 

501 row.object_permission_id, 

502 tuple(row.agents or ()), 

503 tuple(dict.fromkeys(self.stable_agent_id(agent_id) for agent_id in row.agents or ())), 

504 ) 

505 for row in rows 

506 ) 

507 counts: Final = await asyncio.gather( 

508 *( 

509 table.update_many( 

510 where={"object_permission_id": object_permission_id, "agents": {"equals": snapshot_agents}}, 

511 data={"agents": translated_agents}, 

512 ) 

513 for object_permission_id, snapshot_agents, translated_agents in updates 

514 ) 

515 ) 

516 rewritten: Final = sum(counts) 

517 return GrantMigrationResult(rewritten=rewritten, missed=len(updates) - rewritten) 

518 

519 ########################################################### 

520 ########### DB management helpers for agents ########### 

521 ############################################################ 

522 async def add_agent_to_db( 

523 self, 

524 agent: AgentConfig, 

525 prisma_client: PrismaClient, 

526 created_by: str, 

527 agent_id: str | None = None, 

528 ) -> AgentResponse: 

529 """ 

530 Add an agent to the database. 

531 

532 If ``agent_id`` is provided, it is used as the primary key for the new 

533 row (otherwise the DB generates a UUID). Callers pass an explicit ID 

534 when the agent_card_params must reference the agent's own URL before 

535 the row exists, e.g. the A2A merge in ``create_agent``. 

536 """ 

537 try: 

538 agent_name: Final = agent.get("agent_name") 

539 

540 # Serialize litellm_params. A create has no stored row to restore a 

541 # secret behind, so a sensitive key submitted as the redaction 

542 # marker (e.g. a stray client re-post) is dropped rather than 

543 # persisted as the literal placeholder string. 

544 litellm_params_obj: Final = agent.get("litellm_params", {}) 

545 litellm_params_dict: Final = _restore_redacted_litellm_params( 

546 _dump_agent_params(litellm_params_obj), _EMPTY_LITELLM_PARAMS 

547 ) 

548 litellm_params: Final[str] = safe_dumps(litellm_params_dict) 

549 

550 # Serialize agent_card_params 

551 agent_card_params_obj: Final = agent.get("agent_card_params", {}) 

552 agent_card_params_dict: Final[dict[str, object]] = _dump_agent_params(agent_card_params_obj) 

553 agent_card_params: Final[str] = safe_dumps(agent_card_params_dict) 

554 

555 # Handle object_permission (MCP tool access for agent) 

556 object_permission_id: str | None = None 

557 if agent.get("object_permission") is not None: 

558 agent_copy: Final = dict(agent) 

559 object_permission_id = await handle_update_object_permission_common(agent_copy, None, prisma_client) 

560 

561 # Serialize static_headers 

562 static_headers_obj: Final = agent.get("static_headers") 

563 static_headers_val: Final[str | None] = safe_dumps(dict(static_headers_obj)) if static_headers_obj else None 

564 

565 extra_headers_val: Final = agent.get("extra_headers") 

566 access_group_ids_val: Final = agent.get("access_group_ids") 

567 

568 create_data: Final[dict[str, object]] = { 

569 "agent_name": agent_name, 

570 "litellm_params": litellm_params, 

571 "agent_card_params": agent_card_params, 

572 "kill_switch": serialize_agent_kill_switch(agent.get("kill_switch"), None), 

573 "created_by": created_by, 

574 "updated_by": created_by, 

575 "created_at": datetime.now(timezone.utc), 

576 "updated_at": datetime.now(timezone.utc), 

577 } 

578 if agent_id is not None: 578 ↛ 580line 578 didn't jump to line 580 because the condition on line 578 was always true

579 create_data["agent_id"] = agent_id 

580 if static_headers_val is not None: 

581 create_data["static_headers"] = static_headers_val 

582 if extra_headers_val is not None: 

583 create_data["extra_headers"] = extra_headers_val 

584 if access_group_ids_val is not None: 

585 create_data["access_group_ids"] = tuple(dict.fromkeys(access_group_ids_val)) 

586 if object_permission_id is not None: 

587 create_data["object_permission_id"] = object_permission_id 

588 

589 for rate_field in ( 

590 "tpm_limit", 

591 "rpm_limit", 

592 "session_tpm_limit", 

593 "session_rpm_limit", 

594 ): 

595 _val = agent.get(rate_field) 

596 if _val is not None: 

597 create_data[rate_field] = _val 

598 

599 # Create agent in DB 

600 created_agent: Final = await agents_table(prisma_client).create( 

601 data=create_data, 

602 include={"object_permission": True}, 

603 ) 

604 

605 created_agent_dict: Final = created_agent.model_dump() 

606 if created_agent.object_permission is not None: 

607 try: 

608 created_agent_dict["object_permission"] = created_agent.object_permission.model_dump() 

609 except Exception: 

610 created_agent_dict["object_permission"] = created_agent.object_permission.dict() 

611 return AgentResponse(**created_agent_dict) 

612 except Exception as e: 

613 raise Exception(f"Error adding agent to DB: {e}") 

614 

615 async def delete_agent_from_db(self, agent_id: str, prisma_client: PrismaClient) -> Mapping[str, object]: 

616 """ 

617 Delete an agent from the database 

618 """ 

619 try: 

620 deleted_agent: Final = await agents_table(prisma_client).delete(where={"agent_id": agent_id}) 

621 if deleted_agent is None: 621 ↛ 622line 621 didn't jump to line 622 because the condition on line 621 was never true

622 raise ValueError(f"Agent not found, passed agent_id={agent_id}") 

623 return dict(deleted_agent) 

624 except Exception as e: 

625 raise Exception(f"Error deleting agent from DB: {e}") 

626 

627 async def patch_agent_in_db( 

628 self, 

629 agent_id: str, 

630 agent: PatchAgentRequest, 

631 prisma_client: PrismaClient, 

632 updated_by: str, 

633 ) -> AgentResponse: 

634 """ 

635 Patch an agent in the database. 

636 

637 Get the existing agent from the database and patch it with the new values. 

638 

639 Args: 

640 agent_id: The ID of the agent to patch 

641 agent: The new agent values to patch 

642 prisma_client: The Prisma client to use 

643 updated_by: The user ID of the user who is patching the agent 

644 

645 Returns: 

646 The patched agent 

647 """ 

648 try: 

649 existing_record: Final = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id}) 

650 if existing_record is None: 650 ↛ 651line 650 didn't jump to line 651 because the condition on line 650 was never true

651 raise Exception(f"Agent with ID {agent_id} not found") 

652 existing_agent: Final[Mapping[str, object]] = dict(existing_record) 

653 

654 augment_agent: Final = {**existing_agent, **agent} 

655 update_data: Final[dict[str, object]] = { 

656 **_patched_access_group_ids(agent), 

657 **_patched_kill_switch(agent, existing_agent.get("kill_switch")), 

658 } 

659 if augment_agent.get("agent_name"): 

660 update_data["agent_name"] = augment_agent.get("agent_name") 

661 if "litellm_params" in agent: 

662 existing_litellm_params: Final = parse_agent_litellm_params(existing_agent.get("litellm_params")) 

663 update_data["litellm_params"] = safe_dumps( 

664 _restore_redacted_litellm_params( 

665 _dump_agent_params(agent.get("litellm_params") or _EMPTY_LITELLM_PARAMS), 

666 existing_litellm_params, 

667 ) 

668 ) 

669 if augment_agent.get("agent_card_params"): 669 ↛ 672line 669 didn't jump to line 672 because the condition on line 669 was always true

670 update_data["agent_card_params"] = safe_dumps(augment_agent.get("agent_card_params")) 

671 

672 for rate_field in ( 

673 "tpm_limit", 

674 "rpm_limit", 

675 "session_tpm_limit", 

676 "session_rpm_limit", 

677 ): 

678 if rate_field in agent: 

679 update_data[rate_field] = agent.get(rate_field) 

680 if "static_headers" in agent: 

681 headers_value: Final = agent.get("static_headers") 

682 update_data["static_headers"] = safe_dumps(dict(headers_value) if headers_value is not None else {}) 

683 if "extra_headers" in agent: 

684 extra_headers_value: Final = agent.get("extra_headers") 

685 update_data["extra_headers"] = extra_headers_value if extra_headers_value is not None else [] 

686 if agent.get("object_permission") is not None: 

687 agent_copy: Final = dict(augment_agent) 

688 existing_object_permission_id: Final = existing_record.object_permission_id 

689 object_permission_id: Final = await handle_update_object_permission_common( 

690 agent_copy, 

691 existing_object_permission_id, 

692 prisma_client, 

693 ) 

694 if object_permission_id is not None: 694 ↛ 697line 694 didn't jump to line 697 because the condition on line 694 was always true

695 update_data["object_permission_id"] = object_permission_id 

696 # Patch agent in DB 

697 patched_agent: Final = await agents_table(prisma_client).update( 

698 where={"agent_id": agent_id}, 

699 data={ 

700 **update_data, 

701 "updated_by": updated_by, 

702 "updated_at": datetime.now(timezone.utc), 

703 }, 

704 include={"object_permission": True}, 

705 ) 

706 if patched_agent is None: 706 ↛ 707line 706 didn't jump to line 707 because the condition on line 706 was never true

707 raise ValueError(f"Agent not found, passed agent_id={agent_id}") 

708 patched_agent_dict: Final = patched_agent.model_dump() 

709 if patched_agent.object_permission is not None: 

710 try: 

711 patched_agent_dict["object_permission"] = patched_agent.object_permission.model_dump() 

712 except Exception: 

713 patched_agent_dict["object_permission"] = patched_agent.object_permission.dict() 

714 return AgentResponse(**patched_agent_dict) 

715 except Exception as e: 

716 raise Exception(f"Error patching agent in DB: {e}") 

717 

718 async def update_agent_in_db( 

719 self, 

720 agent_id: str, 

721 agent: AgentConfig, 

722 prisma_client: PrismaClient, 

723 updated_by: str, 

724 ) -> AgentResponse: 

725 """ 

726 Update an agent in the database 

727 """ 

728 try: 

729 agent_name: Final = agent.get("agent_name") 

730 

731 # A PUT fully replaces litellm_params from the request body, so the 

732 # existing row is read up front to restore any sensitive key the 

733 # caller echoed back redacted (or omitted) rather than persisting 

734 # the marker -- or nothing -- over the real stored credential. 

735 existing_row: Final = await agents_table(prisma_client).find_unique( 

736 where={"agent_id": agent_id} # mutable-ok: prisma's query builder rejects a Mapping/MappingProxyType 

737 ) 

738 existing_litellm_params: Final = parse_agent_litellm_params( 

739 existing_row.litellm_params if existing_row is not None else None 

740 ) 

741 

742 # Serialize litellm_params 

743 litellm_params_obj: Final = agent.get("litellm_params", {}) 

744 litellm_params_dict: Final = _restore_redacted_litellm_params( 

745 _dump_agent_params(litellm_params_obj), existing_litellm_params 

746 ) 

747 litellm_params: Final[str] = safe_dumps(litellm_params_dict) 

748 

749 # Serialize agent_card_params 

750 agent_card_params_obj: Final = agent.get("agent_card_params", {}) 

751 agent_card_params_dict: Final[dict[str, object]] = _dump_agent_params(agent_card_params_obj) 

752 agent_card_params: Final[str] = safe_dumps(agent_card_params_dict) 

753 

754 # Serialize static_headers for update 

755 static_headers_obj_u: Final = agent.get("static_headers") 

756 static_headers_val_u: Final[str] = ( 

757 safe_dumps(dict(static_headers_obj_u)) if static_headers_obj_u is not None else safe_dumps({}) 

758 ) 

759 extra_headers_val_u: Final = agent.get("extra_headers") or [] 

760 access_group_ids_val_u: Final = tuple(dict.fromkeys(agent.get("access_group_ids") or ())) 

761 kill_switch_val_u: Final = serialize_agent_kill_switch( 

762 agent.get("kill_switch"), existing_row.kill_switch if existing_row is not None else None 

763 ) 

764 

765 update_data: Final[dict[str, object]] = { 

766 "agent_name": agent_name, 

767 "litellm_params": litellm_params, 

768 "agent_card_params": agent_card_params, 

769 "static_headers": static_headers_val_u, 

770 "extra_headers": extra_headers_val_u, 

771 "kill_switch": kill_switch_val_u, 

772 "access_group_ids": access_group_ids_val_u, 

773 "updated_by": updated_by, 

774 "updated_at": datetime.now(timezone.utc), 

775 } 

776 

777 for rate_field in ( 

778 "tpm_limit", 

779 "rpm_limit", 

780 "session_tpm_limit", 

781 "session_rpm_limit", 

782 ): 

783 _val = agent.get(rate_field) 

784 if _val is not None: