Minor updates due to internal API changes