litellm/tests/integration/authorization/test_memory_ownership.py

import uuid
from typing import Final

import httpx
from integration._support.client import Gateway, Scenario
from integration._support.database import read_rows, write_rows
from pydantic import JsonValue, TypeAdapter

from litellm.types.memory_management import LiteLLM_MemoryRow, MemoryDeleteResponse, MemoryListResponse

_MEMORY_ROUTE: Final = ["/v1/memory"]
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_CREATE_METADATA: Final[dict[str, JsonValue]] = {"source": "create", "labels": ["alpha", "beta"], "depth": 2}
_PUT_METADATA: Final[dict[str, JsonValue]] = {"source": "put", "nested": {"enabled": True, "weights": [0.5, 1]}}
_UPDATED_METADATA: Final[dict[str, JsonValue]] = {"source": "update", "labels": []}


def _admin_only_error(route: str, user_id: str) -> dict[str, JsonValue]:
    masked_user_id: Final = f"{user_id[:6]}{'*' * (len(user_id) - 8)}{user_id[-2:]}"
    return {
        "error": {
            "message": (
                "Authentication Error, Only proxy admin can be used to generate, delete, update info for new "
                f"keys/users/teams. Route={route}. Your role=internal_user. Your user_id={masked_user_id}"
            ),
            "type": "auth_error",
            "param": "None",
            "code": "401",
        }
    }


def _delete_memory_row(key: str) -> None:
    write_rows('DELETE FROM "LiteLLM_MemoryTable" WHERE key = %s', (key,))


def _memory_rows(key: str) -> list[dict[str, JsonValue]]:
    return read_rows(
        'SELECT key, value, user_id, team_id, created_by, updated_by FROM "LiteLLM_MemoryTable" WHERE key = %s',
        (key,),
    )


def _memory_create_rows(key: str) -> list[dict[str, JsonValue]]:
    return read_rows(
        "SELECT memory_id, key, value, metadata, user_id, team_id, created_by, updated_by "
        'FROM "LiteLLM_MemoryTable" WHERE key = %s',
        (key,),
    )


def _memory_metadata(key: str) -> list[dict[str, JsonValue]]:
    return read_rows('SELECT metadata FROM "LiteLLM_MemoryTable" WHERE key = %s', (key,))


def _prefix_listing(gateway: Gateway, prefix: str, caller: str) -> tuple[int, list[str]]:
    response: Final = gateway.request("GET", "/v1/memory", key=caller, params={"key_prefix": prefix})
    assert response.status_code == 200, response.text
    listing: Final = MemoryListResponse.model_validate_json(response.content)
    return listing.total, sorted(row.key for row in listing.memories)


def _memory_request(
    gateway: Gateway,
    method: str,
    key: str,
    *,
    caller: str,
    body: dict[str, JsonValue] | None = None,
) -> httpx.Response:
    return gateway.request(method, f"/v1/memory/{key}", body, key=caller)


def _register_memory_cleanup(scenario: Scenario, key: str) -> None:
    scenario.cleanups.callback(_delete_memory_row, key)


