Entity handling fixes