Small functionality fixes