1 files changed, 2 insertions, 2 deletions

diff --git a/tools/idevicediagnostics.c b/tools/idevicediagnostics.c
index 464df20..216ebf1 100644
--- a/tools/idevicediagnostics.c
+++ b/tools/idevicediagnostics.c
@@ -203,7 +203,7 @@ int main(int argc, char **argv)
                                         }
                                 break;
                                 case CMD_RESTART:
-                                        if (diagnostics_relay_restart(diagnostics_client, 0) == DIAGNOSTICS_RELAY_E_SUCCESS) {
+                                        if (diagnostics_relay_restart(diagnostics_client, DIAGNOSTICS_RELAY_ACTION_FLAG_WAIT_FOR_DISCONNECT) == DIAGNOSTICS_RELAY_E_SUCCESS) {
                                                 printf("Restarting device.\n");
                                                 result = EXIT_SUCCESS;
                                         } else {
@@ -211,7 +211,7 @@ int main(int argc, char **argv)
                                         }
                                 break;
                                 case CMD_SHUTDOWN:
-                                        if (diagnostics_relay_shutdown(diagnostics_client, 0) == DIAGNOSTICS_RELAY_E_SUCCESS) {
+                                        if (diagnostics_relay_shutdown(diagnostics_client, DIAGNOSTICS_RELAY_ACTION_FLAG_WAIT_FOR_DISCONNECT) == DIAGNOSTICS_RELAY_E_SUCCESS) {
                                                 printf("Shutting down device.\n");
                                                 result = EXIT_SUCCESS;
                                         } else {