fix: eliminate N+1 performance bug in RemoteSandboxService with batch endpoint (#12105)

Co-authored-by: openhands <openhands@all-hands.dev>

tests/unit/app_server/test_remote_sandbox_service.py

@@ -331,7 +331,7 @@ class TestSandboxInfoConversion:
         runtime_data = create_runtime_data(status='running', pod_status='ready')
 
         # Execute
-        sandbox_info = await remote_sandbox_service._to_sandbox_info(
+        sandbox_info = remote_sandbox_service._to_sandbox_info(
             stored_sandbox, runtime_data
         )
 
@@ -358,7 +358,7 @@ class TestSandboxInfoConversion:
         runtime_data = create_runtime_data(status='running', pod_status='pending')
 
         # Execute
-        sandbox_info = await remote_sandbox_service._to_sandbox_info(
+        sandbox_info = remote_sandbox_service._to_sandbox_info(
             stored_sandbox, runtime_data
         )
 
@@ -367,23 +367,6 @@ class TestSandboxInfoConversion:
         assert sandbox_info.session_api_key == 'test-session-key'
         assert sandbox_info.exposed_urls is None
 
-    @pytest.mark.asyncio
-    async def test_to_sandbox_info_without_runtime(self, remote_sandbox_service):
-        """Test conversion to SandboxInfo without runtime data."""
-        # Setup
-        stored_sandbox = create_stored_sandbox()
-        remote_sandbox_service._get_runtime = AsyncMock(
-            side_effect=Exception('Runtime not found')
-        )
-
-        # Execute
-        sandbox_info = await remote_sandbox_service._to_sandbox_info(stored_sandbox)
-
-        # Verify
-        assert sandbox_info.status == SandboxStatus.MISSING
-        assert sandbox_info.session_api_key is None
-        assert sandbox_info.exposed_urls is None
-
     @pytest.mark.asyncio
     async def test_to_sandbox_info_loads_runtime_when_none_provided(
         self, remote_sandbox_service
@@ -391,15 +374,12 @@ class TestSandboxInfoConversion:
         """Test that runtime data is loaded when not provided."""
         # Setup
         stored_sandbox = create_stored_sandbox()
-        runtime_data = create_runtime_data()
-        remote_sandbox_service._get_runtime = AsyncMock(return_value=runtime_data)
 
         # Execute
-        sandbox_info = await remote_sandbox_service._to_sandbox_info(stored_sandbox)
+        sandbox_info = remote_sandbox_service._to_sandbox_info(stored_sandbox, None)
 
         # Verify
-        remote_sandbox_service._get_runtime.assert_called_once_with('test-sandbox-123')
-        assert sandbox_info.status == SandboxStatus.RUNNING
+        assert sandbox_info.status == SandboxStatus.MISSING
 
 
 class TestSandboxLifecycle:
@@ -677,15 +657,18 @@ class TestSandboxSearch:
         mock_result = MagicMock()
         mock_result.scalars.return_value = mock_scalars
         remote_sandbox_service.db_session.execute = AsyncMock(return_value=mock_result)
-        remote_sandbox_service._to_sandbox_info = AsyncMock(
-            side_effect=lambda stored: SandboxInfo(
-                id=stored.id,
-                created_by_user_id=stored.created_by_user_id,
-                sandbox_spec_id=stored.sandbox_spec_id,
-                status=SandboxStatus.RUNNING,
-                session_api_key='test-key',
-                created_at=stored.created_at,
-            )
+
+        # Mock the batch endpoint response
+        mock_batch_response = MagicMock()
+        mock_batch_response.raise_for_status.return_value = None
+        mock_batch_response.json.return_value = {
+            'runtimes': [
+                create_runtime_data('sb1'),
+                create_runtime_data('sb2'),
+            ]
+        }
+        remote_sandbox_service.httpx_client.request = AsyncMock(
+            return_value=mock_batch_response
         )
 
         # Execute
@@ -697,6 +680,14 @@ class TestSandboxSearch:
         assert result.items[0].id == 'sb1'
         assert result.items[1].id == 'sb2'
 
+        # Verify that the batch endpoint was called
+        remote_sandbox_service.httpx_client.request.assert_called_once_with(
+            'GET',
+            'https://api.example.com/sessions/batch',
+            headers={'X-API-Key': 'test-api-key'},
+            params=[('ids', 'sb1'), ('ids', 'sb2')],
+        )
+
     @pytest.mark.asyncio
     async def test_search_sandboxes_with_pagination(self, remote_sandbox_service):
         """Test sandbox search with pagination."""
@@ -710,15 +701,15 @@ class TestSandboxSearch:
         mock_result = MagicMock()
         mock_result.scalars.return_value = mock_scalars
         remote_sandbox_service.db_session.execute = AsyncMock(return_value=mock_result)
-        remote_sandbox_service._to_sandbox_info = AsyncMock(
-            side_effect=lambda stored: SandboxInfo(
-                id=stored.id,
-                created_by_user_id=stored.created_by_user_id,
-                sandbox_spec_id=stored.sandbox_spec_id,
-                status=SandboxStatus.RUNNING,
-                session_api_key='test-key',
-                created_at=stored.created_at,
-            )
+
+        # Mock the batch endpoint response
+        mock_batch_response = MagicMock()
+        mock_batch_response.raise_for_status.return_value = None
+        mock_batch_response.json.return_value = {
+            'runtimes': [create_runtime_data(f'sb{i}') for i in range(6)]
+        }
+        remote_sandbox_service.httpx_client.request = AsyncMock(
+            return_value=mock_batch_response
         )
 
         # Execute
@@ -739,15 +730,15 @@ class TestSandboxSearch:
         mock_result = MagicMock()
         mock_result.scalars.return_value = mock_scalars
         remote_sandbox_service.db_session.execute = AsyncMock(return_value=mock_result)
-        remote_sandbox_service._to_sandbox_info = AsyncMock(
-            side_effect=lambda stored: SandboxInfo(
-                id=stored.id,
-                created_by_user_id=stored.created_by_user_id,
-                sandbox_spec_id=stored.sandbox_spec_id,
-                status=SandboxStatus.RUNNING,
-                session_api_key='test-key',
-                created_at=stored.created_at,
-            )
+
+        # Mock the batch endpoint response
+        mock_batch_response = MagicMock()
+        mock_batch_response.raise_for_status.return_value = None
+        mock_batch_response.json.return_value = {
+            'runtimes': [create_runtime_data('sb1')]
+        }
+        remote_sandbox_service.httpx_client.request = AsyncMock(
+            return_value=mock_batch_response
         )
 
         # Execute
@@ -757,6 +748,80 @@ class TestSandboxSearch:
         # Note: We can't easily verify the exact SQL query, but we can verify the method was called
         remote_sandbox_service.db_session.execute.assert_called_once()
 
+    @pytest.mark.asyncio
+    async def test_get_runtimes_batch_success(self, remote_sandbox_service):
+        """Test successful batch runtime retrieval."""
+        # Setup
+        sandbox_ids = ['sb1', 'sb2', 'sb3']
+        mock_response = MagicMock()
+        mock_response.raise_for_status.return_value = None
+        mock_response.json.return_value = {
+            'runtimes': [
+                create_runtime_data('sb1'),
+                create_runtime_data('sb2'),
+                create_runtime_data('sb3'),
+            ]
+        }
+        remote_sandbox_service.httpx_client.request = AsyncMock(
+            return_value=mock_response
+        )
+
+        # Execute
+        result = await remote_sandbox_service._get_runtimes_batch(sandbox_ids)
+
+        # Verify
+        assert len(result) == 3
+        assert 'sb1' in result
+        assert 'sb2' in result
+        assert 'sb3' in result
+        assert result['sb1']['session_id'] == 'sb1'
+
+        # Verify the correct API call was made
+        remote_sandbox_service.httpx_client.request.assert_called_once_with(
+            'GET',
+            'https://api.example.com/sessions/batch',
+            headers={'X-API-Key': 'test-api-key'},
+            params=[('ids', 'sb1'), ('ids', 'sb2'), ('ids', 'sb3')],
+        )
+
+    @pytest.mark.asyncio
+    async def test_get_runtimes_batch_empty_list(self, remote_sandbox_service):
+        """Test batch runtime retrieval with empty sandbox list."""
+        # Execute
+        result = await remote_sandbox_service._get_runtimes_batch([])
+
+        # Verify
+        assert result == {}
+        # Verify no API call was made
+        remote_sandbox_service.httpx_client.request.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_get_runtimes_batch_partial_results(self, remote_sandbox_service):
+        """Test batch runtime retrieval with partial results (some sandboxes not found)."""
+        # Setup
+        sandbox_ids = ['sb1', 'sb2', 'sb3']
+        mock_response = MagicMock()
+        mock_response.raise_for_status.return_value = None
+        mock_response.json.return_value = {
+            'runtimes': [
+                create_runtime_data('sb1'),
+                create_runtime_data('sb3'),
+                # sb2 is missing from the response
+            ]
+        }
+        remote_sandbox_service.httpx_client.request = AsyncMock(
+            return_value=mock_response
+        )
+
+        # Execute
+        result = await remote_sandbox_service._get_runtimes_batch(sandbox_ids)
+
+        # Verify
+        assert len(result) == 2
+        assert 'sb1' in result
+        assert 'sb2' not in result  # Missing from response
+        assert 'sb3' in result
+
     @pytest.mark.asyncio
     async def test_get_sandbox_exists(self, remote_sandbox_service):
         """Test getting an existing sandbox."""
@@ -765,7 +830,7 @@ class TestSandboxSearch:
         remote_sandbox_service._get_stored_sandbox = AsyncMock(
             return_value=stored_sandbox
         )
-        remote_sandbox_service._to_sandbox_info = AsyncMock(
+        remote_sandbox_service._to_sandbox_info = MagicMock(
             return_value=SandboxInfo(
                 id='test-sandbox-123',
                 created_by_user_id='test-user-123',