def test_memory_create_enforces_personal_and_team_scope(gateway: Gateway) -> None:
    with gateway.scenario() as scenario:
        team: Final = scenario.team()
        other_team: Final = scenario.team()
        member_a: Final = scenario.member(team)
        member_b: Final = scenario.member(team)
        key_a: Final = scenario.key(user_id=member_a, team_id=team, allowed_routes=_MEMORY_ROUTE)
        own_key: Final = f"memory-create-own-{uuid.uuid4().hex}"
        other_user_key: Final = f"memory-create-other-user-{uuid.uuid4().hex}"
        other_team_key: Final = f"memory-create-other-team-{uuid.uuid4().hex}"
        admin_key: Final = f"memory-create-admin-{uuid.uuid4().hex}"

        _register_memory_cleanup(scenario, own_key)
        own_response: Final = gateway.request(
            "POST",
            "/v1/memory",
            {"key": own_key, "value": "own", "metadata": _CREATE_METADATA},
            key=key_a,
        )
        assert own_response.status_code == 200, own_response.text
        own: Final = LiteLLM_MemoryRow.model_validate_json(own_response.content)
        assert (
            own.key,
            own.value,
            own.metadata,
            own.user_id,
            own.team_id,
            own.created_by,
            own.updated_by,
        ) == (own_key, "own", _CREATE_METADATA, member_a, team, member_a, member_a), own_response.text
        assert own.created_at is not None and own.updated_at is not None, own_response.text
        assert _memory_create_rows(own_key) == [
            {
                "memory_id": own.memory_id,
                "key": own_key,
                "value": "own",
                "metadata": _CREATE_METADATA,
                "user_id": member_a,
                "team_id": team,
                "created_by": member_a,
                "updated_by": member_a,
            }
        ], own_response.text

        _register_memory_cleanup(scenario, other_user_key)
        other_user_response: Final = gateway.request(
            "POST",
            "/v1/memory",
            {"key": other_user_key, "value": "x", "user_id": member_b},
            key=key_a,
        )
        assert other_user_response.status_code == 403, other_user_response.text
        assert other_user_response.text == '{"detail":"Only proxy admins may set user_id to a different user."}', (
            other_user_response.text
        )
        assert _memory_rows(other_user_key) == []

        _register_memory_cleanup(scenario, other_team_key)
        other_team_response: Final = gateway.request(
            "POST",
            "/v1/memory",
            {"key": other_team_key, "value": "x", "team_id": other_team},
            key=key_a,
        )
        assert other_team_response.status_code == 403, other_team_response.text
        assert other_team_response.text == '{"detail":"Only proxy admins may set team_id to a different team."}', (
            other_team_response.text
        )
        assert _memory_rows(other_team_key) == []

        _register_memory_cleanup(scenario, admin_key)
        admin_response: Final = gateway.request(
            "POST",
            "/v1/memory",
            {"key": admin_key, "value": "x", "user_id": member_b},
            key=gateway.key,
        )
        assert admin_response.status_code == 200, admin_response.text
        admin_memory: Final = LiteLLM_MemoryRow.model_validate_json(admin_response.content)
        assert (
            admin_memory.key,
            admin_memory.value,
            admin_memory.user_id,
            admin_memory.created_by,
            admin_memory.updated_by,
        ) == (admin_key, "x", member_b, "default_user_id", "default_user_id"), admin_response.text
        assert admin_memory.created_at is not None and admin_memory.updated_at is not None, admin_response.text
        assert _memory_create_rows(admin_key) == [
            {
                "memory_id": admin_memory.memory_id,
                "key": admin_key,
                "value": "x",
                "metadata": None,
                "user_id": member_b,
                "team_id": None,
                "created_by": "default_user_id",
                "updated_by": "default_user_id",
            }
        ], admin_response.text


