feat(sdk): implement close_session method in Sandbox client #449

rock/sdk/sandbox/client.py

@@ -882,8 +882,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,26 @@
+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(sandbox_instance: Sandbox):
+    session_name = "test-close-session"
+    await sandbox_instance.create_session(CreateBashSessionRequest(session=session_name, session_type="bash"))
+
+    obs = await sandbox_instance.arun(cmd="echo 'alive'", session=session_name)
+    assert "alive" in obs.output
+
+    resp = await sandbox_instance.close_session(CloseBashSessionRequest(session=session_name, session_type="bash"))
+    assert resp is not None
+
+    with pytest.raises(Exception):
+        await sandbox_instance.arun(cmd="echo 'should fail'", session=session_name)