Improve O1 agent UT coverage