feat(sdk): add close_session method to Sandbox clientClose session

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()

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()

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())