perf: don't disconnect + wipe cache on command rejection

CHANGELOG.md

@@ -1,8 +1,9 @@
-# unreleased
+# Unreleased
 
 ## Performance
 
 - pyats broker: removed the per-command SSH liveness probe from the connection broker; dead sessions are now recovered by reconnecting and retrying the command once. Removes ~0.77s of event-loop blocking per test and the fleet-wide slowdown it caused.
+- pyats broker: command rejections (SubCommandFailure) no longer trigger a full SSH disconnect and cache wipe; only transport-level failures do. Eliminates a 4.7x per-failure penalty on devices with unsupported commands.
 
 # 2.1.0b1
 

nac_test/pyats_core/broker/connection_broker.py

@@ -292,16 +292,29 @@ async def _run_and_cache(self, hostname: str, connection: Any, cmd: str) -> str:
         """Run a command on a connection once and cache its output.
 
         Raises:
-            Exception: Whatever the device layer raised, after tearing the
-                connection down so it is not handed to the next caller.
+            SubCommandFailure: Re-raised as-is. The device answered and rejected
+                the command, so the session and its cache are left intact.
+            Exception: Any other failure, after tearing the connection down so it
+                is not handed to the next caller.
         """
         # Execute command in thread pool (since Unicon is synchronous)
         loop = get_or_create_event_loop()
+        try:
+            from unicon.core.errors import SubCommandFailure
+        except ImportError:
+            SubCommandFailure = None  # type: ignore[assignment,misc]
+
         try:
             output = await loop.run_in_executor(None, connection.execute, cmd)
         except Exception as e:
+            # SubCommandFailure is a fast path, not a health guarantee - unicon wraps
+            # transport errors in it too. Those get neither a disconnect here nor a retry
+            # in _execute_command; recovery relies on unicon re-establishing the session
+            # on the next execute(), which it does transparently.
+            if SubCommandFailure is not None and isinstance(e, SubCommandFailure):
+                logger.warning(f"Command rejected by {hostname} (session intact): {e}")
+                raise
             logger.error(f"Command execution failed on {hostname}: {e}")
-            # Try to reconnect on failure
             await self._disconnect_device(hostname)
             raise
 

tests/integration/test_connection_broker_failure_modes.py

@@ -27,6 +27,7 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
+from unicon.core.errors import SubCommandFailure
 
 from nac_test.pyats_core.broker.broker_client import BrokerClient
 from nac_test.pyats_core.broker.connection_broker import ConnectionBroker
@@ -680,3 +681,144 @@ async def _run() -> None:
 
         assert broker.connected_devices == {}
         assert not broker.socket_path.exists()
+
+
+class TestCommandRejectionVsTransportFailure:
+    """Contrast SubCommandFailure (device rejects command, session healthy) with
+    transport-level failures (SSH dies, cache must be wiped).
+
+    Both tests populate the command cache with a successful first command, then
+    trigger a failure on a second command and verify the divergent outcomes.
+    """
+
+    def test_sub_command_failure_preserves_connection_and_cache(
+        self,
+        make_broker: Any,
+        good_device: MagicMock,
+    ) -> None:
+        """SubCommandFailure re-raises to the client but does NOT disconnect
+        the device or wipe the command cache."""
+        broker: ConnectionBroker = make_broker({"router-1": good_device})
+
+        async def _run() -> None:
+            loop = asyncio.get_event_loop()
+
+            async def _body() -> None:
+                with patch(
+                    "nac_test.pyats_core.broker.connection_broker.get_or_create_event_loop",
+                    return_value=loop,
+                ):
+                    with _patch_executor(loop):
+                        async with BrokerClient(
+                            socket_path=broker.socket_path
+                        ) as client:
+                            # 1) Successful command — populates cache
+                            good_device.execute.return_value = "good output"
+                            r1 = await asyncio.wait_for(
+                                client._send_request(
+                                    {
+                                        "command": "execute",
+                                        "hostname": "router-1",
+                                        "cmd": "show version",
+                                    }
+                                ),
+                                timeout=2.0,
+                            )
+                            assert r1["status"] == "success"
+                            assert r1["result"] == "good output"
+
+                            # 2) SubCommandFailure on second command
+                            good_device.execute.side_effect = SubCommandFailure(
+                                "Invalid command at '^' marker"
+                            )
+                            err = await _expect_broker_error(
+                                client,
+                                {
+                                    "command": "execute",
+                                    "hostname": "router-1",
+                                    "cmd": "show bgp all",
+                                },
+                                timeout=2.0,
+                            )
+                            assert "Invalid command" in err
+
+                            # Connection still present
+                            assert "router-1" in broker.connected_devices
+
+                            # First command's cache entry still intact
+                            cache = broker.command_cache.get("router-1")
+                            assert cache is not None
+                            assert cache.get("show version") == "good output"
+
+                            # Device was never disconnected
+                            good_device.disconnect.assert_not_called()
+
+            await _run_broker(broker, _body())
+
+        asyncio.run(_run())
+
+    def test_transport_failure_disconnects_and_wipes_cache(
+        self,
+        make_broker: Any,
+        good_device: MagicMock,
+    ) -> None:
+        """A transport-level exception (e.g. OSError) disconnects the device
+        and wipes the command cache — regression guard for existing behavior."""
+        broker: ConnectionBroker = make_broker({"router-1": good_device})
+
+        async def _run() -> None:
+            loop = asyncio.get_event_loop()
+
+            async def _body() -> None:
+                with patch(
+                    "nac_test.pyats_core.broker.connection_broker.get_or_create_event_loop",
+                    return_value=loop,
+                ):
+                    with _patch_executor(loop):
+                        async with BrokerClient(
+                            socket_path=broker.socket_path
+                        ) as client:
+                            # 1) Successful command — populates cache
+                            good_device.execute.return_value = "good output"
+                            r1 = await asyncio.wait_for(
+                                client._send_request(
+                                    {
+                                        "command": "execute",
+                                        "hostname": "router-1",
+                                        "cmd": "show version",
+                                    }
+                                ),
+                                timeout=2.0,
+                            )
+                            assert r1["status"] == "success"
+                            assert r1["result"] == "good output"
+
+                            # 2) Transport failure on second command
+                            good_device.execute.side_effect = OSError(
+                                "Socket is closed"
+                            )
+                            err = await _expect_broker_error(
+                                client,
+                                {
+                                    "command": "execute",
+                                    "hostname": "router-1",
+                                    "cmd": "show interfaces",
+                                },
+                                timeout=2.0,
+                            )
+                            assert "Socket is closed" in err
+
+                            # OSError is a transport failure, so this exercises the full #899 path:
+                            # fail -> disconnect -> reconnect -> fail -> disconnect -> raise
+                            assert good_device.execute.call_count == 3
+                            assert good_device.disconnect.call_count == 2
+
+                            # Connection removed
+                            assert "router-1" not in broker.connected_devices
+
+                            # Command cache wiped for this device
+                            assert "router-1" not in broker.command_cache
+
+            await _run_broker(broker, _body())
+
+        asyncio.run(_run())