Show KeeperAI-closed banner for pam tunnel start command

keepercommander/commands/tunnel/port_forward/tunnel_helpers.py

@@ -167,6 +167,21 @@ def print_above_keeper_prompt(msg):
             pass
 
 
+def _print_ai_closed_banner():
+    """User-facing notice for a tunnel closed with the AIClosed reason.
+
+    Wording matches `pam_launch.launch`'s `_print_close_reason_notice` for the
+    same condition — duplicated rather than shared across the two commands,
+    per this codebase's preference for small per-command copies over a
+    premature cross-module helper.
+    """
+    print_above_keeper_prompt(
+        f"\n{bcolors.FAIL}Session terminated by KeeperAI.{bcolors.ENDC}\n"
+        "Critical activity was detected during this session.\n"
+        "Contact your administrator to unlock this record."
+    )
+
+
 # Constants
 NONCE_LENGTH = 12
 MAIN_NONCE_LENGTH = 16
@@ -2028,6 +2043,27 @@ def _signal_from_rust_inner(self, response: dict):
                 elif close_reason.is_retryable():
                     logging.debug(f"{bcolors.WARNING}Tunnel closed with retryable error - '{tube_id}': {close_reason.name}{bcolors.ENDC}")
 
+                elif close_reason.code == CloseConnectionReasons.AIClosed:
+                    # Explicit branch (rather than relying on the generic `else`
+                    # below, where this currently lands) so the user actually
+                    # sees why their session ended — mirrors `pam launch`'s
+                    # handling of the same close reason.
+                    logging.debug(f"{bcolors.OKBLUE}Tunnel '{tube_id}' closed by KeeperAI{bcolors.ENDC}")
+                    _print_ai_closed_banner()
+                    if self.record_uid:
+                        try:
+                            from ...tunnel_and_connections import (
+                                _AI_CLOSED_EXIT_CODE_BY_RECORD,
+                                _LEASE_SHUTDOWN_EVENTS_BY_RECORD,
+                                EXIT_CODE_AI_TERMINATED,
+                            )
+                            _AI_CLOSED_EXIT_CODE_BY_RECORD[self.record_uid] = EXIT_CODE_AI_TERMINATED
+                            shutdown_event = _LEASE_SHUTDOWN_EVENTS_BY_RECORD.get(self.record_uid)
+                            if shutdown_event is not None:
+                                shutdown_event.set()
+                        except Exception:
+                            logging.debug("Could not wake --foreground wait for AI-closed tunnel", exc_info=True)
+
                 else:
                     logging.debug(f"{bcolors.OKBLUE}Tunnel '{tube_id}' closed with reason: {close_reason.name}{bcolors.ENDC}")
 

keepercommander/commands/tunnel_and_connections.py

@@ -39,6 +39,7 @@
 from .pam_import.keeper_ai_settings import get_resource_jit_settings
 from .pam_import.nsf_helpers import sync_down_preserving_nsf_keys
 from .pam_launch.terminal_connection import _extract_user_record_credentials
