Add missing includes to fix build break