[build fix] Fix type annotation