Update to work with mercurial ui refactorings