Make LoginManager close its clients on close()

Now that clients provide a `close()` method, the LoginManager can use it
to implement a better closing semantic. When a login manager creates a
client object, the client is tracked in a set. When the login manager is
closed, all its clients are closed and the set is cleared.

Author: sirosen

changelog.d/20251205_165556_sirosen_close_every_client.md

@@ -0,0 +1,3 @@
+### Bugfixes
+
+* Improved internal resource cleanup for network connections.

src/globus_cli/login_manager/manager.py

@@ -48,14 +48,26 @@
 P = ParamSpec("P")
 R = t.TypeVar("R")
 
+# string-ized annotation to avoid triggering eager imports
+CLIENT_T = t.TypeVar("CLIENT_T", bound="globus_sdk.BaseClient")
+
 
 class LoginManager:
     def __init__(self) -> None:
         self.storage = CLIStorage()
         self._nonstatic_requirements: dict[str, list[Scope]] = {}
 
+        self._client_pool: set[globus_sdk.BaseClient] = set()
+
     def close(self) -> None:
         self.storage.close()
+        for c in self._client_pool:
+            c.close()
+        self._client_pool.clear()
+
+    def _into_client_pool(self, client: CLIENT_T) -> CLIENT_T:
+        self._client_pool.add(client)
+        return client
 
     def add_requirement(self, rs_name: str, scopes: t.Sequence[Scope]) -> None:
         self._nonstatic_requirements[rs_name] = list(scopes)
@@ -373,25 +385,35 @@ def get_transfer_client(self) -> CustomTransferClient:
         from ..services.transfer import CustomTransferClient
 
         authorizer = self._get_client_authorizer(TransferScopes.resource_server)
-        return CustomTransferClient(authorizer=authorizer, app_name=version.app_name)
+        return self._into_client_pool(
+            CustomTransferClient(authorizer=authorizer, app_name=version.app_name)
+        )
 
     def get_auth_client(self) -> CustomAuthClient:
         from ..services.auth import CustomAuthClient
 
         authorizer = self._get_client_authorizer(AuthScopes.resource_server)
-        return CustomAuthClient(authorizer=authorizer, app_name=version.app_name)
+        return self._into_client_pool(
+            CustomAuthClient(authorizer=authorizer, app_name=version.app_name)
+        )
 
     def get_groups_client(self) -> globus_sdk.GroupsClient:
         authorizer = self._get_client_authorizer(GroupsScopes.resource_server)
-        return globus_sdk.GroupsClient(authorizer=authorizer, app_name=version.app_name)
+        return self._into_client_pool(
+            globus_sdk.GroupsClient(authorizer=authorizer, app_name=version.app_name)
+        )
 
     def get_flows_client(self) -> globus_sdk.FlowsClient:
         authorizer = self._get_client_authorizer(FlowsScopes.resource_server)
-        return globus_sdk.FlowsClient(authorizer=authorizer, app_name=version.app_name)
+        return self._into_client_pool(
+            globus_sdk.FlowsClient(authorizer=authorizer, app_name=version.app_name)
+        )
 
     def get_search_client(self) -> globus_sdk.SearchClient:
         authorizer = self._get_client_authorizer(SearchScopes.resource_server)
-        return globus_sdk.SearchClient(authorizer=authorizer, app_name=version.app_name)
+        return self._into_client_pool(
+            globus_sdk.SearchClient(authorizer=authorizer, app_name=version.app_name)
+        )
 
     def get_timer_client(
         self, *, flow_id: uuid.UUID | None = None
@@ -404,7 +426,9 @@ def get_timer_client(
             self._assert_requester_has_timer_flow_consent(flow_id)
 
         authorizer = self._get_client_authorizer(TimersScopes.resource_server)
-        return globus_sdk.TimersClient(authorizer=authorizer, app_name=version.app_name)
+        return self._into_client_pool(
+            globus_sdk.TimersClient(authorizer=authorizer, app_name=version.app_name)
+        )
 
     def _assert_requester_has_timer_flow_consent(self, flow_id: uuid.UUID) -> None:
         flow_scope = SpecificFlowScopes(flow_id).user
@@ -445,7 +469,9 @@ def get_specific_flow_client(
     ) -> globus_sdk.SpecificFlowClient:
         # Create a SpecificFlowClient without an authorizer
         # to take advantage of its scope creation code.
-        client = globus_sdk.SpecificFlowClient(flow_id, app_name=version.app_name)
+        client = self._into_client_pool(
+            globus_sdk.SpecificFlowClient(flow_id, app_name=version.app_name)
+        )
         assert client.scopes is not None
         self.add_requirement(client.scopes.resource_server, [client.scopes.user])
 
@@ -531,11 +557,13 @@ def get_gcs_client(
                 f"Please run:\n\n  {login_context.login_command}\n"
             ),
         )
-        return CustomGCSClient(
-            epish.get_gcs_address(),
-            source_epish=epish,
-            authorizer=authorizer,
-            app_name=version.app_name,
+        return self._into_client_pool(
+            CustomGCSClient(
+                epish.get_gcs_address(),
+                source_epish=epish,
+                authorizer=authorizer,
+                app_name=version.app_name,
+            )
         )
 
     def get_current_identity_id(self) -> str:

tests/unit/test_login_manager.py

@@ -4,6 +4,7 @@
 import uuid
 from unittest import mock
 
+import click
 import globus_sdk
 import globus_sdk.scopes
 import jwt
@@ -428,3 +429,34 @@ def test_immature_signature_during_jwt_decode_skips_notice_if_date_cannot_parse(
 
     stderr = capsys.readouterr().err
     assert "This may indicate a clock skew problem." not in stderr
+
+
+def test_login_manager_closes_created_clients_after_context_exit(test_token_storage):
+    from globus_cli.commands import main
+
+    def add_close_spy(client):
+        real_close = client.close
+        spy = mock.Mock(side_effect=real_close)
+        client.close = spy
+        return spy
+
+    with click.Context(main):
+        spy1, spy2 = None, None
+
+        @LoginManager.requires_login("auth", "transfer")
+        def dummy_command(login_manager):
+            transfer_client = login_manager.get_transfer_client()
+            auth_client = login_manager.get_auth_client()
+
+            nonlocal spy1, spy2
+            spy1 = add_close_spy(transfer_client)
+            spy2 = add_close_spy(auth_client)
+            return True
+
+        assert dummy_command()
+
+        spy1.assert_not_called()
+        spy2.assert_not_called()
+
+    spy1.assert_called_once()
+    spy2.assert_called_once()