feat(vpn): auto-disconnect VPN on project and organization switch

riocli/organization/select.py

@@ -20,7 +20,12 @@
 from riocli.constants import Colors, Symbols
 from riocli.organization.util import name_to_guid
 from riocli.utils.context import get_root_context
-from riocli.vpn.util import cleanup_hosts_file
+from riocli.vpn.util import (
+    cleanup_hosts_file,
+    is_tailscale_up,
+    should_disconnect_vpn,
+    stop_tailscale,
+)
 
 
 @click.command(
@@ -30,6 +35,13 @@
     help_options_color=Colors.GREEN,
 )
 @click.argument("organization-name", type=str)
+@click.option(
+    "--keep-vpn",
+    is_flag=True,
+    default=False,
+    help="Keep the VPN connected after switching organizations. Skips both "
+    "VPN disconnect and hosts file cleanup.",
+)
 @click.option(
     "--interactive/--no-interactive",
     is_flag=True,
@@ -51,6 +63,7 @@ def select_organization(
     organization_name: str,
     organization_guid: str,
     organization_short_id: str,
+    keep_vpn: bool,
     interactive: bool,
     silent: bool,
 ) -> None:
@@ -103,13 +116,25 @@ def select_organization(
 
     ctx.obj.save()
 
-    try:
-        cleanup_hosts_file()
-    except Exception as e:
-        click.secho(
-            f"{Symbols.WARNING} Failed to clean up hosts file: {str(e)}",
-            fg=Colors.YELLOW,
-        )
+    if should_disconnect_vpn(ctx.obj.data, keep_vpn):
+        if is_tailscale_up():
+            if stop_tailscale():
+                click.secho(
+                    f"{Symbols.SUCCESS} VPN disconnected.",
+                    fg=Colors.GREEN,
+                )
+            else:
+                click.secho(
+                    f"{Symbols.WARNING} Failed to disconnect VPN.",
+                    fg=Colors.YELLOW,
+                )
+        try:
+            cleanup_hosts_file()
+        except Exception as e:
+            click.secho(
+                f"{Symbols.WARNING} Failed to clean up hosts file: {str(e)}",
+                fg=Colors.YELLOW,
+            )
 
     if ctx.obj.data.get("project_id"):
         from riocli.ssh import refresh_ssh_cert

riocli/project/select.py

@@ -17,7 +17,12 @@
 from riocli.constants import Colors, Symbols
 from riocli.project.util import name_to_guid
 from riocli.utils.context import get_root_context
-from riocli.vpn.util import cleanup_hosts_file
+from riocli.vpn.util import (
+    cleanup_hosts_file,
+    is_tailscale_up,
+    should_disconnect_vpn,
+    stop_tailscale,
+)
 
 
 @click.command(
@@ -27,31 +32,57 @@
     help_options_color=Colors.GREEN,
 )
 @click.argument("project-name", type=str)
+@click.option(
+    "--keep-vpn",
+    is_flag=True,
+    default=False,
+    help="Keep the VPN connected after switching projects. Skips both "
+    "VPN disconnect and hosts file cleanup.",
+)
 @name_to_guid
 @click.pass_context
 def select_project(
     ctx: click.Context,
     project_name: str,
     project_guid: str,
+    keep_vpn: bool,
 ) -> None:
     """Switch to a different project in the current organization.
 
     The project will be set in the CLI's context and will be used
     for all the subsequent commands.
+
+    By default, if a VPN is active it will be disconnected and the
+    hosts file will be cleaned up. Use --keep-vpn to suppress this,
+    for example when you have an active SSH session into a device on
+    the previous project. You can also set ``auto_disconnect_vpn: false``
+    in ~/.rio-cli/config.json to permanently suppress auto-disconnect.
     """
     ctx = get_root_context(ctx)
 
     ctx.obj.data["project_id"] = project_guid
     ctx.obj.data["project_name"] = project_name
     ctx.obj.save()
 
-    try:
-        cleanup_hosts_file()
-    except Exception as e:
-        click.secho(
-            f"{Symbols.WARNING} Failed to clean up hosts file: {str(e)}",
-            fg=Colors.YELLOW,
-        )
+    if should_disconnect_vpn(ctx.obj.data, keep_vpn):
+        if is_tailscale_up():
+            if stop_tailscale():
+                click.secho(
+                    f"{Symbols.SUCCESS} VPN disconnected.",
+                    fg=Colors.GREEN,
+                )
+            else:
+                click.secho(
+                    f"{Symbols.WARNING} Failed to disconnect VPN.",
+                    fg=Colors.YELLOW,
+                )
+        try:
+            cleanup_hosts_file()
+        except Exception as e:
+            click.secho(
+                f"{Symbols.WARNING} Failed to clean up hosts file: {str(e)}",
+                fg=Colors.YELLOW,
+            )
 
     click.secho(
         f"{Symbols.SUCCESS} Project {project_name} ({project_guid}) is selected!",

riocli/vpn/util.py

@@ -87,6 +87,17 @@ def get_tailscale_status() -> dict:
     return json.loads(output)
 
 
+def should_disconnect_vpn(config: dict, keep_vpn: bool) -> bool:
+    """Returns True if VPN should be auto-disconnected on project/org switch.
+
+    Disconnect is skipped if --keep-vpn flag is passed, or if the user has
+    set auto_disconnect_vpn: false in their CLI config (~/.rio-cli/config.json).
+    """
+    if keep_vpn:
+        return False
+    return config.get("auto_disconnect_vpn", True)
+
+
 def install_vpn_tools(force: bool = False) -> None:
     if is_tailscale_installed():
         return

tests/unit/vpn/test_auto_disconnect.py

@@ -0,0 +1,152 @@
+from unittest.mock import MagicMock, patch
+
+from click.testing import CliRunner
+
+from riocli.vpn.util import should_disconnect_vpn
+
+
+class TestShouldDisconnectVpn:
+    def test_disconnects_by_default(self):
+        assert should_disconnect_vpn({}, keep_vpn=False) is True
+
+    def test_keep_vpn_flag_suppresses_disconnect(self):
+        assert should_disconnect_vpn({}, keep_vpn=True) is False
+
+    def test_config_opt_out_suppresses_disconnect(self):
+        assert should_disconnect_vpn({"auto_disconnect_vpn": False}, keep_vpn=False) is False
+
+    def test_config_opt_out_with_keep_vpn_flag(self):
+        assert should_disconnect_vpn({"auto_disconnect_vpn": False}, keep_vpn=True) is False
+
+    def test_config_explicitly_true_disconnects(self):
+        assert should_disconnect_vpn({"auto_disconnect_vpn": True}, keep_vpn=False) is True
+
+    def test_keep_vpn_flag_overrides_config_true(self):
+        assert should_disconnect_vpn({"auto_disconnect_vpn": True}, keep_vpn=True) is False
+
+
+def _make_project_ctx(config_data=None):
+    obj = MagicMock()
+    obj.data = {
+        "project_id": "old-guid",
+        "project_name": "old-project",
+        "organization_id": "org-guid",
+        **(config_data or {}),
+    }
+    return obj
+
+
+def _make_org_ctx(config_data=None):
+    obj = MagicMock()
+    obj.data = {
+        "organization_id": "different-org-guid",  # different so "already in org" check passes
+        "organization_name": "old-org",
+        "organization_short_id": "old-short",
+        **(config_data or {}),
+    }
+    return obj
+
+
+PROJECT_PATCHES = [
+    patch("riocli.project.util.new_v2_client"),
+    patch("riocli.project.util.find_project_guid", return_value="new-project-guid"),
+    patch("riocli.project.util.get_project_name", return_value="new-project"),
+    patch("riocli.project.select.get_root_context"),
+]
+
+ORG_PATCHES = [
+    patch("riocli.organization.util.new_v2_client"),
+    patch("riocli.organization.util.find_organization_guid", return_value=("new-org-guid", "new-short")),
+    patch("riocli.organization.select.get_root_context"),
+]
+
+
+class TestProjectSelectVpnDisconnect:
+    def _invoke(self, args, ctx_obj):
+        from riocli.project.select import select_project
+
+        with (
+            patch("riocli.project.util.new_v2_client"),
+            patch("riocli.project.util.find_project_guid", return_value="new-guid"),
+            patch("riocli.project.util.get_project_name", return_value="new-project"),
+            patch("riocli.project.select.get_root_context") as mock_get_ctx,
+        ):
+            mock_get_ctx.return_value.obj = ctx_obj
+            # Pass obj so click.get_current_context().obj works inside name_to_guid
+            return CliRunner().invoke(select_project, args, obj=ctx_obj), mock_get_ctx
+
+    @patch("riocli.project.select.is_tailscale_up", return_value=True)
+    @patch("riocli.project.select.stop_tailscale", return_value=True)
+    @patch("riocli.project.select.cleanup_hosts_file")
+    def test_disconnects_vpn_when_tailscale_up(self, mock_cleanup, mock_stop, mock_is_up):
+        result, _ = self._invoke(["new-project"], _make_project_ctx())
+        assert result.exit_code == 0
+        mock_stop.assert_called_once()
+        mock_cleanup.assert_called_once()
+
+    @patch("riocli.project.select.is_tailscale_up", return_value=False)
+    @patch("riocli.project.select.stop_tailscale")
+    @patch("riocli.project.select.cleanup_hosts_file")
+    def test_skips_stop_when_tailscale_not_up(self, mock_cleanup, mock_stop, mock_is_up):
+        result, _ = self._invoke(["new-project"], _make_project_ctx())
+        assert result.exit_code == 0
+        mock_stop.assert_not_called()
+        mock_cleanup.assert_called_once()
+
+    @patch("riocli.project.select.is_tailscale_up", return_value=True)
+    @patch("riocli.project.select.stop_tailscale")
+    @patch("riocli.project.select.cleanup_hosts_file")
+    def test_keep_vpn_flag_skips_disconnect_and_cleanup(self, mock_cleanup, mock_stop, mock_is_up):
+        result, _ = self._invoke(["new-project", "--keep-vpn"], _make_project_ctx())
+        assert result.exit_code == 0
+        mock_stop.assert_not_called()
+        mock_cleanup.assert_not_called()
+
+    @patch("riocli.project.select.is_tailscale_up", return_value=True)
+    @patch("riocli.project.select.stop_tailscale")
+    @patch("riocli.project.select.cleanup_hosts_file")
+    def test_config_opt_out_skips_disconnect_and_cleanup(self, mock_cleanup, mock_stop, mock_is_up):
+        result, _ = self._invoke(["new-project"], _make_project_ctx({"auto_disconnect_vpn": False}))
+        assert result.exit_code == 0
+        mock_stop.assert_not_called()
+        mock_cleanup.assert_not_called()
+
+
+class TestOrgSelectVpnDisconnect:
+    def _invoke(self, args, ctx_obj):
+        from riocli.organization.select import select_organization
+
+        with (
+            patch("riocli.organization.util.new_v2_client"),
+            patch("riocli.organization.util.find_organization_guid", return_value=("new-org-guid", "new-short")),
+            patch("riocli.organization.select.get_root_context") as mock_get_ctx,
+        ):
+            mock_get_ctx.return_value.obj = ctx_obj
+            return CliRunner().invoke(select_organization, args), mock_get_ctx
+
+    @patch("riocli.organization.select.is_tailscale_up", return_value=True)
+    @patch("riocli.organization.select.stop_tailscale", return_value=True)
+    @patch("riocli.organization.select.cleanup_hosts_file")
+    def test_disconnects_vpn_when_tailscale_up(self, mock_cleanup, mock_stop, mock_is_up):
+        result, _ = self._invoke(["new-org", "--no-interactive"], _make_org_ctx())
+        assert result.exit_code == 0
+        mock_stop.assert_called_once()
+        mock_cleanup.assert_called_once()
+
+    @patch("riocli.organization.select.is_tailscale_up", return_value=True)
+    @patch("riocli.organization.select.stop_tailscale")
+    @patch("riocli.organization.select.cleanup_hosts_file")
+    def test_keep_vpn_flag_skips_disconnect_and_cleanup(self, mock_cleanup, mock_stop, mock_is_up):
+        result, _ = self._invoke(["new-org", "--keep-vpn", "--no-interactive"], _make_org_ctx())
+        assert result.exit_code == 0
+        mock_stop.assert_not_called()
+        mock_cleanup.assert_not_called()
+
+    @patch("riocli.organization.select.is_tailscale_up", return_value=True)
+    @patch("riocli.organization.select.stop_tailscale")
+    @patch("riocli.organization.select.cleanup_hosts_file")
+    def test_config_opt_out_skips_disconnect_and_cleanup(self, mock_cleanup, mock_stop, mock_is_up):
+        result, _ = self._invoke(["new-org", "--no-interactive"], _make_org_ctx({"auto_disconnect_vpn": False}))
+        assert result.exit_code == 0
+        mock_stop.assert_not_called()
+        mock_cleanup.assert_not_called()