diff --git a/python/openshell/sandbox.py b/python/openshell/sandbox.py index 19a7ba3b8..e9c840e48 100644 --- a/python/openshell/sandbox.py +++ b/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, diff --git a/python/openshell/sandbox_cleanup_test.py b/python/openshell/sandbox_cleanup_test.py index 7d38fadf3..4cf558251 100644 --- a/python/openshell/sandbox_cleanup_test.py +++ b/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