+from .pam_launch.launch import EXIT_CODE_AI_TERMINATED
 from .tunnel_registry import (
     PARENT_GRACE_SECONDS,
     is_pid_alive,
@@ -67,6 +68,11 @@
 # read by the lease-expiry callback. Default interactive mode does NOT register
 # (it has no blocking wait to interrupt; user SSH session continues naturally).
 _LEASE_SHUTDOWN_EVENTS_BY_RECORD: Dict[str, threading.Event] = {}
+# Maps record_uid -> exit code set by tunnel_helpers' channel_closed handler when
+# the Gateway/Rust reports an AIClosed close reason. Read by the --foreground
+# block below after fg_shutdown wakes, so the process exits with the same code
+# `pam launch` uses for the same condition.
+_AI_CLOSED_EXIT_CODE_BY_RECORD: Dict[str, int] = {}
 
 
 def _coerce_settings_subdicts(entry, *keys):
@@ -1699,11 +1705,13 @@ def _fg_signal_handler(signum, _frame):
                         print(f"{bcolors.WARNING}Note: On Windows, tunnel stop uses hard termination. "
                               f"WebRTC cleanup is best-effort.{bcolors.ENDC}\n")
 
+                    ai_exit_code = None
                     try:
                         fg_shutdown.wait()
                     except KeyboardInterrupt:
                         pass
                     finally:
+                        ai_exit_code = _AI_CLOSED_EXIT_CODE_BY_RECORD.pop(record_uid, None)
                         _LEASE_SHUTDOWN_EVENTS_BY_RECORD.pop(record_uid, None)
                         unregister_tunnel()
                         signal.signal(signal.SIGTERM, prev_sigterm)
@@ -1727,6 +1735,14 @@ def _fg_signal_handler(signum, _frame):
                                     os.remove(pid_file)
                                 except OSError:
                                     pass
+                    # Mirror `pam launch`'s exit code for the same condition, but only
+                    # in batch / scripted mode. In the interactive shell a SystemExit
+                    # would tear down the whole shell (and its other sessions); the
+                    # user should land back at the Keeper> prompt. The banner itself
+                    # is printed from tunnel_helpers' close_reason handling as soon as
+                    # the signal arrives, not here.
+                    if ai_exit_code is not None and getattr(params, 'batch_mode', False):
+                        raise SystemExit(ai_exit_code)
         else:
             # Print failure message
             error_msg = result.get("error", "Unknown error") if result else "Failed to start tunnel"

unit-tests/pam/test_pam_tunnel_ai_closed_banner.py

@@ -0,0 +1,93 @@
+"""
+Unit tests for `pam tunnel start`'s handling of the AIClosed (code 15) close
+reason — the signal already reaches `TunnelSignalHandler` today, but used to
+fall through to a silent `logging.debug(...)` with nothing shown to the user.
+"""
+import io
+import json
+import threading
+import unittest
+from contextlib import redirect_stdout
+
+from keepercommander.commands.pam_launch.launch import EXIT_CODE_AI_TERMINATED
+from keepercommander.commands.tunnel.port_forward.tunnel_helpers import (
+    TunnelSignalHandler,
+    _print_ai_closed_banner,
+)
+from keepercommander.commands.tunnel_and_connections import (
+    _AI_CLOSED_EXIT_CODE_BY_RECORD,
+    _LEASE_SHUTDOWN_EVENTS_BY_RECORD,
+)
+
+
+class TestPrintAIClosedBanner(unittest.TestCase):
+
+    def test_banner_text(self):
+        buf = io.StringIO()
+        with redirect_stdout(buf):
+            _print_ai_closed_banner()
+        out = buf.getvalue()
+        self.assertIn('Session terminated by KeeperAI.', out)
+        self.assertIn('Critical activity was detected during this session.', out)
+        self.assertIn('Contact your administrator to unlock this record.', out)
+
+
+class TestTunnelSignalHandlerAIClosed(unittest.TestCase):
+
+    def setUp(self):
+        self.record_uid = 'test-ai-closed-record-uid'
+        self.addCleanup(_AI_CLOSED_EXIT_CODE_BY_RECORD.pop, self.record_uid, None)
+        self.addCleanup(_LEASE_SHUTDOWN_EVENTS_BY_RECORD.pop, self.record_uid, None)
+
+    def _make_handler(self):
+        return TunnelSignalHandler(
+            params=None,
+            record_uid=self.record_uid,
+            gateway_uid='gateway-uid',
+            symmetric_key=b'0' * 32,
+            base64_nonce='nonce',
+            conversation_id='conversation-id',
+            tube_registry=None,
+        )
+
+    def _send_ai_closed(self, handler):
+        payload = {
+            'kind': 'channel_closed',
+            'tube_id': 'fake-tube-id',
+            'conversation_id': 'conversation-id',
+            'data': json.dumps({'close_reason': {'code': 15}}),
+        }
+        buf = io.StringIO()
+        with redirect_stdout(buf):
+            handler.signal_from_rust(payload)
+        return buf.getvalue()
+
+    def test_prints_banner_and_sets_exit_code(self):
+        handler = self._make_handler()
+        out = self._send_ai_closed(handler)
+        self.assertIn('Session terminated by KeeperAI.', out)
+        self.assertEqual(
+            _AI_CLOSED_EXIT_CODE_BY_RECORD.get(self.record_uid),
+            EXIT_CODE_AI_TERMINATED,
+        )
+
+    def test_wakes_registered_foreground_shutdown_event(self):
+        shutdown_event = threading.Event()
+        _LEASE_SHUTDOWN_EVENTS_BY_RECORD[self.record_uid] = shutdown_event
+        handler = self._make_handler()
+        self._send_ai_closed(handler)
+        self.assertTrue(shutdown_event.is_set())
+
+    def test_no_shutdown_event_registered_does_not_raise(self):
+        handler = self._make_handler()
+        # No entry in _LEASE_SHUTDOWN_EVENTS_BY_RECORD for this record_uid —
+        # mirrors interactive-shell mode, which never registers one.
+        self._send_ai_closed(handler)
+        self.assertEqual(
+            _AI_CLOSED_EXIT_CODE_BY_RECORD.get(self.record_uid),
+            EXIT_CODE_AI_TERMINATED,
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()