connectors.ssh: add workaround for paramiko no session error

See: paramiko/paramiko#1390

Author: dfaerch

src/pyinfra/connectors/ssh.py

@@ -9,6 +9,7 @@
 
 import click
 from paramiko import AuthenticationException, BadHostKeyException, SFTPClient, SSHException
+from paramiko.agent import Agent
 from typing_extensions import TypedDict, Unpack, override
 
 from pyinfra import logger
@@ -286,10 +287,64 @@ def _connect(self) -> None:
                 f"Host key for {e.hostname} does not match.",
             )
 
+        except SSHException as e:
+            if self._retry_paramiko_agent_keys(hostname, kwargs, e):
+                return
+            raise
+
     @override
     def disconnect(self) -> None:
         self.get_file_transfer_connection.cache.clear()
 
+    def _retry_paramiko_agent_keys(
+        self,
+        hostname: str,
+        kwargs: dict[str, Any],
+        error: SSHException,
+    ) -> bool:
+        # Workaround for Paramiko multi-key bug (paramiko/paramiko#1390).
+        if "no existing session" not in str(error).lower():
+            return False
+
+        if not kwargs.get("allow_agent"):
+            return False
+
+        try:
+            agent_keys = list(Agent().get_keys())
+        except Exception:
+            return False
+
+        if not agent_keys:
+            return False
+
+        # Skip the first agent key, since Paramiko already attempted it
+        attempt_keys = agent_keys[1:] if len(agent_keys) > 1 else agent_keys
+
+        for agent_key in attempt_keys:
+            if self.client is not None:
+                try:
+                    self.client.close()
+                except Exception:
+                    pass
+
+            self.client = SSHClient()
+
+            single_key_kwargs = dict(kwargs)
+            single_key_kwargs["allow_agent"] = False
+            single_key_kwargs["pkey"] = agent_key
+
+            try:
+                self.client.connect(hostname, **single_key_kwargs)
+                return True
+            except AuthenticationException:
+                continue
+            except SSHException as retry_error:
+                if "no existing session" in str(retry_error).lower():
+                    continue
+                raise retry_error
+
+        return False
+
     @override
     def run_shell_command(
         self,

tests/test_connectors/test_ssh.py

@@ -10,6 +10,7 @@
 from pyinfra.api.connect import connect_all
 from pyinfra.api.exceptions import ConnectError, PyinfraError
 from pyinfra.context import ctx_state
+from pyinfra.connectors import ssh
 
 from ..util import make_inventory
 
@@ -119,6 +120,94 @@ def test_connect_with_rsa_ssh_key(self):
 
         connect_all(second_state)
 
+    def test_retry_paramiko_agent_keys_single_key(self):
+        connector = ssh.SSHConnector.__new__(ssh.SSHConnector)
+        connector.client = mock.Mock()
+
+        attempts = []
+        connect_outcomes = [None]
+
+        def make_client():
+            client = mock.Mock()
+
+            def fake_connect(hostname, **kwargs):
+                attempts.append(dict(kwargs))
+                outcome = connect_outcomes.pop(0)
+                if isinstance(outcome, Exception):
+                    raise outcome
+
+            client.connect.side_effect = fake_connect
+            client.close = mock.Mock()
+            return client
+
+        with (
+            mock.patch("pyinfra.connectors.ssh.Agent") as fake_agent,
+            mock.patch("pyinfra.connectors.ssh.SSHClient", side_effect=make_client),
+        ):
+            fake_agent.return_value.get_keys.return_value = ["key-one"]
+
+            result = connector._retry_paramiko_agent_keys(
+                "host",
+                {"allow_agent": True},
+                SSHException("No existing session"),
+            )
+
+        self.assertTrue(result)
+        self.assertEqual(
+            attempts,
+            [
+                {"allow_agent": False, "pkey": "key-one"},
+            ],
+        )
+
+    def test_retry_paramiko_agent_keys_returns_false_without_keys(self):
+        connector = ssh.SSHConnector.__new__(ssh.SSHConnector)
+        connector.client = mock.Mock()
+
+        with mock.patch("pyinfra.connectors.ssh.Agent") as fake_agent:
+            fake_agent.return_value.get_keys.return_value = []
+
+            result = connector._retry_paramiko_agent_keys(
+                "host",
+                {"allow_agent": True},
+                SSHException("No existing session"),
+            )
+
+        self.assertFalse(result)
+
+    @mock.patch("pyinfra.connectors.ssh.Agent")
+    def test_connect_retries_agent_keys_after_paramiko_failure(self, fake_agent):
+        key_one = mock.Mock(name="agent-key-1")
+        key_two = mock.Mock(name="agent-key-2")
+        fake_agent.return_value.get_keys.return_value = [key_one, key_two]
+
+        connect_calls = []
+
+        def fake_connect(hostname, **kwargs):
+            connect_calls.append((hostname, dict(kwargs)))
+            if len(connect_calls) == 1:
+                raise SSHException("No existing session")
+
+        self.fake_connect_mock.side_effect = fake_connect
+
+        inventory = make_inventory(hosts=("somehost",))
+        state = State(inventory, Config())
+
+        connect_all(state)
+
+        self.assertEqual(len(state.active_hosts), 1)
+        self.assertEqual(len(connect_calls), 2)
+
+        first_hostname, first_kwargs = connect_calls[0]
+        self.assertEqual(first_hostname, "somehost")
+        self.assertTrue(first_kwargs.get("allow_agent"))
+        self.assertNotIn("pkey", first_kwargs)
+
+        second_hostname, second_kwargs = connect_calls[1]
+        self.assertEqual(second_hostname, "somehost")
+        self.assertFalse(second_kwargs.get("allow_agent"))
+        self.assertIs(second_kwargs.get("pkey"), key_two)
+
     def test_connect_with_rsa_ssh_key_password(self):
         state = State(
             make_inventory(