fix(access groups): sync assigned_key_ids from the key write paths - #36843
Conversation
|
|
Greptile SummaryThe PR synchronizes the two sides of virtual-key/access-group membership across key creation, updates, deletion, and regeneration.
Confidence Score: 5/5The PR appears safe to merge. No blocking failure remains.
|
| Filename | Overview |
|---|---|
| litellm/proxy/management_helpers/access_group_key_sync.py | Implements atomic guarded attach, detach, and regeneration-repoint operations with targeted access-group cache invalidation; the previously reported concurrency defects are addressed. |
| litellm/proxy/management_endpoints/key_management_endpoints.py | Integrates reverse membership synchronization into all covered key lifecycle paths after credential-cache invalidation. |
| tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py | Adds focused mocked coverage for lifecycle synchronization, authorization effects, idempotency, live-state regeneration, exact SQL write shapes, and invalidation ordering. |
Reviews (6): Last reviewed commit: "fix(access groups): batch membership wri..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 4 · PR risk: 0/10 |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
5764bc3 to
9b39502
Compare
|
@greptileai please re-review 9b39502. Took the non-atomic finding: attach now uses Prisma's atomic push. Reasoning is in the description under Review notes. |
9b39502 to
edf0130
Compare
|
@greptileai please re-review edf0130. Detach now uses an atomic array_remove, and the sync moved after credential invalidation. Reasoning under Review notes. |
edf0130 to
149dcfc
Compare
|
@greptileai please re-review 149dcfc. Attach is now a guarded array_append, so concurrent identical attaches cannot duplicate. Measured on Postgres 16. |
055f0e6 to
7f44000
Compare
|
@greptileai please re-review 7f44000. Took the last open finding: the sync now runs after credential invalidation on the update paths too. |
|
@greptileai please re-review 395bdee. Both remaining findings taken: set-based membership writes, and regeneration swaps in place with no snapshot. |
An access group and a key each store their half of the same relationship, and both halves are read: the grant check in auth_checks authorizes a key only when the group's assigned_key_ids lists that key's token. The access-group endpoints maintain both halves, but nothing on the key side ever wrote the group's copy, so attaching a group through /key/generate or /key/update recorded the key's half and left the group's half empty. The 200 came back with the group attached and the key still got none of the group's models. Adds a reverse sync the key write paths call: create, update (which also covers both bulk update routes), regenerate, and delete. Regeneration replaces the token, so the group's copy is re-pointed from the old hash to the new one. Detaching withdraws the token, and the group's cache entry is invalidated on every write so auth stops reading a stale row.
…from live rows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
395bdee to
17a1b6c
Compare
efbdb69
into
litellm_internal_staging
3 participants
TLDR
Problem this solves:
How it solves it:
User Flow
Before: an admin grants a team-restricted model to one key through an access group, the API accepts it, and the key still cannot call the model
{"access_group_name": "premium", "access_model_names": ["claude-haiku-4-5"]}and get back 201 with anaccess_group_idgpt-4o-mini, then POST https://litellm-domain/key/generate on that team, and get back a keyclaude-haiku-4-5with that key and gets 400 "team not allowed to access model", which is expected at this point{"key": "sk-...", "access_group_ids": ["<id>"]}and gets 200 back showing the group attached"assigned_key_ids": [], so the access group page shows the group attached to no keys at allclaude-haiku-4-5and gets the identical 400. The grant they were told they had never took effectAfter: the same steps, and the key can call the model the group grants
"assigned_key_ids"carrying that key, so the access group page lists itclaude-haiku-4-5and gets 200 with a real completion"access_group_ids": []takes it away again, and the group's page stops listing the keyRelevant issues
Linear ticket
Resolves LIT-5521
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
Same proxy, same Postgres, same commands on both legs, against real OpenAI and Anthropic APIs. The team is restricted to one model and the access group grants the other, so the only thing that can change the outcome is whether the group's own list learned about the key.
Before, captured at
262ed530f8(staging without this commit,grep -c sync_key_access_group_membershipreturns 0):After, captured at
7f44000410:The other three write paths, same proxy and same commit, so the group's list is verified live through the whole key lifecycle rather than only through /key/update:
Re-run at
17a1b6ce12, the commit that batches the writes and drops the regeneration snapshot. Same live proxy and Postgres, real Anthropic calls, team restricted toanthropic-haiku-4-5and the group grantinganthropic-sonnet-4-5. Attach still grants on the very next request, a cache-warmed detach still denies on the very next request with no restart, and create, re-attach, regenerate and delete all still behave as above. What is new is the cost and the concurrency case.Five groups attached and detached in one /key/update each way, with
log_statement=allshowing one statement per direction carrying all five ids in a single array parameter:Regeneration against a group edited out of band, run twice with only the helper swapped, so the difference is the fix and nothing else. The key row still lists G1 while the old hash has been moved into G2 directly in Postgres, then the key is regenerated:
Type
🐛 Bug Fix
Caveats
Review notes
Why the value is recorded rather than derived. The group's list and the key's list are two halves of one handshake, not a denormalized cache: a key is authorized only when it names the group AND the group names the key. Deriving the group's half from "every key naming this group" collapses that to a single side and grants strictly more than today.
tests/proxy_unit_tests/test_auth_checks.pypins the case a derivation would break, a key naming a group whose own list excludes that key's token, and asserts denial.Where the ticket and the code disagree. LIT-5521 says a key detached from the key side keeps receiving the group's grants. It does not: every reader gates on the key's own
access_group_idsfirst, so a detach revokes immediately, and the leftover entry is a stale listing rather than a live grant. The authorization defect runs the other way, on attach, which is what the proof above captures. The detach direction is still fixed here, for the listing and so a regenerated or deleted key stops leaving entries behind.What is covered and what is not. Every key write path that can change membership is covered: create (which serves /key/generate, /key/service_account/generate and the key /user/new mints), update (one shared helper reached by /key/update and by both bulk routes through
_process_single_key_update), regenerate, and delete. Three cascade paths delete key rows in bulk without going through the key delete helper, /team/delete, /user/delete, and the org path behind it, and they still leave stale token hashes. There is no authorization consequence there, the key row is gone so the token cannot authenticate, and deleting the access group already self-heals out-of-sync entries. Worth a follow-up, deliberately not folded in here.Greptile's non-atomic membership finding, taken. I reproduced it against a real Postgres 16 rather than reasoning about it, and the result changed what the fix should be: two concurrent read-modify-writes on the same row lose one membership even when both run inside a transaction, because under READ COMMITTED the second write is derived from a value read before the first committed. So wrapping the two writes in a transaction would have looked like a fix without being one. The same two writes expressed as a single
array_appendkeep both. Both halves are now single guarded statements,array_append ... WHERE NOT (token = ANY(...))andarray_remove ... WHERE token = ANY(...), which is whatTeamRepository.remove_memberalready tells you to reach for when the read-modify-write matters. It matters here: a lost detach puts an already revoked token back into a group and restores its grants, and a lost attach silently drops a grant an admin just made. TheWHEREguards also make each statement idempotent, which a check-then-push cannot be. I measured all three shapes on Postgres 16 before choosing: two concurrent read-modify-writes leave{keyA}and lose the other, two guarded appends of different tokens leave{keyA,keyB}, and two concurrent appends of the SAME token leave{tok}guarded versus{tok,tok}when the check and the append are separate statements. Tests pin both write shapes, so a regression to array replacement in either direction fails, and the mocked executor emulates the guards rather than assuming them. On partial failure the group write throwing after the key row commits leaves exactly the pre-PR state and surfaces as an error to the caller, so nothing is newly corrupted, anddelete_access_groupalready reconciles out-of-sync entries by unioning both sides.Greptile's revocation-ordering finding, taken. The sync ran before credential-cache invalidation on both delete and regenerate, so a failure inside it would have returned an error with the old key still authenticating from cache. All four call sites now run after invalidation, update and bulk update included, so nothing fallible sits between changing a key and revoking what it had. The cached auth object still carries the key's old
access_group_ids, so syncing first meant a failed sync returned an error with the key still authenticating against groups it had just lost. A test pins the ordering.veria's unbounded-work finding, taken. The caller picks the size of
access_group_ids, and one statement per group let a single /key/update hold a worker and a connection for as many sequential writes as the list is long, even for ids that do not exist. Both directions are now set-based,WHERE access_group_id = ANY($2::text[]), so the cost is one statement for the attaches and one for the detaches whatever the list size, with the guards and idempotence unchanged. Each statement returns the ids it actually moved, so cache invalidation stays scoped to the groups that changed rather than every id in the request. A cardinality cap was the other option and it is strictly worse here: it puts an arbitrary number in the request schema, breaks callers legitimately assigning many groups, and would still leave the per-group loop behind it. A test drives 60 attaches and 60 detaches through /key/update and asserts exactly two statements, so a regression to per-group writes fails.Greptile's regeneration-transition finding, taken. Regeneration was replaying the key row it read before the new token existed, so an access-group edit landing in between was either resurrected, the group had dropped the key and the replay put the new token back, or lost, the group had just attached the key and kept the dead hash. It now swaps in place from whatever the rows hold at write time,
SET assigned_key_ids = array_append(array_remove(array_remove(assigned_key_ids, $1), $2), $2) WHERE $1 = ANY(assigned_key_ids), which needs no snapshot and no lock, and the doublearray_removekeeps a retry from duplicating the new token. A regenerate request that also changesaccess_group_idsstill applies that delta afterwards, since that intent is expressed against the key's own list. The test sets the group rows to disagree with the key row in both directions and asserts the revoked group stays empty while the newly attached one follows the new token.Consistency with the team-side fix. This mirrors PR #36825 field for field: same helper module shape, same previous/updated delta signature, same cache invalidation, and the same reason for living outside
access_group_endpoints.py, which is a lazily registered feature router that must not be imported eagerly. It does not import from that PR's module, so the two can merge in either order.Final Attestation
Link to Devin session: https://app.devin.ai/sessions/d4faac79f9a84e0fba54b0c90f2265fc
Requested by: @yassin-berriai