feat(e2b): support listing sandboxes by metadata

Author: zhongwen666

rock/admin/core/sandbox_table.py

@@ -5,8 +5,9 @@
 import datetime
 from typing import TYPE_CHECKING, Any
 
-from sqlalchemy import select
+from sqlalchemy import and_, select
 
+from rock.actions.sandbox.response import State
 from rock.admin.core.db_provider import DatabaseProvider, retry_on_disconnect
 from rock.admin.core.schema import SandboxRecord
 from rock.admin.metrics.decorator import monitor_metastore_operation
@@ -190,6 +191,27 @@ def _list_by_in_sync(
             result = session.execute(stmt)
             return [_merge_status_blob(r.to_dict()) for r in result.scalars().all()]
 
+    @retry_on_disconnect
+    @monitor_metastore_operation
+    async def list_by_metadata(self, metadata: dict[str, str]) -> list[dict]:
+        """Return running sandboxes whose persisted metadata contains all requested pairs."""
+        return await self._db.run(self._list_by_metadata_sync, metadata)
+
+    def _list_by_metadata_sync(self, metadata: dict[str, str]) -> list[dict]:
+        if not metadata:
+            raise ValueError("metadata filter must not be empty")
+
+        stmt = (
+            select(SandboxRecord)
+            .where(SandboxRecord.state == State.RUNNING.value)
+            .where(_metadata_filter_expression(metadata, dialect_name=self._db.engine.dialect.name))
+            .order_by(SandboxRecord.sandbox_id.asc())
+        )
+
+        with self._db.session_factory() as session:
+            result = session.execute(stmt)
+            return [_merge_status_blob(record.to_dict()) for record in result.scalars().all()]
+
     @retry_on_disconnect
     @monitor_metastore_operation
     async def list_expired_by(
@@ -235,6 +257,12 @@ def _pick_columns(data: dict[str, Any]) -> dict[str, Any]:
     return result
 
 
+def _metadata_filter_expression(metadata: dict[str, str], dialect_name: str):
+    if dialect_name == "postgresql":
+        return SandboxRecord.labels.op("@>")(metadata)
+    return and_(*(SandboxRecord.labels[key].as_string() == value for key, value in metadata.items()))
+
+
 def _parse_aware_datetime(value: Any) -> datetime.datetime | None:
     if value is None:
         return None

rock/admin/entrypoints/e2b_proxy_api.py

@@ -1,11 +1,13 @@
-from fastapi import APIRouter, Request
+from typing import Annotated
+
+from fastapi import APIRouter, Query, Request
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse
 from fastapi.routing import APIRoute
 
-from rock.admin.proto.response import E2BSandboxDetail
+from rock.admin.proto.response import E2BListedSandbox, E2BSandboxDetail
+from rock.admin.service.e2b_proxy_service import E2BProxyService
 from rock.logger import init_logger
-from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService
 from rock.sdk.common.exceptions import BadRequestRockError, E2BSandboxNotFoundError
 
 logger = init_logger(__name__)
@@ -25,18 +27,21 @@ async def handler(request: Request):
                 return _error_response(400, message)
             except E2BSandboxNotFoundError as error:
                 return _error_response(404, str(error))
+            except BadRequestRockError as error:
+                logger.warning("E2B proxy request rejected: %s", error)
+                return _error_response(400, str(error))
             except Exception:
-                logger.exception("E2B get sandbox failed")
+                logger.exception("E2B proxy request failed")
                 return _error_response(500, "Internal server error")
 
         return handler
 
 
 e2b_proxy_router = APIRouter(route_class=E2BProxyAPIRoute)
-e2b_proxy_service: SandboxProxyService
+e2b_proxy_service: E2BProxyService
 
 
-def set_e2b_proxy_service(service: SandboxProxyService) -> None:
+def set_e2b_proxy_service(service: E2BProxyService) -> None:
     global e2b_proxy_service
     e2b_proxy_service = service
 
@@ -45,16 +50,22 @@ def _error_response(status_code: int, message: str) -> JSONResponse:
     return JSONResponse(status_code=status_code, content={"code": status_code, "message": message})
 
 
+@e2b_proxy_router.get(
+    "/v2/sandboxes",
+    response_model=list[E2BListedSandbox],
+    response_model_by_alias=True,
+    response_model_exclude_none=True,
+)
+async def list_sandboxes(
+    metadata: Annotated[str, Query(min_length=1)],
+) -> list[E2BListedSandbox]:
+    return await e2b_proxy_service.list_sandboxes(metadata)
+
+
 @e2b_proxy_router.get(
     "/sandboxes/{sandboxID}",
     response_model=E2BSandboxDetail,
     response_model_by_alias=True,
 )
 async def get_sandbox(sandboxID: str) -> E2BSandboxDetail:
-    try:
-        sandbox_status = await e2b_proxy_service.get_status(sandboxID, include_all_states=True)
-    except BadRequestRockError as error:
-        if str(error) == f"Sandbox {sandboxID} not found":
-            raise E2BSandboxNotFoundError(str(error)) from None
-        raise
-    return E2BSandboxDetail.from_sandbox_status(sandboxID, sandbox_status)
+    return await e2b_proxy_service.get_sandbox(sandboxID)

rock/admin/main.py

@@ -45,6 +45,7 @@
 from rock.admin.scheduler.tasks.sandbox_log_archive_task import (
     set_sandbox_table_provider as set_archive_sandbox_table_provider,
 )
+from rock.admin.service.e2b_proxy_service import E2BProxyService
 from rock.admin.service.ops_service import OpsService
 from rock.common.exception import request_validation_exception_handler
 from rock.config import DatabaseConfig, RockConfig, SchedulerConfig
@@ -263,7 +264,7 @@ async def lifespan(app: FastAPI):
 
     else:
         sandbox_manager = create_sandbox_proxy_service(rock_config=rock_config, meta_store=meta_store)
-        set_e2b_proxy_service(sandbox_manager)
+        set_e2b_proxy_service(E2BProxyService(sandbox_service=sandbox_manager, meta_store=meta_store))
         set_sandbox_proxy_service(sandbox_manager)
         proxy_service_ref = sandbox_manager
 

rock/admin/proto/response.py

@@ -1,6 +1,3 @@
-import datetime
-import math
-from ipaddress import ip_address
 from typing import Literal
 
 from pydantic import BaseModel, ConfigDict, Field
@@ -9,10 +6,7 @@
 from rock.actions.sandbox.response import State, StateTransitionRecord
 from rock.actions.sandbox.sandbox_info import SandboxInfo
 from rock.admin.proto.request import TaskSetSpec
-from rock.common.constants import E2B_CLIENT_ID, E2B_ENVD_VERSION, E2B_SANDBOX_IP_METADATA_KEY, E2B_STATE_BY_ROCK_STATE
 from rock.sandbox.utils.timeout import SandboxTimeoutHelper
-from rock.sdk.common.exceptions import E2BSandboxNotFoundError
-from rock.utils.format import parse_size_to_bytes
 
 
 class E2BCreateSandboxResponse(BaseModel):
@@ -22,6 +16,14 @@ class E2BCreateSandboxResponse(BaseModel):
     template_id: str = Field(alias="templateID")
 
 
+class E2BListedSandbox(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+
+    sandbox_id: str = Field(alias="sandboxID")
+    metadata: dict[str, str]
+    state: Literal["running", "paused"]
+
+
 class SandboxStartResponse(SandboxResponse):
     sandbox_id: str | None = None
     host_name: str | None = None
@@ -115,68 +117,6 @@ class E2BSandboxDetail(BaseModel):
     started_at: str = Field(alias="startedAt")
     end_at: str = Field(alias="endAt")
 
-    @staticmethod
-    def _state(sandbox_id: str, state: State | str | None) -> Literal["running", "paused"]:
-        try:
-            rock_state = state if isinstance(state, State) else State(state)
-            return E2B_STATE_BY_ROCK_STATE[rock_state.value]
-        except (KeyError, TypeError, ValueError):
-            raise E2BSandboxNotFoundError(f"Sandbox {sandbox_id} not found") from None
-
-    @staticmethod
-    def _iso8601_timestamp(sandbox_id: str, field: str, value: object) -> str:
-        if not isinstance(value, str):
-            raise ValueError(f"Sandbox {sandbox_id} {field} is invalid")
-        try:
-            parsed = datetime.datetime.fromisoformat(value.replace("Z", "+00:00"))
-        except ValueError:
-            raise ValueError(f"Sandbox {sandbox_id} {field} is invalid") from None
-        if parsed.tzinfo is None:
-            raise ValueError(f"Sandbox {sandbox_id} {field} must include a timezone")
-        return parsed.isoformat(timespec="seconds")
-
-    @classmethod
-    def from_sandbox_status(
-        cls,
-        sandbox_id: str,
-        sandbox_status: SandboxStatusResponse,
-    ) -> "E2BSandboxDetail":
-        state = cls._state(sandbox_id, sandbox_status.state)
-        end_at = (
-            sandbox_status.auto_stop_time
-            if state == "running"
-            else sandbox_status.auto_delete_time or sandbox_status.archive_time
-        )
-
-        metadata = sandbox_status.metadata
-        if not isinstance(metadata, dict) or not all(
-            isinstance(key, str) and isinstance(value, str) for key, value in metadata.items()
-        ):
-            raise ValueError(f"Sandbox {sandbox_id} metadata is invalid")
-
-        host_ip = sandbox_status.host_ip
-        if not isinstance(host_ip, str) or not host_ip.strip():
-            raise ValueError(f"Sandbox {sandbox_id} IP is missing")
-        ip_address(host_ip)
-
-        return cls(
-            sandboxID=sandbox_id,
-            metadata={**metadata, E2B_SANDBOX_IP_METADATA_KEY: host_ip},
-            state=state,
-            clientID=E2B_CLIENT_ID,
-            templateID=str(sandbox_status.image),
-            envdVersion=E2B_ENVD_VERSION,
-            cpuCount=max(1, math.ceil(float(sandbox_status.cpus))),
-            memoryMB=parse_size_to_bytes(str(sandbox_status.memory)) // (1024**2),
-            diskSizeMB=parse_size_to_bytes(str(sandbox_status.disk)) // (1024**2),
-            startedAt=cls._iso8601_timestamp(
-                sandbox_id,
-                "start time",
-                sandbox_status.start_time or sandbox_status.create_time,
-            ),
-            endAt=cls._iso8601_timestamp(sandbox_id, "end time", end_at),
-        )
-
 
 class SandboxListStatusResponse(SandboxStatusResponse):
     rock_authorization_encrypted: str | None = None