From 43db9a125476023fa86a411d32f55084e126b6a9 Mon Sep 17 00:00:00 2001 From: bird2426 Date: Sat, 7 Feb 2026 20:56:41 +0800 Subject: [PATCH] feat(sdk): implement close_session method in Sandbox client Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-opencode) Co-authored-by: Sisyphus --- rock/sdk/sandbox/client.py | 18 ++- .../sdk/sandbox/test_close_session.py | 53 ++++++++ tests/unit/test_close_session_logic.py | 113 ++++++++++++++++++ 3 files changed, 182 insertions(+), 2 deletions(-) create mode 100644 tests/integration/sdk/sandbox/test_close_session.py create mode 100644 tests/unit/test_close_session_logic.py diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index 3278b46e..7c84087d 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -874,8 +874,22 @@ def _generate_utc_iso_time(self): return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") async def close_session(self, request: CloseSessionRequest) -> CloseSessionResponse: - # TODO: implement this - pass + url = f"{self._url}/close_session" + headers = self._build_headers() + data = { + "sandbox_id": self.sandbox_id, + **request.model_dump(), + } + try: + response = await HttpUtils.post(url, headers, data) + except Exception as e: + raise Exception(f"Failed to close session: {str(e)}, post url {url}") + + logging.debug(f"Close session response: {response}") + if "Success" != response.get("status"): + raise Exception(f"Failed to close session: {response}") + result: dict = response.get("result") # type: ignore + return CloseSessionResponse(**result) async def close(self) -> CloseResponse: await self.stop() diff --git a/tests/integration/sdk/sandbox/test_close_session.py b/tests/integration/sdk/sandbox/test_close_session.py new file mode 100644 index 00000000..650a5545 --- /dev/null +++ b/tests/integration/sdk/sandbox/test_close_session.py @@ -0,0 +1,53 @@ +import pytest +import os +from pathlib import Path +from rock.sdk.sandbox.client import Sandbox +from rock.actions import CreateBashSessionRequest, CloseBashSessionRequest, BashAction +from tests.integration.conftest import SKIP_IF_NO_DOCKER + +# Set writable status directory for sandbox deployment +os.environ["ROCK_SERVICE_STATUS_DIR"] = "/tmp/rock_status" + + +@pytest.mark.need_admin +@SKIP_IF_NO_DOCKER +@pytest.mark.asyncio +async def test_sdk_close_session(admin_remote_server): + # Use low resource config to avoid Ray scheduling issues on local machine + config_params = {"image": "python:3.11", "memory": "512m", "cpus": 0.5} + + # We can't easily pass params to sandbox_instance fixture if we want to customize it heavily, + # so we'll create our own sandbox here using the server provided by the fixture. + from rock.sdk.sandbox.config import SandboxConfig + + config = SandboxConfig( + image=config_params["image"], + memory=config_params["memory"], + cpus=config_params["cpus"], + base_url=f"{admin_remote_server.endpoint}:{admin_remote_server.port}", + ) + + sandbox = Sandbox(config) + try: + # 1. Start + await sandbox.start() + + # 2. Create session + session_name = "test-close-session" + await sandbox.create_session(CreateBashSessionRequest(session=session_name, session_type="bash")) + + # 3. Verify alive (use arun instead of deprecated run_in_session) + obs = await sandbox.arun(cmd="echo 'alive'", session=session_name) + assert "alive" in obs.output + + # 4. Close session + resp = await sandbox.close_session(CloseBashSessionRequest(session=session_name, session_type="bash")) + assert resp is not None + + # 5. Verify closed (expecting error) + with pytest.raises(Exception): + await sandbox.arun(cmd="echo 'should fail'", session=session_name) + + finally: + if sandbox.sandbox_id: + await sandbox.stop() diff --git a/tests/unit/test_close_session_logic.py b/tests/unit/test_close_session_logic.py new file mode 100644 index 00000000..efa1a980 --- /dev/null +++ b/tests/unit/test_close_session_logic.py @@ -0,0 +1,113 @@ +""" +Unit test for close_session - tests the core logic without Docker +""" +import pytest +from rock.rocklet.local_sandbox import LocalSandboxRuntime, BashSession +from rock.admin.proto.request import SandboxBashAction as BashAction +from rock.admin.proto.request import SandboxCloseBashSessionRequest as CloseBashSessionRequest +from rock.admin.proto.request import SandboxCreateBashSessionRequest as CreateBashSessionRequest + + +@pytest.mark.asyncio +async def test_close_session_logic(): + """Test close_session logic directly with LocalSandboxRuntime""" + + # Create runtime (no extra parameters allowed) + runtime = LocalSandboxRuntime() + + # 1. Create a session + session_name = "test-session" + create_req = CreateBashSessionRequest(session=session_name, session_type="bash") + create_resp = await runtime.create_session(create_req) + assert create_resp.session_type == "bash" + assert session_name in runtime.sessions + print(f"✅ Session '{session_name}' created successfully") + + # 2. Run a command in the session + action = BashAction(command="echo 'hello'", session=session_name) + obs = await runtime.run_in_session(action) + assert "hello" in obs.output + print(f"✅ Command executed in session: {obs.output.strip()}") + + # 3. Close the session + close_req = CloseBashSessionRequest(session=session_name, session_type="bash") + close_resp = await runtime.close_session(close_req) + assert close_resp.session_type == "bash" + assert session_name not in runtime.sessions + print(f"✅ Session '{session_name}' closed successfully") + + # 4. Verify session is closed (should raise exception) + with pytest.raises(Exception) as exc_info: + await runtime.run_in_session(action) + + assert "does not exist" in str(exc_info.value) + print(f"✅ Verified: Closed session cannot be used") + + print("\n" + "=" * 60) + print("✅ All unit tests passed!") + print("=" * 60) + + +@pytest.mark.asyncio +async def test_close_nonexistent_session(): + """Test closing a session that doesn't exist raises SessionDoesNotExistError""" + from rock.rocklet.local_sandbox import SessionDoesNotExistError + + runtime = LocalSandboxRuntime() + + # Try to close a session that was never created + close_req = CloseBashSessionRequest(session="nonexistent-session", session_type="bash") + + with pytest.raises(SessionDoesNotExistError) as exc_info: + await runtime.close_session(close_req) + + assert "nonexistent-session" in str(exc_info.value) + print("✅ Correctly raised SessionDoesNotExistError for non-existent session") + + +@pytest.mark.asyncio +async def test_close_session_twice(): + """Test that closing a session twice raises an error""" + from rock.rocklet.local_sandbox import SessionDoesNotExistError + + runtime = LocalSandboxRuntime() + + # Create a session + session_name = "test-double-close" + create_req = CreateBashSessionRequest(session=session_name, session_type="bash") + await runtime.create_session(create_req) + + # Close the session first time - should succeed + close_req = CloseBashSessionRequest(session=session_name, session_type="bash") + await runtime.close_session(close_req) + + # Close the session second time - should raise error + with pytest.raises(SessionDoesNotExistError) as exc_info: + await runtime.close_session(close_req) + + assert session_name in str(exc_info.value) + print("✅ Correctly raised SessionDoesNotExistError when closing session twice") + + +@pytest.mark.asyncio +async def test_close_session_with_default_name(): + """Test closing the default session""" + runtime = LocalSandboxRuntime() + + # Create a session with default name + create_req = CreateBashSessionRequest(session_type="bash") # uses default session name + await runtime.create_session(create_req) + assert "default" in runtime.sessions + + # Close the default session + close_req = CloseBashSessionRequest(session_type="bash") # uses default session name + close_resp = await runtime.close_session(close_req) + + assert close_resp.session_type == "bash" + assert "default" not in runtime.sessions + print("✅ Successfully closed default session") + + +if __name__ == "__main__": + import asyncio + asyncio.run(test_close_session_logic())