def test_memory_routes_enforce_personal_and_team_write_ownership(gateway: Gateway) -> None:
    with gateway.scenario() as scenario:
        team: Final = scenario.team()
        member_a: Final = scenario.member(team)
        member_b: Final = scenario.member(team)
        team_admin: Final = scenario.member(team, role="admin")
        outsider: Final = scenario.user(user_role="internal_user")
        default_key: Final = scenario.key(user_id=member_a)
        key_a: Final = scenario.key(user_id=member_a, team_id=team, allowed_routes=_MEMORY_ROUTE)
        key_b: Final = scenario.key(user_id=member_b, team_id=team, allowed_routes=_MEMORY_ROUTE)
        admin_key: Final = scenario.key(user_id=team_admin, team_id=team, allowed_routes=_MEMORY_ROUTE)
        team_key: Final = scenario.key(team_id=team, allowed_routes=_MEMORY_ROUTE)
        outsider_key: Final = scenario.key(user_id=outsider, allowed_routes=_MEMORY_ROUTE)
        personal_key: Final = f"agent/{team}-{member_a}-alpha"
        team_only_key: Final = f"agent/{team}-shared"

        default_access: Final = _memory_request(
            gateway,
            "PUT",
            personal_key,
            caller=default_key,
            body={"value": "default access"},
        )
        assert default_access.status_code == 401, default_access.text
        assert _JSON_OBJECT.validate_json(default_access.content) == _admin_only_error(
            f"/v1/memory/{personal_key}", member_a
        ), default_access.text

        created: Final = _memory_request(
            gateway,
            "PUT",
            personal_key,
            caller=key_a,
            body={"value": "a1", "metadata": _PUT_METADATA},
        )
        assert created.status_code == 200, created.text
        _register_memory_cleanup(scenario, personal_key)
        created_row: Final = LiteLLM_MemoryRow.model_validate_json(created.content)
        assert (created_row.key, created_row.value, created_row.user_id, created_row.team_id) == (
            personal_key,
            "a1",
            member_a,
            team,
        ), created.text
        assert created_row.metadata == _PUT_METADATA, created.text
        assert _memory_metadata(personal_key) == [{"metadata": _PUT_METADATA}]
        assert _memory_rows(personal_key) == [
            {
                "key": personal_key,
                "value": "a1",
                "user_id": member_a,
                "team_id": team,
                "created_by": member_a,
                "updated_by": member_a,
            }
        ]

        visible_to_member: Final = _memory_request(gateway, "GET", personal_key, caller=key_b)
        assert visible_to_member.status_code == 200, visible_to_member.text
        assert LiteLLM_MemoryRow.model_validate_json(visible_to_member.content).value == "a1", visible_to_member.text
        overwrite_refused: Final = _memory_request(
            gateway,
            "PUT",
            personal_key,
            caller=key_b,
            body={"value": "b-overwrite"},
        )
        assert overwrite_refused.status_code == 403, overwrite_refused.text
        assert overwrite_refused.text == ('{"detail":"You do not have permission to modify this memory entry."}'), (
            overwrite_refused.text
        )
        delete_refused: Final = _memory_request(gateway, "DELETE", personal_key, caller=key_b)
        assert delete_refused.status_code == 403, delete_refused.text
        assert _memory_rows(personal_key) == [
            {
                "key": personal_key,
                "value": "a1",
                "user_id": member_a,
                "team_id": team,
                "created_by": member_a,
                "updated_by": member_a,
            }
        ]

        outsider_read: Final = _memory_request(gateway, "GET", personal_key, caller=outsider_key)
        assert outsider_read.status_code == 404, outsider_read.text
        assert outsider_read.text == f"""{{"detail":"Memory with key '{personal_key}' not found"}}""", (
            outsider_read.text
        )
        outsider_list: Final = gateway.request("GET", "/v1/memory", key=outsider_key)
        assert outsider_list.status_code == 200, outsider_list.text
        outsider_memories: Final = MemoryListResponse.model_validate_json(outsider_list.content)
        assert outsider_memories.total == 0 and outsider_memories.memories == [], outsider_list.text

        overwritten: Final = _memory_request(
            gateway,
            "PUT",
            personal_key,
            caller=key_a,
            body={"value": "a2"},
        )
        assert overwritten.status_code == 200, overwritten.text
        assert LiteLLM_MemoryRow.model_validate_json(overwritten.content).value == "a2", overwritten.text
        assert _memory_metadata(personal_key) == [{"metadata": _PUT_METADATA}]
        metadata_updated: Final = _memory_request(
            gateway,
            "PUT",
            personal_key,
            caller=key_a,
            body={"metadata": _UPDATED_METADATA},
        )
        assert metadata_updated.status_code == 200, metadata_updated.text
        assert LiteLLM_MemoryRow.model_validate_json(metadata_updated.content).model_dump(
            include={"value", "metadata"}
        ) == {"value": "a2", "metadata": _UPDATED_METADATA}, metadata_updated.text
        assert _memory_metadata(personal_key) == [{"metadata": _UPDATED_METADATA}]
        assert _memory_rows(personal_key) == [
            {
                "key": personal_key,
                "value": "a2",
                "user_id": member_a,
                "team_id": team,
                "created_by": member_a,
                "updated_by": member_a,
            }
        ]

        team_created: Final = _memory_request(
            gateway,
            "PUT",
            team_only_key,
            caller=team_key,
            body={"value": "team value"},
        )
        assert team_created.status_code == 200, team_created.text
        _register_memory_cleanup(scenario, team_only_key)
        team_row: Final = LiteLLM_MemoryRow.model_validate_json(team_created.content)
        assert (team_row.value, team_row.user_id, team_row.team_id) == ("team value", None, team), team_created.text
        assert _memory_rows(team_only_key) == [
            {
                "key": team_only_key,
                "value": "team value",
                "user_id": None,
                "team_id": team,
                "created_by": None,
                "updated_by": None,
            }
        ]
        member_team_read: Final = _memory_request(gateway, "GET", team_only_key, caller=key_b)
        assert member_team_read.status_code == 200, member_team_read.text
        assert (
            LiteLLM_MemoryRow.model_validate_json(member_team_read.content).model_dump(include={"key", "value"})
        ) == {"key": team_only_key, "value": "team value"}, member_team_read.text
        member_list: Final = gateway.request("GET", "/v1/memory", key=key_b)
        assert member_list.status_code == 200, member_list.text
        member_memories: Final = MemoryListResponse.model_validate_json(member_list.content)
        assert (member_memories.total, sorted((row.key, row.value) for row in member_memories.memories)) == (
            2,
            sorted(((personal_key, "a2"), (team_only_key, "team value"))),
        ), member_list.text
        outsider_team_read: Final = _memory_request(gateway, "GET", team_only_key, caller=outsider_key)
        assert outsider_team_read.status_code == 404, outsider_team_read.text

        team_prefix: Final = f"agent/{team}-"
        outsider_prefixed_key: Final = f"{team_prefix}outsider"
        member_unprefixed_key: Final = f"notes/{team}-{member_a}"
        outsider_prefixed: Final = _memory_request(
            gateway,
            "PUT",
            outsider_prefixed_key,
            caller=outsider_key,
            body={"value": "outsider value"},
        )
        assert outsider_prefixed.status_code == 200, outsider_prefixed.text
        _register_memory_cleanup(scenario, outsider_prefixed_key)
        member_unprefixed: Final = _memory_request(
            gateway,
            "PUT",
            member_unprefixed_key,
            caller=key_a,
            body={"value": "notes value"},
        )
        assert member_unprefixed.status_code == 200, member_unprefixed.text
        _register_memory_cleanup(scenario, member_unprefixed_key)
        assert _prefix_listing(gateway, team_prefix, key_b) == (2, sorted((personal_key, team_only_key)))
        assert _prefix_listing(gateway, team_prefix, outsider_key) == (1, [outsider_prefixed_key])
        assert _prefix_listing(gateway, "notes/", outsider_key) == (0, [])

        member_team_write: Final = _memory_request(
            gateway,
            "PUT",
            team_only_key,
            caller=key_a,
            body={"value": "member overwrite"},
        )
        assert member_team_write.status_code == 403, member_team_write.text
        assert _memory_rows(team_only_key) == [
            {
                "key": team_only_key,
                "value": "team value",
                "user_id": None,
                "team_id": team,
                "created_by": None,
                "updated_by": None,
            }
        ]
        admin_write: Final = _memory_request(
            gateway,
            "PUT",
            team_only_key,
            caller=admin_key,
            body={"value": "admin overwrite"},
        )
        assert admin_write.status_code == 200, admin_write.text
        assert LiteLLM_MemoryRow.model_validate_json(admin_write.content).value == "admin overwrite", admin_write.text

        team_admin_personal_write: Final = _memory_request(
            gateway,
            "PUT",
            personal_key,
            caller=admin_key,
            body={"value": "team admin overwrite"},
        )
        assert team_admin_personal_write.status_code == 403, team_admin_personal_write.text
        assert team_admin_personal_write.text == (
            '{"detail":"You do not have permission to modify this memory entry."}'
        ), team_admin_personal_write.text
        assert _memory_rows(personal_key) == [
            {
                "key": personal_key,
                "value": "a2",
                "user_id": member_a,
                "team_id": team,
                "created_by": member_a,
                "updated_by": member_a,
            }
        ]

        master_read: Final = _memory_request(gateway, "GET", personal_key, caller=gateway.key)
        assert master_read.status_code == 200, master_read.text
        assert LiteLLM_MemoryRow.model_validate_json(master_read.content).value == "a2", master_read.text
        master_write: Final = _memory_request(
            gateway,
            "PUT",
            personal_key,
            caller=gateway.key,
            body={"value": "master overwrite"},
        )
        assert master_write.status_code == 200, master_write.text
        assert _memory_rows(personal_key) == [
            {
                "key": personal_key,
                "value": "master overwrite",
                "user_id": member_a,
                "team_id": team,
                "created_by": member_a,
                "updated_by": "default_user_id",
            }
        ]

        deleted: Final = _memory_request(gateway, "DELETE", personal_key, caller=key_a)
        assert deleted.status_code == 200, deleted.text
        assert MemoryDeleteResponse.model_validate_json(deleted.content).deleted is True, deleted.text
        assert _memory_rows(personal_key) == []
        missing: Final = _memory_request(gateway, "GET", personal_key, caller=key_a)
        assert missing.status_code == 404, missing.text
        assert missing.text == f"""{{"detail":"Memory with key '{personal_key}' not found"}}""", missing.text