import json import os import re import threading from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor from hashlib import sha256 from pathlib import Path from typing import Final from uuid import uuid4 import httpx import psutil import psycopg import pytest import yaml from integration._support.client import ( JSON_OBJECT, Gateway, delete_key_if_present, eventually, gateway_from_environment, object_value, string_value, ) from integration._support.database import read_rows, write_rows from integration._support.process import CONFIGURATION, OwnedProxy, group_members, owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue from litellm.models.user import LiteLLM_UserTable from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken _OWNED_PROXY_SALT_KEY: Final = "sk-integration-salt" _OWNERSHIP_POOL_LIMIT: Final = 4 _CONNECTION_BUDGET: Final = re.compile( r"Postgres connection budget per pod = \d+ worker\(s\) x \(?writer connection_limit (\d+)" ) _OWNERSHIP_MARKER: Final = re.compile(rb"ownership-[0-9a-f]{32}") def _ownership_config(directory: Path) -> Path: source: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" settings: Final = CONFIGURATION.validate_python(yaml.safe_load(source.read_text())) general: Final = CONFIGURATION.validate_python(settings["general_settings"]) configured: Final = { **settings, "general_settings": {**general, "database_connection_pool_limit": _OWNERSHIP_POOL_LIMIT}, } destination: Final = directory / "ownership-proxy-config.yaml" destination.write_text(yaml.safe_dump(configured, sort_keys=False)) return destination def _project_rows(project_id: str) -> list[dict[str, JsonValue]]: return read_rows( 'SELECT p.project_id, p.project_alias, p.description, p.team_id, p.models, p.blocked, ' 'p.budget_id, b.max_budget FROM "LiteLLM_ProjectTable" AS p ' 'LEFT JOIN "LiteLLM_BudgetTable" AS b ON b.budget_id = p.budget_id ' 'WHERE p.project_id = %s', (project_id,), ) def _key_rows(key: str) -> list[dict[str, JsonValue]]: return read_rows( 'SELECT token, key_alias, team_id, project_id FROM "LiteLLM_VerificationToken" WHERE token = %s', (sha256(key.encode()).hexdigest(),), ) def _key_permission_and_budget_ids(key: str) -> list[dict[str, JsonValue]]: return read_rows( 'SELECT object_permission_id, budget_id FROM "LiteLLM_VerificationToken" WHERE token = %s', (sha256(key.encode()).hexdigest(),), ) def _key_state_rows(key: str) -> list[dict[str, JsonValue]]: return read_rows( 'SELECT to_jsonb(k) AS key_row, to_jsonb(b) AS budget_row FROM "LiteLLM_VerificationToken" AS k ' 'LEFT JOIN "LiteLLM_BudgetTable" AS b ON b.budget_id = k.budget_id WHERE k.token = %s', (sha256(key.encode()).hexdigest(),), ) def _project_state_rows(project_id: str) -> list[dict[str, JsonValue]]: return read_rows( 'SELECT to_jsonb(p) AS project_row, to_jsonb(b) AS budget_row FROM "LiteLLM_ProjectTable" AS p ' 'LEFT JOIN "LiteLLM_BudgetTable" AS b ON b.budget_id = p.budget_id WHERE p.project_id = %s', (project_id,), ) def _object_permission_rows(permission_id: str) -> list[dict[str, JsonValue]]: return read_rows( 'SELECT * FROM "LiteLLM_ObjectPermissionTable" WHERE object_permission_id = %s', (permission_id,), ) def _object_permission_table_rows() -> list[dict[str, JsonValue]]: return read_rows( 'SELECT to_jsonb(p) AS row FROM "LiteLLM_ObjectPermissionTable" AS p ORDER BY object_permission_id', (), ) def _budget_table_rows() -> list[dict[str, JsonValue]]: return read_rows( 'SELECT to_jsonb(b) AS row FROM "LiteLLM_BudgetTable" AS b ORDER BY budget_id', (), ) def _budget_rows(budget_id: str) -> list[dict[str, JsonValue]]: return read_rows( 'SELECT to_jsonb(b) AS row FROM "LiteLLM_BudgetTable" AS b WHERE budget_id = %s', (budget_id,), ) def _clear_key_object_permission(key: str, permission_id: str) -> None: write_rows( 'UPDATE "LiteLLM_VerificationToken" SET object_permission_id = NULL WHERE token = %s', (sha256(key.encode()).hexdigest(),), ) write_rows( 'DELETE FROM "LiteLLM_ObjectPermissionTable" WHERE object_permission_id = %s', (permission_id,), ) def _deleted_key_rows(key: str) -> list[dict[str, JsonValue]]: return read_rows( 'SELECT to_jsonb(d) AS row FROM "LiteLLM_DeletedVerificationToken" AS d WHERE d.token = %s', (sha256(key.encode()).hexdigest(),), ) def _deprecated_key_rows(key: str) -> list[dict[str, JsonValue]]: return read_rows( 'SELECT to_jsonb(d) AS row FROM "LiteLLM_DeprecatedVerificationToken" AS d WHERE d.token = %s', (sha256(key.encode()).hexdigest(),), ) def _cli_session_token( user_id: str, team_id: str | None, *, monkeypatch: pytest.MonkeyPatch, max_budget: float | None = None, ) -> str: monkeypatch.setenv("LITELLM_SALT_KEY", _OWNED_PROXY_SALT_KEY) user: Final = LiteLLM_UserTable( user_id=user_id, user_role="internal_user", teams=[team_id] if team_id is not None else [], models=[], max_budget=max_budget, ) return ExperimentalUIJWTToken.get_cli_jwt_auth_token( user_info=user, team_id=team_id, team_alias="ownership-team" if team_id is not None else None, max_budget=max_budget, ) def _discard_unexpected_key(candidate: Gateway, response: httpx.Response) -> None: if response.status_code != 200: return body: Final = JSON_OBJECT.validate_json(response.content) candidate.post("/key/delete", {"keys": [string_value(body["key"])]}) def _ownership_chat_reply(identity: str, stream: bool) -> Reply: if not stream: return Reply( body=json.dumps( { "id": identity, "object": "chat.completion", "created": 1, "model": "gpt-4.1-mini", "choices": [ { "index": 0, "message": {"role": "assistant", "content": "ownership ok"}, "finish_reason": "stop", } ], "usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9}, } ).encode() ) return Reply( content_type="text/event-stream", chunks=( f"data: {json.dumps({'id': identity, 'object': 'chat.completion.chunk', 'created': 1, 'model': 'gpt-4.1-mini', 'choices': [{'index': 0, 'delta': {'role': 'assistant', 'content': 'ownership'}}]})}\n\n".encode(), f"data: {json.dumps({'id': identity, 'object': 'chat.completion.chunk', 'created': 1, 'model': 'gpt-4.1-mini', 'choices': [{'index': 0, 'delta': {'content': ' ok'}, 'finish_reason': 'stop'}], 'usage': {'prompt_tokens': 7, 'completion_tokens': 2, 'total_tokens': 9}})}\n\n".encode(), b"data: [DONE]\n\n", ), ) def _ownership_responses_reply(identity: str, stream: bool) -> Reply: response: Final = { "id": identity, "object": "response", "created_at": 1, "status": "completed", "model": "gpt-4.1-mini", "output": [ { "id": f"msg_{identity}", "type": "message", "role": "assistant", "status": "completed", "content": [{"type": "output_text", "text": "ownership ok", "annotations": []}], } ], "usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9}, } if not stream: return Reply(body=json.dumps(response).encode()) events: Final = ( { "type": "response.created", "sequence_number": 0, "response": {**response, "status": "in_progress", "output": []}, }, { "type": "response.output_text.delta", "sequence_number": 1, "item_id": f"msg_{identity}", "output_index": 0, "content_index": 0, "delta": "ownership ok", }, {"type": "response.completed", "sequence_number": 2, "response": response}, ) return Reply( content_type="text/event-stream", chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events), ) def _ownership_upstream(request: Request) -> Reply: found: Final = _OWNERSHIP_MARKER.search(request.body) if found is None: return Reply(status=400, body=b'{"error":"missing ownership marker"}') marker: Final = found.group(0).decode() body: Final = JSON_OBJECT.validate_json(request.body) stream: Final = body.get("stream") is True if request.target.endswith("/responses"): return _ownership_responses_reply(f"resp_{marker}", stream) return _ownership_chat_reply(f"chatcmpl-{marker}", stream) def _ownership_sse_events(text: str) -> tuple[dict[str, JsonValue], ...]: return tuple( JSON_OBJECT.validate_json(line[6:]) for line in text.splitlines() if line.startswith("data: ") and line != "data: [DONE]" ) def _ownership_request_marker(request: Request) -> str: found: Final = _OWNERSHIP_MARKER.search(request.body) assert found is not None, request.body return found.group(0).decode() def _ownership_request_payload( path: str, model: str, marker: str, stream: bool ) -> tuple[dict[str, JsonValue], dict[str, str]]: if path.endswith("/messages"): return ( { "model": model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}], "stream": stream, }, {"anthropic-version": "2023-06-01"}, ) if path.endswith("/responses"): return {"model": model, "input": marker, "stream": stream}, {} return {"model": model, "messages": [{"role": "user", "content": marker}], "stream": stream}, {} def _ownership_response_id(response: httpx.Response, path: str) -> str: if not response.headers.get("content-type", "").startswith("text/event-stream"): return string_value(JSON_OBJECT.validate_json(response.content)["id"]) events: Final = _ownership_sse_events(response.text) identities: Final = ( tuple( string_value(object_value(event["message"])["id"]) for event in events if event.get("type") == "message_start" ) if path.endswith("/messages") else tuple( string_value(object_value(event["response"])["id"]) for event in events if event.get("type") == "response.completed" ) if path.endswith("/responses") else tuple(string_value(event["id"]) for event in events if "id" in event) ) unique_identities: Final = frozenset(identities) assert len(unique_identities) == 1, response.text return next(iter(unique_identities)) def _ownership_serving_call( candidate: Gateway, key: str, model: str, index: int, marker: str, ) -> tuple[int, str, str]: stream: Final = index % 2 == 0 route: Final = index % 3 paths: Final = ("/v1/chat/completions", "/v1/responses", "/v1/messages") path: Final = paths[route] body, headers = _ownership_request_payload(path, model, marker, stream) response: Final = candidate.request("POST", path, body, key=key, headers=headers) response.read() if response.status_code != 200: return response.status_code, "", response.text return response.status_code, _ownership_response_id(response, path), response.text def _create_mcp_server(candidate: Gateway, server_id: str, server_name: str, alias: str) -> None: response: Final = candidate.request( "POST", "/v1/mcp/server", { "server_id": server_id, "server_name": server_name, "alias": alias, "transport": "sse", "url": "http://127.0.0.1:9/mcp", }, ) assert response.status_code == 201, response.text def _delete_mcp_server(candidate: Gateway, server_id: str) -> None: response: Final = candidate.request("DELETE", f"/v1/mcp/server/{server_id}") assert response.status_code == 202, response.text def test_service_account_generate_rejects_foreign_team_project_without_writing_key(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project_b: Final = scenario.project(team_b, models=[model]) alias: Final = f"service-account-{uuid4().hex}" vector_store: Final = f"vector-store-{uuid4().hex}" budget_rows_before: Final = _budget_table_rows() object_permission_rows_before: Final = _object_permission_table_rows() response: Final = ownership_gateway.request( "POST", "/key/service-account/generate", { "team_id": team_a, "project_id": project_b, "key_alias": alias, "models": [model], "soft_budget": 3.5, "object_permission": {"vector_stores": [vector_store]}, }, ) _discard_unexpected_key(ownership_gateway, response) assert response.status_code == 400, response.text assert _budget_table_rows() == budget_rows_before assert _object_permission_table_rows() == object_permission_rows_before assert ( read_rows( 'SELECT token FROM "LiteLLM_VerificationToken" WHERE project_id = %s AND key_alias = %s', (project_b, alias), ) == [] ) def test_key_generate_unowned_project_accepts_any_team(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project: Final = scenario.project(team_a, models=[model]) write_rows('UPDATE "LiteLLM_ProjectTable" SET team_id = NULL WHERE project_id = %s', (project,)) team_key: Final = scenario.key(team_id=team_b, project_id=project, models=[model]) teamless_key: Final = scenario.key(project_id=project, models=[model]) chats: Final = tuple( ownership_gateway.request( "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "legacy unowned"}]}, key=key, ) for key in (team_key, teamless_key) ) assert all(chat.status_code == 200 for chat in chats), tuple(chat.text for chat in chats) @pytest.mark.parametrize( ("path", "stream"), ( ("/v1/chat/completions", False), ("/v1/chat/completions", True), ("/v1/messages", False), ("/v1/messages", True), ("/v1/responses", False), ("/v1/responses", True), ), ) def test_legacy_mismatched_key_keeps_serving_and_stays_editable( ownership_gateway: Gateway, path: str, stream: bool ) -> None: def respond(request: Request) -> Reply: return _ownership_upstream(request) with wire_server(respond) as upstream: with ownership_gateway.scenario() as scenario: model: Final = scenario.model(api_base=f"{upstream.url}/v1") team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project: Final = scenario.project(team_a, models=[model]) key: Final = scenario.key(team_id=team_a, project_id=project, models=[model]) write_rows( 'UPDATE "LiteLLM_VerificationToken" SET team_id = %s WHERE token = %s', (team_b, sha256(key.encode()).hexdigest()), ) marker: Final = f"ownership-{uuid4().hex}" body, headers = _ownership_request_payload(path, model, marker, stream) response: Final = ownership_gateway.request("POST", path, body, key=key, headers=headers) response.read() assert response.status_code == 200, response.text requests: Final = upstream.drain() assert len(requests) == 1 expected_target: Final = "/chat/completions" if path.endswith("/chat/completions") else "/responses" assert requests[0].target.endswith(expected_target), requests[0].target assert marker.encode() in requests[0].body assert _ownership_response_id(response, path) != "" def test_stored_team_mismatch_allows_key_edits_and_detach(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project: Final = scenario.project(team_a, models=[model]) key: Final = scenario.key(team_id=team_a, project_id=project, models=[model]) write_rows( 'UPDATE "LiteLLM_VerificationToken" SET team_id = %s WHERE token = %s', (team_b, sha256(key.encode()).hexdigest()), ) alias: Final = f"stored-mismatch-{uuid4().hex}" alias_update: Final = ownership_gateway.request( "POST", "/key/update", {"key": key, "key_alias": alias}, ) assert alias_update.status_code == 200, alias_update.text aliased_key: Final = _key_rows(key) assert aliased_key[0]["key_alias"] == alias unchanged: Final = ownership_gateway.request( "POST", "/key/update", {"key": key, "team_id": team_b}, ) assert unchanged.status_code == 200, unchanged.text assert _key_rows(key) == aliased_key detached: Final = ownership_gateway.request( "POST", "/key/update", {"key": key, "project_id": None}, ) assert detached.status_code == 200, detached.text detached_key: Final = _key_rows(key) assert detached_key[0]["team_id"] == team_b assert detached_key[0]["project_id"] is None def test_key_regenerate_cross_team_with_object_permission_writes_nothing(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project_a: Final = scenario.project(team_a, models=[model]) project_b: Final = scenario.project(team_b, models=[model]) key: Final = scenario.key( team_id=team_a, project_id=project_a, models=[model], object_permission={"vector_stores": ["before"]}, ) ids: Final = _key_permission_and_budget_ids(key) assert len(ids) == 1 permission_id: Final = string_value(ids[0]["object_permission_id"]) scenario.cleanups.callback(_clear_key_object_permission, key, permission_id) before: Final = ( _object_permission_rows(permission_id), _object_permission_table_rows(), _key_rows(key), _key_state_rows(key), _key_permission_and_budget_ids(key), _deleted_key_rows(key), _deprecated_key_rows(key), ) response: Final = ownership_gateway.request( "POST", f"/key/{key}/regenerate", {"project_id": project_b, "object_permission": {"vector_stores": ["after"]}}, ) _discard_unexpected_key(ownership_gateway, response) assert response.status_code == 400, response.text assert ( _object_permission_rows(permission_id), _object_permission_table_rows(), _key_rows(key), _key_state_rows(key), _key_permission_and_budget_ids(key), _deleted_key_rows(key), _deprecated_key_rows(key), ) == before def test_key_bulk_update_cross_team_with_object_permission_preserves_permission(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project: Final = scenario.project(team_a, models=[model]) key: Final = scenario.key( team_id=team_a, project_id=project, models=[model], object_permission={"vector_stores": ["before"]}, ) ids: Final = _key_permission_and_budget_ids(key) assert len(ids) == 1 permission_id: Final = string_value(ids[0]["object_permission_id"]) scenario.cleanups.callback(_clear_key_object_permission, key, permission_id) before: Final = ( _object_permission_rows(permission_id), _object_permission_table_rows(), _key_rows(key), _key_state_rows(key), _key_permission_and_budget_ids(key), _deleted_key_rows(key), _deprecated_key_rows(key), ) response: Final = ownership_gateway.request( "POST", "/key/bulk_update", { "keys": [ { "key": key, "team_id": team_b, "object_permission": {"vector_stores": ["after"]}, } ], }, ) assert response.status_code == 200, response.text failed_updates: Final = JSON_OBJECT.validate_json(response.content)["failed_updates"] assert isinstance(failed_updates, list) assert len(failed_updates) == 1 failed_update: Final = object_value(failed_updates[0]) assert f"Project {project} belongs to team {team_a}" in string_value(failed_update["failed_reason"]) assert ( _object_permission_rows(permission_id), _object_permission_table_rows(), _key_rows(key), _key_state_rows(key), _key_permission_and_budget_ids(key), _deleted_key_rows(key), _deprecated_key_rows(key), ) == before def test_key_update_ambiguous_mcp_permission_error_precedes_project_ownership(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() identifier: Final = f"ambiguous{uuid4().hex}" first_id: Final = f"mcp{uuid4().hex}" second_id: Final = f"mcp{uuid4().hex}" _create_mcp_server(ownership_gateway, first_id, identifier, f"alias{uuid4().hex}") scenario.cleanups.callback(_delete_mcp_server, ownership_gateway, first_id) _create_mcp_server(ownership_gateway, second_id, f"name{uuid4().hex}", f"alias{uuid4().hex}") scenario.cleanups.callback(_delete_mcp_server, ownership_gateway, second_id) write_rows( 'UPDATE "LiteLLM_MCPServerTable" SET alias = %s WHERE server_id = %s', (identifier, second_id), ) team_permissions: Final = {"mcp_servers": [first_id, second_id]} team_a: Final = scenario.team(models=[model], object_permission=team_permissions) team_b: Final = scenario.team(models=[model], object_permission=team_permissions) project: Final = scenario.project(team_a, models=[model]) key: Final = scenario.key( team_id=team_a, project_id=project, models=[model], object_permission={"vector_stores": ["before"]}, ) ids: Final = _key_permission_and_budget_ids(key) assert len(ids) == 1 permission_id: Final = string_value(ids[0]["object_permission_id"]) scenario.cleanups.callback(_clear_key_object_permission, key, permission_id) before: Final = ( _object_permission_rows(permission_id), _object_permission_table_rows(), _key_rows(key), _key_state_rows(key), _key_permission_and_budget_ids(key), _deleted_key_rows(key), _deprecated_key_rows(key), ) response: Final = ownership_gateway.request( "POST", "/key/update", { "key": key, "team_id": team_b, "object_permission": {"mcp_tool_permissions": {identifier: ["tool"]}}, }, ) assert response.status_code == 400, response.text assert "ambiguous" in response.text.lower(), response.text assert "project" not in response.text.lower(), response.text assert ( _object_permission_rows(permission_id), _object_permission_table_rows(), _key_rows(key), _key_state_rows(key), _key_permission_and_budget_ids(key), _deleted_key_rows(key), _deprecated_key_rows(key), ) == before def test_team_key_bulk_update_rejects_foreign_team_project(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project_b: Final = scenario.project(team_b, models=[model]) key: Final = scenario.key(team_id=team_a, models=[model]) before: Final = _key_rows(key) response: Final = ownership_gateway.request( "POST", "/team/key/bulk_update", { "team_id": team_a, "key_ids": [sha256(key.encode()).hexdigest()], "update_fields": {"project_id": project_b, "team_id": team_b}, }, ) assert response.status_code == 422, response.text assert "project_id" in response.text, response.text assert "team_id" in response.text, response.text assert _key_rows(key) == before def test_bulk_update_and_regenerate_new_object_permission_is_served(ownership_gateway: Gateway) -> None: with wire_server(_ownership_upstream) as upstream: with ownership_gateway.scenario() as scenario: model: Final = scenario.model(api_base=f"{upstream.url}/v1") team: Final = scenario.team(models=[model]) project: Final = scenario.project(team, models=[model]) bulk_key: Final = scenario.key(team_id=team, project_id=project, models=[model]) bulk_response: Final = ownership_gateway.request( "POST", "/key/bulk_update", { "keys": [ { "key": bulk_key, "object_permission": {"vector_stores": ["bulk"]}, } ], }, ) assert bulk_response.status_code == 200, bulk_response.text bulk_info: Final = ownership_gateway.request("GET", "/key/info", params={"key": bulk_key}) assert bulk_info.status_code == 200, bulk_info.text bulk_row: Final = object_value(JSON_OBJECT.validate_json(bulk_info.content)["info"]) bulk_permission_id: Final = string_value(bulk_row["object_permission_id"]) assert bulk_permission_id != "" assert len(_object_permission_rows(bulk_permission_id)) == 1 scenario.cleanups.callback(_clear_key_object_permission, bulk_key, bulk_permission_id) bulk_chat: Final = ownership_gateway.chat(model, key=bulk_key, text=f"ownership-{uuid4().hex}") assert string_value(bulk_chat["id"]) != "" regenerated: Final = ownership_gateway.request( "POST", "/key/generate", {"team_id": team, "project_id": project, "models": [model]}, ) assert regenerated.status_code == 200, regenerated.text regenerated_key: Final = string_value(JSON_OBJECT.validate_json(regenerated.content)["key"]) scenario.cleanups.callback(delete_key_if_present, ownership_gateway, regenerated_key) regeneration: Final = ownership_gateway.request( "POST", f"/key/{regenerated_key}/regenerate", {"object_permission": {"vector_stores": ["regenerated"]}}, ) assert regeneration.status_code == 200, regeneration.text new_key: Final = string_value(JSON_OBJECT.validate_json(regeneration.content)["key"]) scenario.cleanups.callback(delete_key_if_present, ownership_gateway, new_key) info: Final = ownership_gateway.request("GET", "/key/info", params={"key": new_key}) assert info.status_code == 200, info.text info_row: Final = object_value(JSON_OBJECT.validate_json(info.content)["info"]) regenerated_permission_id: Final = string_value(info_row["object_permission_id"]) assert regenerated_permission_id != "" assert regenerated_permission_id != bulk_permission_id assert len(_object_permission_rows(regenerated_permission_id)) == 1 scenario.cleanups.callback(_clear_key_object_permission, new_key, regenerated_permission_id) chat: Final = ownership_gateway.chat(model, key=new_key, text=f"ownership-{uuid4().hex}") assert string_value(chat["id"]) != "" def test_ownership_rejections_during_concurrent_traffic_burst(ownership_gateway: Gateway) -> None: with wire_server(_ownership_upstream) as upstream: with ownership_gateway.scenario() as scenario: model: Final = scenario.model(api_base=f"{upstream.url}/v1") team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project_a: Final = scenario.project(team_a, models=[model]) serving_project: Final = scenario.project(team_a, models=[model]) key: Final = scenario.key(team_id=team_a, project_id=project_a, models=[model]) serving_key: Final = scenario.key(team_id=team_a, project_id=serving_project, models=[model]) markers: Final = tuple(f"ownership-{uuid4().hex}" for _ in range(30)) key_before: Final = ( _object_permission_table_rows(), _key_rows(key), _key_state_rows(key), _key_permission_and_budget_ids(key), _deleted_key_rows(key), _deprecated_key_rows(key), ) serving_key_before: Final = ( _key_rows(serving_key), _key_permission_and_budget_ids(serving_key), _deleted_key_rows(serving_key), _deprecated_key_rows(serving_key), ) project_before: Final = _project_state_rows(project_a) def serve(index: int) -> tuple[int, str, str]: return _ownership_serving_call(ownership_gateway, serving_key, model, index, markers[index]) def update_key() -> httpx.Response: return ownership_gateway.request("POST", "/key/update", {"key": key, "team_id": team_b}) def bulk_update() -> httpx.Response: return ownership_gateway.request( "POST", "/key/bulk_update", {"keys": [{"key": key, "team_id": team_b}]}, ) def move_project() -> httpx.Response: return ownership_gateway.request( "POST", "/project/update", {"project_id": project_a, "team_id": team_b} ) with ThreadPoolExecutor(max_workers=33) as pool: serving_futures: Final = tuple(pool.submit(serve, index) for index in range(30)) ownership_futures: Final = ( pool.submit(update_key), pool.submit(bulk_update), pool.submit(move_project), ) serving_results: Final = tuple(future.result(timeout=90) for future in serving_futures) ownership_results: Final = tuple(future.result(timeout=90) for future in ownership_futures) assert all(status == 200 for status, _, _ in serving_results), serving_results assert all(response_id for _, response_id, _ in serving_results), serving_results assert len({response_id for _, response_id, _ in serving_results}) == 30 key_update_response: Final = ownership_results[0] bulk_update_response: Final = ownership_results[1] project_update_response: Final = ownership_results[2] assert key_update_response.status_code == 400, key_update_response.text assert bulk_update_response.status_code == 200, bulk_update_response.text bulk_result: Final = JSON_OBJECT.validate_json(bulk_update_response.content) successful_updates: Final = bulk_result["successful_updates"] failed_updates: Final = bulk_result["failed_updates"] assert isinstance(successful_updates, list), bulk_update_response.text assert successful_updates == [], bulk_update_response.text assert isinstance(failed_updates, list), bulk_update_response.text assert len(failed_updates) == 1, bulk_update_response.text failed_update: Final = object_value(failed_updates[0]) assert f"Project {project_a} belongs to team {team_a}" in string_value(failed_update["failed_reason"]), ( bulk_update_response.text ) assert project_update_response.status_code == 400, project_update_response.text upstream_requests: Final = upstream.drain() assert len(upstream_requests) == 30 received_markers: Final = tuple(_ownership_request_marker(request) for request in upstream_requests) assert tuple(sorted(received_markers)) == tuple(sorted(markers)) assert ( _object_permission_table_rows(), _key_rows(key), _key_state_rows(key), _key_permission_and_budget_ids(key), _deleted_key_rows(key), _deprecated_key_rows(key), ) == key_before assert ( _key_rows(serving_key), _key_permission_and_budget_ids(serving_key), _deleted_key_rows(serving_key), _deprecated_key_rows(serving_key), ) == serving_key_before assert _project_state_rows(project_a) == project_before def test_ownership_checks_wait_for_row_locks_without_failing_readiness( ownership_gateway: Gateway, ownership_proxy: OwnedProxy ) -> None: budget: Final = _CONNECTION_BUDGET.search(ownership_proxy.log.read_text()) assert budget is not None, ownership_proxy.log.read_text() assert int(budget.group(1)) == _OWNERSHIP_POOL_LIMIT with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team: Final = scenario.team(models=[model]) project: Final = scenario.project(team, models=[model], max_budget=7) key: Final = scenario.key(team_id=team, project_id=project, models=[model]) alias: Final = f"locked-{uuid4().hex}" locked: Final = threading.Event() release: Final = threading.Event() def hold_locks() -> None: with psycopg.connect(os.environ["DATABASE_URL"]) as connection: connection.execute("BEGIN") connection.execute( 'SELECT project_id FROM "LiteLLM_ProjectTable" WHERE project_id = %s FOR UPDATE', (project,), ) connection.execute( 'SELECT token FROM "LiteLLM_VerificationToken" WHERE token = %s FOR UPDATE', (sha256(key.encode()).hexdigest(),), ) locked.set() assert release.wait(timeout=30) connection.commit() with ThreadPoolExecutor(max_workers=3) as pool: lock_future: Final = pool.submit(hold_locks) try: assert locked.wait(timeout=10) key_future: Final = pool.submit( ownership_gateway.request, "POST", "/key/update", {"key": key, "key_alias": alias}, ) project_future: Final = pool.submit( ownership_gateway.request, "POST", "/project/update", {"project_id": project, "team_id": team, "max_budget": 19}, ) lock_waiters: Final = eventually( lambda: read_rows( "SELECT pid FROM pg_stat_activity WHERE datname = current_database() " "AND wait_event_type = 'Lock' AND cardinality(pg_blocking_pids(pid)) > 0", (), ), lambda rows: len(rows) >= 2, seconds=10, ) assert len(lock_waiters) >= 2 readiness: Final = eventually( lambda: ownership_gateway.request("GET", "/health/readiness").status_code, lambda status: status == 200, seconds=10, ) assert readiness == 200 finally: release.set() key_response: Final = key_future.result(timeout=90) project_response: Final = project_future.result(timeout=90) lock_future.result(timeout=30) assert key_response.status_code == 200, key_response.text assert project_response.status_code == 200, project_response.text assert _key_rows(key)[0]["key_alias"] == alias assert _project_rows(project)[0]["max_budget"] == 19.0 def test_ownership_enforced_after_worker_kill(ownership_gateway: Gateway, tmp_path: Path) -> None: started: Final = threading.Event() release_stream: Final = threading.Event() def respond(request: Request) -> Reply: started.set() reply: Final = _ownership_upstream(request) return Reply( status=reply.status, body=reply.body, content_type=reply.content_type, chunks=reply.chunks, gate_after_first=release_stream, ) with wire_server(respond) as upstream: with owned_proxy_process( ownership_gateway, tmp_path / "worker-kill", {"LITELLM_SALT_KEY": _OWNED_PROXY_SALT_KEY}, workers=2, ) as owned: with owned.gateway.scenario() as scenario: model: Final = scenario.model(api_base=f"{upstream.url}/v1") team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project_a: Final = scenario.project(team_a, models=[model]) project_b: Final = scenario.project(team_b, models=[model]) key: Final = scenario.key(team_id=team_a, project_id=project_a, models=[model]) markers: Final = tuple(f"ownership-{uuid4().hex}" for _ in range(24)) def serve(index: int) -> tuple[int, int, str]: try: response: Final = owned.gateway.request( "POST", "/v1/chat/completions", { "model": model, "messages": [{"role": "user", "content": markers[index]}], "stream": True, }, key=key, ) except httpx.TransportError as error: return index, 0, str(error) return index, response.status_code, response.text candidate_port: Final = owned.gateway.client.base_url.port assert candidate_port is not None workers: Final = tuple( process for process in group_members(owned.process.pid) if process.pid != owned.process.pid and any( connection.laddr.port == candidate_port and connection.status == psutil.CONN_LISTEN for connection in process.net_connections(kind="inet") ) ) assert len(workers) == 2, tuple(process.pid for process in workers) victim: Final = workers[0] survivor: Final = workers[1] with owned.replacing_worker(victim.pid) as kill, ThreadPoolExecutor(max_workers=24) as pool: futures: Final = tuple(pool.submit(serve, index) for index in range(24)) try: assert started.wait(timeout=10) kill() release_stream.set() finally: release_stream.set() burst: Final = tuple(future.result(timeout=90) for future in futures) assert psutil.pid_exists(survivor.pid), survivor.pid assert len(burst) == 24 burst_requests: Final = upstream.drain() empty_body_count: Final = sum(not request.body for request in burst_requests) burst_markers: Final = tuple( _ownership_request_marker(request) for request in burst_requests if request.body ) successful_markers: Final = tuple(markers[index] for index, status, _ in burst if status == 200) assert burst_markers, f"empty_body_captures={empty_body_count}; burst={burst}" assert all(burst_markers.count(marker) == 1 for marker in successful_markers), ( f"empty_body_captures={empty_body_count}; " f"successful_markers={successful_markers}; upstream_markers={burst_markers}; burst={burst}" ) cross_team: Final = owned.gateway.request( "POST", "/key/generate", {"team_id": team_a, "project_id": project_b, "models": [model]}, ) _discard_unexpected_key(owned.gateway, cross_team) assert cross_team.status_code == 400, cross_team.text marker: Final = f"ownership-{uuid4().hex}" chat: Final = owned.gateway.request( "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": marker}]}, key=key, ) assert chat.status_code == 200, chat.text post_kill_requests: Final = upstream.drain() post_kill_empty_body_count: Final = sum(not request.body for request in post_kill_requests) post_kill_markers: Final = tuple( _ownership_request_marker(request) for request in post_kill_requests if request.body ) assert post_kill_markers.count(marker) == 1, ( f"empty_body_captures={post_kill_empty_body_count}; " f"post_kill_markers={post_kill_markers}; response={chat.text}" ) @pytest.fixture(scope="module") def ownership_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: directory: Final = tmp_path_factory.mktemp("project-team-ownership") with gateway_from_environment() as gateway: with owned_proxy_process( gateway, directory, {"LITELLM_SALT_KEY": _OWNED_PROXY_SALT_KEY}, config=_ownership_config(directory), workers=2, ) as owned: yield owned @pytest.fixture(scope="module") def ownership_gateway(ownership_proxy: OwnedProxy) -> Gateway: return ownership_proxy.gateway @pytest.mark.covers("mgmt.project.new.real_route_persists") def test_project_new_persists_real_state(gateway: Gateway) -> None: with gateway.scenario() as scenario: model: Final = scenario.model() team: Final = scenario.team(models=[model]) budget: Final = scenario.budget(max_budget=7) project: Final = scenario.project( team, project_alias="new-project", budget_id=budget, models=[model], description="new project" ) key: Final = scenario.key(team_id=team, project_id=project, models=[model]) assert object_value(gateway.chat(model, key=key)["usage"])["total_tokens"] == 40 rows: Final = _project_rows(project) assert rows != [] assert len(rows) == 1 row: Final = rows[0] assert row["project_id"] == project assert row["project_alias"] == "new-project" assert row["team_id"] == team assert row["description"] == "new project" assert row["models"] == [model] assert row["budget_id"] == budget assert row["blocked"] is False assert row["max_budget"] == 7.0 @pytest.mark.covers("mgmt.project.update.real_route_persists") def test_project_update_persists_real_state(gateway: Gateway) -> None: with gateway.scenario() as scenario: model: Final = scenario.model() team: Final = scenario.team(models=[model]) budget: Final = scenario.budget(max_budget=3) project: Final = scenario.project(team, budget_id=budget, models=[model], description="before") key: Final = scenario.key(team_id=team, project_id=project, models=[model]) updated: Final = gateway.post( "/project/update", { "project_id": project, "project_alias": "updated-project", "description": "after", "max_budget": 9, "blocked": True, }, ) assert string_value(updated["project_id"]) == project rows: Final = _project_rows(project) assert rows != [] assert len(rows) == 1 row: Final = rows[0] assert row["project_alias"] == "updated-project" assert row["description"] == "after" assert row["team_id"] == team assert row["models"] == [model] assert row["budget_id"] == budget assert row["blocked"] is True assert row["max_budget"] == 9.0 blocked: Final = gateway.request( "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "blocked project"}]}, key=key, ) assert blocked.status_code == 401, blocked.text assert object_value(blocked.json()["error"])["type"] == "auth_error" gateway.post("/project/update", {"project_id": project, "blocked": False}) assert object_value(gateway.chat(model, key=key)["usage"])["total_tokens"] == 40 @pytest.mark.covers("mgmt.project.delete.attached_key_refusal_preserves_state") def test_project_delete_with_attached_key_refuses_and_preserves_state(gateway: Gateway) -> None: with gateway.scenario() as scenario: model: Final = scenario.model() team: Final = scenario.team(models=[model]) budget: Final = scenario.budget() project: Final = scenario.project( team, budget_id=budget, project_alias="delete-project", models=[model] ) key: Final = scenario.key(team_id=team, project_id=project, models=[model]) digest: Final = sha256(key.encode()).hexdigest() project_before: Final = _project_rows(project) key_before: Final = read_rows( 'SELECT token, key_alias, models, metadata, max_budget, team_id, project_id, budget_id ' 'FROM "LiteLLM_VerificationToken" WHERE token = %s', (digest,), ) assert len(project_before) == 1 assert len(key_before) == 1 assert key_before[0]["project_id"] == project assert key_before[0]["team_id"] == team denied: Final = gateway.request("DELETE", "/project/delete", {"project_ids": [project]}) assert denied.status_code == 400, denied.text assert _project_rows(project) == project_before assert read_rows( 'SELECT token, key_alias, models, metadata, max_budget, team_id, project_id, budget_id ' 'FROM "LiteLLM_VerificationToken" WHERE token = %s', (digest,), ) == key_before def test_key_generate_rejects_foreign_team_project(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project_b: Final = scenario.project(team_b, models=[model]) alias: Final = f"cross-team-{uuid4().hex}" vector_store: Final = f"vector-store-{uuid4().hex}" budget_rows_before: Final = _budget_table_rows() object_permission_rows_before: Final = _object_permission_table_rows() cross_team: Final = ownership_gateway.request( "POST", "/key/generate", { "team_id": team_a, "project_id": project_b, "key_alias": alias, "models": [model], "soft_budget": 3.5, "object_permission": {"vector_stores": [vector_store]}, }, ) _discard_unexpected_key(ownership_gateway, cross_team) assert cross_team.status_code == 400, cross_team.text assert _budget_table_rows() == budget_rows_before assert _object_permission_table_rows() == object_permission_rows_before assert ( read_rows( 'SELECT token FROM "LiteLLM_VerificationToken" WHERE project_id = %s AND team_id = %s', (project_b, team_a), ) == [] ) def test_key_generate_rejects_missing_team_for_owned_project(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_b: Final = scenario.team(models=[model]) project_b: Final = scenario.project(team_b, models=[model]) unbound: Final = ownership_gateway.request( "POST", "/key/generate", {"project_id": project_b, "models": [model]} ) _discard_unexpected_key(ownership_gateway, unbound) assert unbound.status_code == 400, unbound.text assert read_rows( 'SELECT token FROM "LiteLLM_VerificationToken" WHERE project_id = %s AND team_id IS NULL', (project_b,), ) == [] def test_key_generate_same_team_project_key_can_chat(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team: Final = scenario.team(models=[model]) project: Final = scenario.project(team, models=[model]) generated: Final = ownership_gateway.request( "POST", "/key/generate", {"team_id": team, "project_id": project, "models": [model]} ) assert generated.status_code == 200, generated.text key: Final = string_value(JSON_OBJECT.validate_json(generated.content)["key"]) scenario.cleanups.callback(scenario.delete_key, key) generated_rows: Final = _key_rows(key) assert len(generated_rows) == 1 assert generated_rows[0]["team_id"] == team assert generated_rows[0]["project_id"] == project chat: Final = ownership_gateway.request( "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "control"}]}, key=key, ) assert chat.status_code == 200, chat.text assert object_value(JSON_OBJECT.validate_json(chat.content)["usage"])["total_tokens"] == 40 def test_key_update_rejects_team_change_and_allows_unchanged_values(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project_a: Final = scenario.project(team_a, models=[model]) key: Final = scenario.key(team_id=team_a, project_id=project_a, models=[model], soft_budget=3.0) permission_seed: Final = ownership_gateway.request( "POST", "/key/update", {"key": key, "object_permission": {"vector_stores": ["existing-store"]}}, ) assert permission_seed.status_code == 200, permission_seed.text key_permission_state: Final = _key_permission_and_budget_ids(key) assert len(key_permission_state) == 1 permission_id: Final = string_value(key_permission_state[0]["object_permission_id"]) budget_id: Final = string_value(key_permission_state[0]["budget_id"]) scenario.cleanups.callback(_clear_key_object_permission, key, permission_id) permission_before: Final = _object_permission_rows(permission_id) budget_before: Final = _budget_rows(budget_id) assert len(permission_before) == 1 assert len(budget_before) == 1 aliased: Final = ownership_gateway.request("POST", "/key/update", {"key": key, "key_alias": "updated"}) assert aliased.status_code == 200, aliased.text after_alias: Final = _key_rows(key) assert len(after_alias) == 1 assert after_alias[0]["key_alias"] == "updated" assert after_alias[0]["team_id"] == team_a assert after_alias[0]["project_id"] == project_a unchanged_team: Final = ownership_gateway.request("POST", "/key/update", {"key": key, "team_id": team_a}) assert unchanged_team.status_code == 200, unchanged_team.text after_unchanged_team: Final = _key_rows(key) assert after_unchanged_team == after_alias before_reassignment: Final = _key_rows(key) assert len(before_reassignment) == 1 reassigned: Final = ownership_gateway.request( "POST", "/key/update", { "key": key, "team_id": team_b, "soft_budget": 5.0, "object_permission": {"vector_stores": ["replacement-store"]}, }, ) assert reassigned.status_code == 400, reassigned.text assert _key_rows(key) == before_reassignment assert _object_permission_rows(permission_id) == permission_before assert _budget_rows(budget_id) == budget_before def test_key_update_can_detach_project_and_change_team(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project_a: Final = scenario.project(team_a, models=[model]) key: Final = scenario.key(team_id=team_a, project_id=project_a, models=[model]) detached: Final = ownership_gateway.request( "POST", "/key/update", {"key": key, "project_id": None, "team_id": team_b} ) assert detached.status_code == 200, detached.text after_detach: Final = _key_rows(key) assert len(after_detach) == 1 assert after_detach[0]["project_id"] is None assert after_detach[0]["team_id"] == team_b assert after_detach[0]["key_alias"] is None def test_key_regenerate_rejects_foreign_project_without_changing_key(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project_b: Final = scenario.project(team_b, models=[model]) generated: Final = ownership_gateway.request( "POST", "/key/generate", {"team_id": team_a, "models": [model]} ) assert generated.status_code == 200, generated.text key: Final = string_value(JSON_OBJECT.validate_json(generated.content)["key"]) scenario.cleanups.callback(scenario.delete_key, key) before: Final = _key_rows(key) before_deleted: Final = _deleted_key_rows(key) before_deprecated: Final = _deprecated_key_rows(key) assert len(before) == 1 assert before_deleted == [] assert before_deprecated == [] response: Final = ownership_gateway.request( "POST", f"/key/{key}/regenerate", {"project_id": project_b, "grace_period": "1h"} ) _discard_unexpected_key(ownership_gateway, response) assert response.status_code == 400, response.text assert _key_rows(key) == before assert _deleted_key_rows(key) == before_deleted assert _deprecated_key_rows(key) == before_deprecated def _regenerate_as(candidate: Gateway, route: str, target: str, caller: str, body: dict[str, JsonValue]) -> httpx.Response: if route == "body": return candidate.request("POST", "/key/regenerate", {"key": target, **body}, key=caller) return candidate.request("POST", f"/key/{target}/regenerate", body, key=caller) @pytest.mark.parametrize("route", ["path", "body"]) def test_key_regenerate_rejects_outsider_moving_own_key_into_foreign_team_project( ownership_gateway: Gateway, route: str ) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_b: Final = scenario.team(models=[model]) project_b: Final = scenario.project(team_b, models=[model]) outsider: Final = scenario.user(user_role="internal_user") caller: Final = scenario.key(user_id=outsider) generated: Final = ownership_gateway.request("POST", "/key/generate", {"user_id": outsider, "models": [model]}) assert generated.status_code == 200, generated.text target: Final = string_value(JSON_OBJECT.validate_json(generated.content)["key"]) scenario.cleanups.callback(delete_key_if_present, ownership_gateway, target) before: Final = _key_rows(target) assert len(before) == 1 assert before[0]["team_id"] is None assert before[0]["project_id"] is None response: Final = _regenerate_as( ownership_gateway, route, target, caller, {"team_id": team_b, "project_id": project_b} ) _discard_unexpected_key(ownership_gateway, response) assert response.status_code == 403, response.text assert _key_rows(target) == before assert _deleted_key_rows(target) == [] assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE project_id = %s', (project_b,)) == [] chat: Final = ownership_gateway.request( "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "rejected regeneration keeps the old key"}]}, key=target, ) assert chat.status_code == 200, chat.text @pytest.mark.parametrize("route", ["path", "body"]) def test_key_regenerate_rejects_team_admin_moving_team_key_into_foreign_team_project( ownership_gateway: Gateway, route: str ) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project_b: Final = scenario.project(team_b, models=[model]) team_a_admin: Final = scenario.member(team_a, "admin") caller: Final = scenario.key(user_id=team_a_admin) generated: Final = ownership_gateway.request("POST", "/key/generate", {"team_id": team_a, "models": [model]}) assert generated.status_code == 200, generated.text target: Final = string_value(JSON_OBJECT.validate_json(generated.content)["key"]) scenario.cleanups.callback(delete_key_if_present, ownership_gateway, target) before: Final = _key_rows(target) assert len(before) == 1 assert before[0]["team_id"] == team_a response: Final = _regenerate_as( ownership_gateway, route, target, caller, {"team_id": team_b, "project_id": project_b} ) _discard_unexpected_key(ownership_gateway, response) assert response.status_code == 403, response.text assert _key_rows(target) == before assert _deleted_key_rows(target) == [] assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE project_id = %s', (project_b,)) == [] chat: Final = ownership_gateway.request( "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "rejected team key regeneration keeps the key"}]}, key=target, ) assert chat.status_code == 200, chat.text @pytest.mark.parametrize(("route", "role"), [("path", "user"), ("body", "admin")]) def test_key_regenerate_lets_project_team_member_move_own_key_into_project( ownership_gateway: Gateway, route: str, role: str ) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_b: Final = scenario.team(models=[model]) project_b: Final = scenario.project(team_b, models=[model]) member: Final = scenario.member(team_b, role) caller: Final = scenario.key(user_id=member) generated: Final = ownership_gateway.request("POST", "/key/generate", {"user_id": member, "models": [model]}) assert generated.status_code == 200, generated.text target: Final = string_value(JSON_OBJECT.validate_json(generated.content)["key"]) scenario.cleanups.callback(delete_key_if_present, ownership_gateway, target) response: Final = _regenerate_as( ownership_gateway, route, target, caller, {"team_id": team_b, "project_id": project_b} ) assert response.status_code == 200, response.text new_key: Final = string_value(JSON_OBJECT.validate_json(response.content)["key"]) scenario.cleanups.callback(delete_key_if_present, ownership_gateway, new_key) rows: Final = _key_rows(new_key) assert len(rows) == 1 assert rows[0]["team_id"] == team_b assert rows[0]["project_id"] == project_b assert _key_rows(target) == [] chat: Final = ownership_gateway.request( "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "member regeneration serves"}]}, key=new_key, ) assert chat.status_code == 200, chat.text def test_key_generation_rejects_missing_project_without_writing_key(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team: Final = scenario.team(models=[model]) missing_project_id: Final = f"missing-{uuid4()}" key_generation: Final = ownership_gateway.request( "POST", "/key/generate", {"team_id": team, "project_id": missing_project_id, "models": [model]}, ) assert key_generation.status_code == 404, key_generation.text assert ( read_rows( 'SELECT token FROM "LiteLLM_VerificationToken" WHERE project_id = %s', (missing_project_id,), ) == [] ) service_account_generation: Final = ownership_gateway.request( "POST", "/key/service-account/generate", {"team_id": team, "project_id": missing_project_id, "models": [model]}, ) assert service_account_generation.status_code == 404, service_account_generation.text assert ( read_rows( 'SELECT token FROM "LiteLLM_VerificationToken" WHERE project_id = %s', (missing_project_id,), ) == [] ) def test_key_regenerate_routes_reject_missing_project_without_changing_key(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team: Final = scenario.team(models=[model]) key: Final = scenario.key(team_id=team, models=[model]) missing_project_id: Final = f"missing-{uuid4()}" before: Final = _key_rows(key) responses: Final = ( ownership_gateway.request( "POST", "/key/regenerate", {"key": key, "project_id": missing_project_id}, ), ownership_gateway.request( "POST", f"/key/{key}/regenerate", {"project_id": missing_project_id}, ), ) for response in responses: assert response.status_code == 404, response.text assert _key_rows(key) == before assert _deleted_key_rows(key) == [] assert _deprecated_key_rows(key) == [] assert ( read_rows( 'SELECT token FROM "LiteLLM_VerificationToken" WHERE project_id = %s', (missing_project_id,), ) == [] ) chat: Final = ownership_gateway.request( "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "rejected regeneration preserves key"}]}, key=key, ) assert chat.status_code == 200, chat.text def test_key_generate_nonmember_organization_error_precedes_project_ownership( ownership_gateway: Gateway, monkeypatch: pytest.MonkeyPatch ) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() organization_id: Final = scenario.organization() owner_team: Final = scenario.team(models=[model]) project: Final = scenario.project(owner_team, models=[model]) caller_id: Final = scenario.user(user_role="internal_user") caller_token: Final = _cli_session_token(caller_id, None, monkeypatch=monkeypatch) response: Final = ownership_gateway.request( "POST", "/key/generate", { "project_id": project, "organization_id": organization_id, "models": [model], }, key=caller_token, ) assert response.status_code == 403, response.text assert f"Caller is not a member of organization_id={organization_id}" in response.text assert ( read_rows( 'SELECT token FROM "LiteLLM_VerificationToken" WHERE project_id = %s', (project,), ) == [] ) def test_key_generate_duplicate_alias_error_precedes_project_ownership( ownership_gateway: Gateway, monkeypatch: pytest.MonkeyPatch ) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() caller_team: Final = scenario.team(models=[model]) owner_team: Final = scenario.team(models=[model]) project: Final = scenario.project(owner_team, models=[model]) key_alias: Final = f"duplicate-{uuid4()}" existing_key: Final = scenario.key(team_id=caller_team, key_alias=key_alias, models=[model]) caller_id: Final = scenario.member(caller_team, role="admin") caller_token: Final = _cli_session_token(caller_id, caller_team, monkeypatch=monkeypatch) existing_alias_rows: Final = read_rows( 'SELECT token FROM "LiteLLM_VerificationToken" WHERE key_alias = %s', (key_alias,), ) response: Final = ownership_gateway.request( "POST", "/key/generate", { "team_id": caller_team, "project_id": project, "key_alias": key_alias, "models": [model], }, key=caller_token, ) assert response.status_code == 400, response.text assert f"Key with alias '{key_alias}' already exists" in response.text assert len(existing_alias_rows) == 1 assert read_rows( 'SELECT token FROM "LiteLLM_VerificationToken" WHERE key_alias = %s', (key_alias,), ) == existing_alias_rows assert len(_key_rows(existing_key)) == 1 assert ( read_rows( 'SELECT token FROM "LiteLLM_VerificationToken" WHERE project_id = %s', (project,), ) == [] ) def test_key_generate_budget_ceiling_precedes_project_ownership( ownership_gateway: Gateway, monkeypatch: pytest.MonkeyPatch ) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() caller_team: Final = scenario.team(models=[model]) owner_team: Final = scenario.team(models=[model]) project: Final = scenario.project(owner_team, models=[model]) caller_id: Final = scenario.member(caller_team, role="admin") caller_token: Final = _cli_session_token( caller_id, caller_team, monkeypatch=monkeypatch, max_budget=1 ) response: Final = ownership_gateway.request( "POST", "/key/generate", { "team_id": caller_team, "project_id": project, "max_budget": 5, "models": [model], }, key=caller_token, ) assert response.status_code == 400, response.text assert "max_budget (5.0) cannot exceed the caller's own max_budget (1.0)" in response.text assert ( read_rows( 'SELECT token FROM "LiteLLM_VerificationToken" WHERE project_id = %s', (project,), ) == [] ) def test_key_regenerate_output_estimate_admin_error_precedes_project_ownership( ownership_gateway: Gateway, monkeypatch: pytest.MonkeyPatch ) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() project_team: Final = scenario.team(models=[model]) destination_team: Final = scenario.team(models=[model]) project: Final = scenario.project(project_team, models=[model]) caller_id: Final = scenario.member(project_team, role="admin") key: Final = scenario.key(team_id=project_team, user_id=caller_id, models=[model]) caller_token: Final = _cli_session_token(caller_id, project_team, monkeypatch=monkeypatch) before: Final = _key_rows(key) response: Final = ownership_gateway.request( "POST", f"/key/{key}/regenerate", { "team_id": destination_team, "project_id": project, "metadata": {"default_estimated_output_tokens": 1}, }, key=caller_token, ) assert response.status_code == 403, response.text assert "Only proxy admins can set" in response.text assert _key_rows(key) == before chat: Final = ownership_gateway.request( "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "rejected regeneration keeps key valid"}]}, key=key, ) assert chat.status_code == 200, chat.text def test_key_bulk_update_rejects_foreign_team_project_and_preserves_key(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project_a: Final = scenario.project(team_a, models=[model]) key: Final = scenario.key(team_id=team_a, project_id=project_a, models=[model]) before: Final = _key_rows(key) assert len(before) == 1 assert before[0]["team_id"] == team_a assert before[0]["project_id"] == project_a response: Final = ownership_gateway.request( "POST", "/key/bulk_update", { "keys": [ {"key": key, "team_id": team_b}, {"key": key, "max_budget": 10, "tags": ["bulk-update"]}, ] }, ) assert response.status_code == 200, response.text body: Final = JSON_OBJECT.validate_json(response.content) failed_updates: Final = body["failed_updates"] successful_updates: Final = body["successful_updates"] assert isinstance(failed_updates, list), response.text assert isinstance(successful_updates, list), response.text assert len(failed_updates) == 1 assert len(successful_updates) == 1 failed_update: Final = object_value(failed_updates[0]) successful_update: Final = object_value(successful_updates[0]) assert string_value(failed_update["key"]) == key assert f"Project {project_a} belongs to team {team_a}" in string_value(failed_update["failed_reason"]) assert string_value(successful_update["key"]) == key after: Final = _key_rows(key) assert len(after) == 1 assert after[0]["team_id"] == before[0]["team_id"] assert after[0]["project_id"] == before[0]["project_id"] def test_project_update_rejects_moving_project_with_attached_key(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) budget: Final = scenario.budget(max_budget=3) attached_project: Final = scenario.project(team_a, budget_id=budget, models=[model]) key: Final = scenario.key(team_id=team_a, project_id=attached_project, models=[model]) project_before: Final = _project_rows(attached_project) key_before: Final = _key_rows(key) budget_before: Final = _budget_rows(budget) moved_with_key: Final = ownership_gateway.request( "POST", "/project/update", { "project_id": attached_project, "team_id": team_b, "max_budget": 11, }, ) assert moved_with_key.status_code == 400, moved_with_key.text assert _project_rows(attached_project) == project_before assert _key_rows(key) == key_before assert _budget_rows(budget) == budget_before def test_project_update_allows_moving_project_without_keys(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) unattached_project: Final = scenario.project(team_a, models=[model]) moved_without_key: Final = ownership_gateway.request( "POST", "/project/update", {"project_id": unattached_project, "team_id": team_b} ) assert moved_without_key.status_code == 200, moved_without_key.text moved_project: Final = _project_rows(unattached_project) assert len(moved_project) == 1 assert moved_project[0]["team_id"] == team_b assert read_rows( 'SELECT token FROM "LiteLLM_VerificationToken" WHERE project_id = %s', (unattached_project,), ) == [] def test_project_update_rejects_moving_project_with_teamless_key(ownership_gateway: Gateway) -> None: with ownership_gateway.scenario() as scenario: model: Final = scenario.model() team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project: Final = scenario.project(team_a, models=[model]) key: Final = scenario.key(team_id=team_a, project_id=project, models=[model]) write_rows( 'UPDATE "LiteLLM_VerificationToken" SET team_id = NULL WHERE token = %s', (sha256(key.encode()).hexdigest(),), ) project_before: Final = _project_rows(project) key_before: Final = _key_rows(key) assert key_before[0]["team_id"] is None moved: Final = ownership_gateway.request("POST", "/project/update", {"project_id": project, "team_id": team_b}) assert moved.status_code == 400, moved.text assert _project_rows(project) == project_before assert _key_rows(key) == key_before