Ansible playbook rebooting the server , wait till reboot complete

I'm facing a issue while executing shell script with ansible playbook.
Issue: I'm losing a connection with remote server because shell script is rebooting the server.

My Ansible playbook

- name: Handle reboot
  hosts: all
  become: yes
  
  tasks:
    - name: Execute the script 
      shell: bash testscript.sh
      args:
        chdir: /home/ubuntu
      notify:
        - Wait for server to restart
        
  handlers:
    - name: Wait for server to restart
      local_action:
        module: wait_for
          host={{ inventory_hostname }}
          port=22
          delay=10
        become: false

My shell script :

#!/bin/bash
echo "Performing some tasks"
echo "rebooting now"
reboot
echo "reboot completed"
echo "Performing some more tasks"

The error I'm getting while the remote server reboot

fatal: [my-ip-address]: UNREACHABLE! => {
    "changed": false,
    "msg": "Failed to connect to the host via ssh: Shared connection to <my-ip-address> closed.",
    "unreachable": true
}

Is it possible to handle a reboot which is done by shell script and wait for connection until the remote server is up again ?

Thanks.

You should use wait_for_connection instead.

tasks:
  - name: Execute the script 
    shell: bash testscript.sh
    args:
      chdir: /home/ubuntu
  - name: wait
    wait_for_connection:
      delay: 10

I'd advise to do this in a task, not a handler. The handler is only executed after all tasks have been finished, so if you have tasks following the task that executes the reboot they will be tried before the playbook even starts to wait.

Alternatively, use the reboot module, which does this automatically.