fix(python): close client when Sandbox entry fails (#4367)

* fix(python): close client when Sandbox entry fails

Signed-off-by: Eric Curtin <eric.curtin@docker.com>

* test(python): share client patch helper

Signed-off-by: Eric Curtin <eric.curtin@docker.com>

---------

Signed-off-by: Eric Curtin <eric.curtin@docker.com>

python/openshell/sandbox.py

@@ -1643,7 +1643,7 @@ class Sandbox:
                 "name, labels, and workload_template cannot be set when attaching to an existing sandbox"
             )
 
-        client = SandboxClient.from_active_cluster(
+        self._client = SandboxClient.from_active_cluster(
             cluster=self._cluster,
             timeout=self._timeout,
             auto_refresh=self._auto_refresh,
@@ -1651,8 +1651,17 @@ class Sandbox:
             insecure=self._insecure,
             client_credentials=self._client_credentials,
         )
-        self._client = client
+        try:
+            self._open_session(self._client)
+        except BaseException:
+            # Python skips __exit__ when __enter__ raises, so release here
+            # and keep the original error even if the close fails.
+            with contextlib.suppress(Exception):
+                self._release()
+            raise
+        return self
 
+    def _open_session(self, client: SandboxClient) -> None:
         if self._sandbox_input is None and self._workload_template is not None:
             self._session = client.create_session_from_template(
                 workspace=self._workspace,
@@ -1684,7 +1693,12 @@ class Sandbox:
         )
         self._session = SandboxSession(client, ready)
 
-        return self
+    def _release(self) -> None:
+        client = self._client
+        self._session = None
+        self._client = None
+        if client is not None:
+            client.close()
 
     def __exit__(self, *args: object) -> None:
         try:
@@ -1708,10 +1722,7 @@ class Sandbox:
                         f"unsupported deletion outcome: {result.outcome}"
                     )
         finally:
-            if self._client is not None:
-                self._client.close()
-            self._session = None
-            self._client = None
+            self._release()
 
     def exec(
         self,

python/openshell/sandbox_cleanup_test.py

@@ -156,3 +156,90 @@ def test_context_cleanup_handles_absence_and_intercepted_errors(
     assert state.closed == [True]
     assert managed._client is None
     assert managed._session is None
+
+
+class _EntryFailureClient:
+    """Fails the named client call and records close()."""
+
+    def __init__(self, fail_at, close_error=None):
+        self.fail_at = fail_at
+        self.close_error = close_error
+        self.closed = 0
+        self.error = RuntimeError(f"{fail_at} failed")
+
+    def _step(self, name):
+        if name == self.fail_at:
+            raise self.error
+        return SimpleNamespace(sandbox=SimpleNamespace(name="s1"))
+
+    def create_session(self, **_kwargs):
+        return self._step("create_session")
+
+    def create_session_from_template(self, **_kwargs):
+        return self._step("create_session_from_template")
+
+    def get_session(self, *_args, **_kwargs):
+        return self._step("get_session")
+
+    def wait_ready(self, *_args, **_kwargs):
+        self._step("wait_ready")
+        return SimpleNamespace(name="s1", workspace="default")
+
+    def close(self):
+        self.closed += 1
+        if self.close_error is not None:
+            raise self.close_error
+
+
+def _use_client(monkeypatch, client):
+    monkeypatch.setattr(
+        SandboxClient, "from_active_cluster", classmethod(lambda _cls, **_kw: client)
+    )
+
+
+@pytest.mark.parametrize(
+    ("fail_at", "kwargs"),
+    [
+        ("create_session", {}),
+        ("create_session_from_template", {"workload_template": "t"}),
+        ("get_session", {"sandbox": "existing"}),
+        ("wait_ready", {}),
+        ("wait_ready", {"sandbox": "existing"}),
+    ],
+)
+def test_context_entry_failure_closes_client(monkeypatch, fail_at, kwargs):
+    client = _EntryFailureClient(fail_at)
+    _use_client(monkeypatch, client)
+    managed = Sandbox(workspace="default", delete_on_exit=False, **kwargs)
+
+    with pytest.raises(RuntimeError) as observed, managed:
+        pass
+
+    assert observed.value is client.error
+    assert client.closed == 1
+    assert managed._client is None
+    assert managed._session is None
+
+
+def test_context_entry_failure_keeps_original_error_when_close_fails(monkeypatch):
+    client = _EntryFailureClient("wait_ready", close_error=OSError("close failed"))
+    _use_client(monkeypatch, client)
+    managed = Sandbox(workspace="default")
+
+    with pytest.raises(RuntimeError) as observed, managed:
+        pass
+
+    assert observed.value is client.error
+    assert client.closed == 1
+    assert managed._client is None
+    assert managed._session is None
+
+
+def test_context_closes_client_once_on_success(monkeypatch):
+    client = _EntryFailureClient(None)
+    _use_client(monkeypatch, client)
+
+    with Sandbox(workspace="default", delete_on_exit=False):
+        assert client.closed == 0
+
+    assert client.closed == 1