diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ac6a2a4a7db..d8e1faf9dc6 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2179,6 +2179,7 @@ async def _process_single_key_update( detail={"error": "Database not connected"}, ) + await _ensure_user_row_for_key_write(prisma_client, non_default_values.get("user_id"), existing_key_row.user_id) _data = {**non_default_values, "token": update_key_request.key} response = await prisma_client.update_data(token=update_key_request.key, data=_data) @@ -2648,6 +2649,7 @@ async def update_key_fn( _data = {**non_default_values, "token": key} if prisma_client is None: raise Exception("Not connected to DB!") + await _ensure_user_row_for_key_write(prisma_client, non_default_values.get("user_id"), existing_key_row.user_id) response = await prisma_client.update_data(token=key, data=_data) # Delete - key from cache, since it's been updated! @@ -3531,6 +3533,24 @@ def _check_model_access_group( return True +async def _ensure_user_row_for_key_write( + prisma_client: PrismaClient, user_id: object, existing_user_id: str | None = None +) -> None: + """Key rows carry a foreign key to LiteLLM_UserTable, so every write path + that sets user_id without provisioning the user (table_name="key" minting + such as JWT auto-register, admin user_id rebinds on update/regenerate) + must create the referenced user row first. Create-only upsert: an existing + user is never modified, and a user_id unchanged from existing_user_id is + already satisfied by the constraint so no lookup is made. + """ + if not user_id or not isinstance(user_id, str) or user_id == existing_user_id: + return + await prisma_client.db.litellm_usertable.upsert( + where={"user_id": user_id}, + data={"create": {"user_id": user_id}, "update": {}}, + ) + + async def generate_key_helper_fn( request_type: Literal["user", "key"], # identifies if this request is from /user/new or /key/generate duration: Optional[str] = None, @@ -3793,6 +3813,8 @@ async def generate_key_helper_fn( "prisma_client: Creating Key= %s", {**key_data, "token": hash_token(token=token)}, ) + if table_name == "key": + await _ensure_user_row_for_key_write(prisma_client, key_data.get("user_id")) create_key_response = await prisma_client.insert_data(data=key_data, table_name="key") key_data["token_id"] = getattr(create_key_response, "token", None) @@ -4520,6 +4542,7 @@ async def _execute_virtual_key_regeneration( grace_period=data.grace_period if data else None, ) + await _ensure_user_row_for_key_write(prisma_client, update_data.get("user_id"), key_in_db.user_id) updated_token = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_api_key}, data=update_data, # type: ignore diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 51f72f91dc3..66b7fdddaa0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -780,6 +780,7 @@ async def test_generate_key_helper_fn_with_access_group_ids(monkeypatch): mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock( return_value=MagicMock(object_permission_id=None) ) + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock() captured_key_data = {} @@ -830,6 +831,7 @@ async def test_generate_key_helper_fn_with_budget_fallbacks(monkeypatch): mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock( return_value=MagicMock(object_permission_id=None) ) + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock() captured_key_data = {} @@ -6255,6 +6257,62 @@ def test_build_key_filter_conditions_agent_id_narrows_visibility(): assert "agent_id" not in json.dumps(where_without) +@pytest.mark.asyncio +async def test_ensure_user_row_for_key_write_create_only_upsert(): + """ + The FK on LiteLLM_VerificationToken.user_id means key writes must + provision the referenced user row: create it when missing, never touch + an existing one, and no-op for empty or non-string user ids. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _ensure_user_row_for_key_write, + ) + + mock_prisma_client = AsyncMock() + mock_upsert = AsyncMock() + mock_prisma_client.db.litellm_usertable.upsert = mock_upsert + + await _ensure_user_row_for_key_write(mock_prisma_client, "ghost-user") + mock_upsert.assert_called_once_with( + where={"user_id": "ghost-user"}, + data={"create": {"user_id": "ghost-user"}, "update": {}}, + ) + + mock_upsert.reset_mock() + await _ensure_user_row_for_key_write(mock_prisma_client, None) + await _ensure_user_row_for_key_write(mock_prisma_client, "") + await _ensure_user_row_for_key_write(mock_prisma_client, 42) + mock_upsert.assert_not_called() + + +@pytest.mark.asyncio +async def test_generate_key_helper_fn_table_name_key_provisions_user_row(): + """ + Regression for the FK on user_id: table_name="key" minting (admin + /key/generate with an arbitrary user_id, JWT auto-register) skips the + user-creation branch, so the helper must provision the user row before + inserting the key or the insert violates the constraint. + """ + mock_prisma_client = AsyncMock() + mock_prisma_client.insert_data = AsyncMock( + return_value=MagicMock(token="hashed-token", litellm_budget_table=None, created_at=None, updated_at=None) + ) + mock_upsert = AsyncMock() + mock_prisma_client.db.litellm_usertable.upsert = mock_upsert + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + await generate_key_helper_fn( + request_type="key", + user_id="ghost-user", + table_name="key", + ) + + assert mock_upsert.call_args.kwargs["where"] == {"user_id": "ghost-user"} + insert_tables = [c.kwargs.get("table_name") for c in mock_prisma_client.insert_data.call_args_list] + assert "user" not in insert_tables, "table_name='key' must not run the full user-creation branch" + assert "key" in insert_tables + + @pytest.mark.asyncio async def test_generate_key_negative_max_budget(): """ @@ -8274,6 +8332,7 @@ async def test_default_key_generate_params_object_permission_not_rejected_for_no mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock( return_value=MagicMock(object_permission_id="objperm-4") ) + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)