Fix edit functionality