Fixup bugs in last patch