Refactor key_manager (#369)