Add stage directive code