From 765804aa696c4feda5cafe4971ee96edad9b7452 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 24 Apr 2026 16:45:15 -0700 Subject: [PATCH] fix(memory): require team admin to modify pure team rows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tightens the write-authorization rule for "pure team rows" (rows with no user_id stamped, only team_id) to match the pattern used by team-management endpoints (`_is_user_team_admin` + `_is_user_org_admin_for_team`): - Plain team members can READ team rows via the OR visibility filter (intentional, unchanged). - Only PROXY_ADMIN, team admins of the row's team_id, or org admins for the team's organization may MODIFY them. Plain members get 403. `_assert_write_access` is now async and takes the prisma_client so it can fetch the team and run the existing `_is_user_team_admin` / `_is_user_org_admin_for_team` helpers from `litellm.proxy.management_endpoints.common_utils`. The org-admin path is best-effort: it calls `get_user_object`, which depends on the proxy_server module being initialized, so any exception there is treated as "not an org admin" rather than crashing the request. Tests: - team admin can modify pure team row → 200 - plain team member cannot modify pure team row → 403 - plain team member cannot delete pure team row → 403 Updates the test fake to add a tiny `litellm_teamtable.find_unique` implementation and a `_make_team(team_id, admin_user_ids=[...])` helper. 27/27 unit tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/proxy/memory/memory_endpoints.py | 86 +++++++++++++++---- .../proxy/memory/test_memory_endpoints.py | 85 +++++++++++++++++- 2 files changed, 152 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py index 0d536903f85..39c5f03031f 100644 --- a/litellm/proxy/memory/memory_endpoints.py +++ b/litellm/proxy/memory/memory_endpoints.py @@ -98,7 +98,9 @@ def _internal_error( return HTTPException(status_code=500, detail=default_detail) -def _assert_write_access(row: Any, user_api_key_dict: UserAPIKeyAuth) -> None: +async def _assert_write_access( + prisma_client: Any, row: Any, user_api_key_dict: UserAPIKeyAuth +) -> None: """ Enforce ownership for mutations (PUT/DELETE). @@ -108,31 +110,83 @@ def _assert_write_access(row: Any, user_api_key_dict: UserAPIKeyAuth) -> None: this check, any team member could overwrite or delete a teammate's personal row whenever both `user_id` and `team_id` are stamped on it. - Rule (non-admin): the row must be authored by the caller's user_id, or be - a "pure team row" (no user_id stamped) within the caller's team. - Admins may modify any row. + Rules (mirroring how key/team management endpoints gate team-scoped writes): + - PROXY_ADMIN: always allowed. + - Personal ownership (`row.user_id == caller.user_id`): allowed. + - Pure team row (`row.user_id is None`, `row.team_id` set): + caller must be a team admin of `row.team_id` (members_with_roles entry + with `role == "admin"`), or an org admin for that team's organization. + Plain team members can only READ team rows, not modify them — same + pattern as `_validate_team_member_add_permissions` etc. + - Anything else: 403. """ if _is_admin(user_api_key_dict): return row_user_id = getattr(row, "user_id", None) row_team_id = getattr(row, "team_id", None) - # Personal ownership: caller authored this row. + # Personal ownership. if row_user_id and row_user_id == user_api_key_dict.user_id: return - # Pure team row (no user_id stamped) inside the caller's team. - if ( - row_user_id is None - and row_team_id is not None - and row_team_id == user_api_key_dict.team_id - ): - return + + # Pure team row — only team admins (or org admins) may write. + if row_user_id is None and row_team_id is not None: + if await _is_team_admin_for(prisma_client, user_api_key_dict, row_team_id): + return + raise HTTPException( status_code=403, detail="You do not have permission to modify this memory entry.", ) +async def _is_team_admin_for( + prisma_client: Any, user_api_key_dict: UserAPIKeyAuth, team_id: str +) -> bool: + """ + True if the caller is a team admin of `team_id`, or an org admin for the + team's organization. Mirrors the auth pattern used by team-management + endpoints (`_is_user_team_admin` + `_is_user_org_admin_for_team`). + + Imported lazily to avoid a circular import with proxy_server during the + memory router's module load. + """ + from litellm.proxy.management_endpoints.common_utils import ( + _is_user_org_admin_for_team, + _is_user_team_admin, + ) + + try: + team_obj = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) + except Exception as e: + verbose_proxy_logger.exception( + "Error loading team for write-auth check (team_id=%s): %s", team_id, e + ) + return False + if team_obj is None: + return False + + if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj): + return True + + # Org-admin path is best-effort: it pulls from the user cache via + # `get_user_object` which depends on the proxy_server module being + # initialized. In tests / non-proxy contexts that import path may fail — + # treat any error as "not an org admin" rather than crashing the request. + try: + if await _is_user_org_admin_for_team( + user_api_key_dict=user_api_key_dict, team_obj=team_obj + ): + return True + except Exception as e: + verbose_proxy_logger.debug( + "Org-admin check skipped during write-auth (team_id=%s): %s", team_id, e + ) + return False + + def _is_unique_violation(exc: Exception) -> bool: """ Detect a Prisma unique-constraint violation. @@ -398,7 +452,7 @@ async def upsert_memory( # owns this row (their user_id matches, or it's a pure team row in # their team) — otherwise a teammate could overwrite a personal # entry through the OR-based visibility filter. - _assert_write_access(existing, user_api_key_dict) + await _assert_write_access(prisma_client, existing, user_api_key_dict) row = await prisma_client.db.litellm_memorytable.update( where={"memory_id": existing.memory_id}, data=data, @@ -443,7 +497,9 @@ async def upsert_memory( detail=f"Memory with key '{key}' already exists.", ) # Same write-authorization check as the non-race path. - _assert_write_access(existing_after_race, user_api_key_dict) + await _assert_write_access( + prisma_client, existing_after_race, user_api_key_dict + ) row = await prisma_client.db.litellm_memorytable.update( where={"memory_id": existing_after_race.memory_id}, data=data, @@ -472,7 +528,7 @@ async def delete_memory( prisma_client = _require_prisma() row = await _find_memory_for_caller(prisma_client, key, user_api_key_dict) # Visibility != write authority — see the upsert handler for the rationale. - _assert_write_access(row, user_api_key_dict) + await _assert_write_access(prisma_client, row, user_api_key_dict) try: await prisma_client.db.litellm_memorytable.delete( where={"memory_id": row.memory_id} diff --git a/tests/test_litellm/proxy/memory/test_memory_endpoints.py b/tests/test_litellm/proxy/memory/test_memory_endpoints.py index 022caddd5aa..4a38837a23c 100644 --- a/tests/test_litellm/proxy/memory/test_memory_endpoints.py +++ b/tests/test_litellm/proxy/memory/test_memory_endpoints.py @@ -139,10 +139,38 @@ class _InMemoryMemoryTable: raise Exception("Not found") +class _InMemoryTeamTable: + """ + Tiny fake of `prisma_client.db.litellm_teamtable` — only `find_unique` + is exercised by the memory router (for team-admin checks). + """ + + def __init__(self): + self.teams: List[Any] = [] + + async def find_unique(self, where: Dict[str, Any]) -> Optional[Any]: + team_id = where["team_id"] + for t in self.teams: + if getattr(t, "team_id", None) == team_id: + return t + return None + + +def _make_team(team_id: str, *, admin_user_ids: List[str]) -> MagicMock: + """Build a team-row stub with `members_with_roles` shaped like Prisma.""" + members = [MagicMock(user_id=uid, role="admin") for uid in admin_user_ids] + team = MagicMock() + team.team_id = team_id + team.organization_id = None # skip org-admin path in tests + team.members_with_roles = members + return team + + def _make_prisma() -> MagicMock: client = MagicMock() client.db = MagicMock() client.db.litellm_memorytable = _InMemoryMemoryTable() + client.db.litellm_teamtable = _InMemoryTeamTable() return client @@ -521,10 +549,10 @@ class TestMemoryEndpoints: assert resp.status_code == 403, resp.text assert len(table.rows) == 1 - def test_put_memory_teammate_can_modify_pure_team_row(self): + def test_put_memory_team_admin_can_modify_pure_team_row(self): """ - A "pure" team row (no user_id stamped, only team_id) is intended to be - shared — any team member can modify it. + Pure team row (no user_id stamped) — only team admins (or org admins) + may modify it, matching the auth pattern in team_endpoints.py. """ table = self.prisma.db.litellm_memorytable table.rows.append( @@ -536,12 +564,61 @@ class TestMemoryEndpoints: team_id="team-shared", ) ) - client = _make_client(_user_auth("user-b", "team-shared")) + # user-admin is registered as a team admin of team-shared. + self.prisma.db.litellm_teamtable.teams.append( + _make_team("team-shared", admin_user_ids=["user-admin"]) + ) + client = _make_client(_user_auth("user-admin", "team-shared")) with _patch_prisma(self.prisma): resp = client.put("/v1/memory/team_playbook", json={"value": "v2"}) assert resp.status_code == 200, resp.text assert table.rows[0].value == "v2" + def test_put_memory_team_member_cannot_modify_pure_team_row(self): + """ + Plain team members can READ team rows (visibility OR-filter), but they + cannot WRITE — only team admins can. + """ + table = self.prisma.db.litellm_memorytable + table.rows.append( + _make_row( + memory_id="m1", + key="team_playbook", + value="v1", + user_id=None, + team_id="team-shared", + ) + ) + # team-shared exists, but user-b is NOT in members_with_roles as admin. + self.prisma.db.litellm_teamtable.teams.append( + _make_team("team-shared", admin_user_ids=["someone-else"]) + ) + client = _make_client(_user_auth("user-b", "team-shared")) + with _patch_prisma(self.prisma): + resp = client.put("/v1/memory/team_playbook", json={"value": "v2"}) + assert resp.status_code == 403, resp.text + assert table.rows[0].value == "v1" + + def test_delete_memory_team_member_cannot_delete_pure_team_row(self): + """Same as above, for DELETE.""" + table = self.prisma.db.litellm_memorytable + table.rows.append( + _make_row( + memory_id="m1", + key="team_playbook", + user_id=None, + team_id="team-shared", + ) + ) + self.prisma.db.litellm_teamtable.teams.append( + _make_team("team-shared", admin_user_ids=["user-admin"]) + ) + client = _make_client(_user_auth("user-b", "team-shared")) + with _patch_prisma(self.prisma): + resp = client.delete("/v1/memory/team_playbook") + assert resp.status_code == 403 + assert len(table.rows) == 1 + def test_admin_can_modify_any_row(self): """Admin bypasses write-authorization.""" table = self.prisma.db.litellm_memorytable