I'm not sure there is a way to reboot a device without losing the connection. By design, wouldn't the device shut down all existing connections when it reboots?
The only thing I can think of right now is to write some custom code to poll for the device using adb and re-instantiate the driver once adb reports the device is available.
Good luck.