SSH exception after reboot

[rsie-dev]

Describe the bug

After a successful reboot the next operation fails with:
Command socket/SSH error: SSHException('SSH session not active',)

To Reproduce

In my case:

reboot = server.reboot( )
apt.packages(
    packages=["ifupdown"],
    present=False,
    _if=reboot.did_change
)

Target system information:

• vagrant: dietpi
• raspberry pi 2: dietpi
  both have the same problem

Expected behavior

After successful reboot, follow operations succeed, too.

Meta

System: Linux
Platform: Linux-6.8.0-41-generic-x86_64-with-glibc2.39
Release: 6.8.0-41-generic
Machine: x86_64
pyinfra: v3.1
click: v8.1.7
configparser: v7.1.0
distro: v1.9.0
gevent: v24.2.1
jinja2: v3.1.4
packaging: v24.1
paramiko: v3.4.1
python-dateutil: v2.9.0.post0
pywinrm: v0.5.0
setuptools: v74.0.0
typeguard: v4.3.0
typing-extensions: v4.12.2
Executable: venv/bin/pyinfra
Python: 3.12.3 (CPython, GCC 13.2.0)

pyinfra installed via: pip
Include debug.log(if one was created)

[uplight-dev]

Also having this issue and it is quite blocking (unless you know some workarounds).
Thanks!

[uplight-dev]

Took a look at the stack trace and pyinfra sources.

I see once the reboot happens the /tmp pyinfra files are lost.
Also the state is lost (Error: could not load fact: deb.DebPackages) and all hosts get invalidated("No hosts remaining").

So I'm not sure if the pyinfra client keeps the state locally or it relies on the remote host's /tmp files to reload it but the reboot causes it to be lost. As the next operation following the reboot cannot access the hosts anymore, it thinks the facts are empty.
In operations/server.py, the reboot op causes the host.connected = True but the facts containing the host list never get updated.

Found this workaround:
So to ensure the pyinfra client's state is fully reloaded after the remote host reboot, I split everything following the reboot in a second py file. And invoke it in a separate CLI command which re-inits the state and hosts correctly.

# this one does a remote host update and reboots it and waits to reconnect
pyinfra -vvv -y /data/inventory.py /data/update1_and_reboot.py

# this second pyinfra CLI command ensures the client is fully reconnected to target host
pyinfra -vvv -y /data/inventory.py /data/update2.py

Since this looks like a workaround to me, is there any way to refresh the state while still in the initial CLI call?
Thanks!