Updated callback logic