diff --git a/rock/actions/sandbox/response.py b/rock/actions/sandbox/response.py index baa087fba8..937a52efa9 100644 --- a/rock/actions/sandbox/response.py +++ b/rock/actions/sandbox/response.py @@ -51,6 +51,7 @@ class SandboxStatusResponse(BaseModel): host_ip: str | None = None is_alive: bool = True image: str | None = None + metadata: dict[str, str] | None = None gateway_version: str | None = None swe_rex_version: str | None = None user_id: str | None = None diff --git a/rock/admin/core/sandbox_table.py b/rock/admin/core/sandbox_table.py index 1e26bae72e..bf47014032 100644 --- a/rock/admin/core/sandbox_table.py +++ b/rock/admin/core/sandbox_table.py @@ -254,7 +254,11 @@ def _merge_status_blob(raw: dict[str, Any]) -> dict[str, Any]: The ``status`` column stores the full SandboxInfo snapshot. Fields that only exist in the blob (e.g. ``state_history``) become top-level keys. - Scalar columns take priority over blob values. + Scalar columns take priority over blob values. The database ``labels`` + column is exposed through the business-facing ``metadata`` field. """ status_blob = raw.pop("status", None) or {} - return {**status_blob, **raw} + result = {**status_blob, **raw} + if "labels" in result: + result["metadata"] = result["labels"] + return result diff --git a/rock/admin/entrypoints/e2b_api.py b/rock/admin/entrypoints/e2b_api.py index dedb1d6d03..37902f6501 100644 --- a/rock/admin/entrypoints/e2b_api.py +++ b/rock/admin/entrypoints/e2b_api.py @@ -1,13 +1,14 @@ import math -from typing import Any +from typing import Annotated -from fastapi import APIRouter, Request +from fastapi import APIRouter, Depends, Request from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse from fastapi.routing import APIRoute -from pydantic import BaseModel, ConfigDict, Field -from rock.common.validation import NonBlankStr +from rock.admin.proto.request import E2BCreateSandboxRequest, StartHeaders +from rock.admin.proto.response import E2BCreateSandboxResponse +from rock.common.constants import AP_SANDBOX_ID_METADATA_KEY, E2B_CLIENT_ID, E2B_ENVD_VERSION from rock.deployments.config import DockerDeploymentConfig from rock.logger import init_logger from rock.sandbox.sandbox_manager import SandboxManager @@ -16,26 +17,6 @@ logger = init_logger(__name__) -class E2BCreateSandboxRequest(BaseModel): - model_config = ConfigDict(extra="forbid", populate_by_name=True) - - template_id: NonBlankStr = Field(alias="templateID") - timeout: int = Field(gt=0, strict=True) - metadata: dict[str, str] - secure: bool | None = None - allow_internet_access: bool | None = None - env_vars: dict[str, str] = Field(default_factory=dict, alias="envVars") - auto_pause: bool | None = Field(default=None, alias="autoPause") - auto_resume: dict[str, Any] | None = Field(default=None, alias="autoResume") - - -class E2BCreateSandboxResponse(BaseModel): - sandbox_id: str = Field(alias="sandboxID") - envd_version: str = Field(alias="envdVersion") - client_id: str = Field(alias="clientID") - template_id: str = Field(alias="templateID") - - class E2BAPIRoute(APIRoute): def get_route_handler(self): route_handler = super().get_route_handler() @@ -79,19 +60,25 @@ def _error_response(status_code: int, message: str) -> JSONResponse: ) async def create_sandbox( request: E2BCreateSandboxRequest, + headers: Annotated[StartHeaders, Depends()], ) -> E2BCreateSandboxResponse: # ROCK stores lifecycle TTLs in whole minutes. Round up so an E2B timeout # never expires a sandbox earlier than the caller requested. config = DockerDeploymentConfig( image=request.template_id, auto_clear_time_minutes=math.ceil(request.timeout / 60), + container_name=request.metadata.get(AP_SANDBOX_ID_METADATA_KEY), metadata=request.metadata, env_vars=request.env_vars, ) - result = await e2b_sandbox_manager.start(config) + result = await e2b_sandbox_manager.start( + config, + user_info=headers.user_info, + cluster_info=headers.cluster_info, + ) return E2BCreateSandboxResponse( sandboxID=result.sandbox_id, - envdVersion="0.1.0", - clientID="rock", + envdVersion=E2B_ENVD_VERSION, + clientID=E2B_CLIENT_ID, templateID=request.template_id, ) diff --git a/rock/admin/entrypoints/e2b_proxy_api.py b/rock/admin/entrypoints/e2b_proxy_api.py new file mode 100644 index 0000000000..9b8e673a7d --- /dev/null +++ b/rock/admin/entrypoints/e2b_proxy_api.py @@ -0,0 +1,60 @@ +from fastapi import APIRouter, 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.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__) + + +class E2BProxyAPIRoute(APIRoute): + def get_route_handler(self): + route_handler = super().get_route_handler() + + async def handler(request: Request): + try: + return await route_handler(request) + except RequestValidationError as error: + message = "; ".join( + f"{'.'.join(str(part) for part in item['loc'])}: {item['msg']}" for item in error.errors() + ) + return _error_response(400, message) + except E2BSandboxNotFoundError as error: + return _error_response(404, str(error)) + except Exception: + logger.exception("E2B get sandbox failed") + return _error_response(500, "Internal server error") + + return handler + + +e2b_proxy_router = APIRouter(route_class=E2BProxyAPIRoute) +e2b_proxy_service: SandboxProxyService + + +def set_e2b_proxy_service(service: SandboxProxyService) -> None: + global e2b_proxy_service + e2b_proxy_service = service + + +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( + "/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) diff --git a/rock/admin/main.py b/rock/admin/main.py index 7dfb880aa8..9703781693 100644 --- a/rock/admin/main.py +++ b/rock/admin/main.py @@ -26,6 +26,7 @@ from rock.admin.core.template_table import TemplateTable from rock.admin.entrypoints.admin_ops_api import admin_ops_router, set_ops_service from rock.admin.entrypoints.e2b_api import e2b_router, set_e2b_sandbox_manager +from rock.admin.entrypoints.e2b_proxy_api import e2b_proxy_router, set_e2b_proxy_service from rock.admin.entrypoints.sandbox_api import sandbox_router, set_sandbox_manager from rock.admin.entrypoints.sandbox_proxy_api import sandbox_proxy_router, set_sandbox_proxy_service from rock.admin.entrypoints.warmup_api import set_warmup_service, warmup_router @@ -262,6 +263,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_sandbox_proxy_service(sandbox_manager) proxy_service_ref = sandbox_manager @@ -351,6 +353,7 @@ def _include_routers(app: FastAPI, role: str) -> None: app.include_router(sandbox_router, prefix="/apis/envs/sandbox/v1", tags=["sandbox"]) app.include_router(admin_ops_router, prefix="/apis/envs/sandbox/v1/ops", tags=["admin-ops"]) else: + app.include_router(e2b_proxy_router, tags=["e2b"]) app.include_router(sandbox_proxy_router, prefix="/apis/envs/sandbox/v1", tags=["sandbox"]) app.include_router(warmup_router, prefix="/apis/envs/sandbox/v1", tags=["warmup"]) app.include_router(gem_router, prefix="/apis/v1/envs/gem", tags=["gem"]) diff --git a/rock/admin/proto/request.py b/rock/admin/proto/request.py index a1a50fdce1..dce71526d1 100644 --- a/rock/admin/proto/request.py +++ b/rock/admin/proto/request.py @@ -1,7 +1,7 @@ -from typing import Annotated, Literal, TypedDict +from typing import Annotated, Any, Literal, TypedDict from fastapi import Header -from pydantic import BaseModel, Field, field_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator from rock import env_vars from rock.actions import ( @@ -12,9 +12,23 @@ ReadFileRequest, WriteFileRequest, ) +from rock.common.constants import BEARER_AUTHORIZATION_PREFIX from rock.common.validation import NonBlankStr +class E2BCreateSandboxRequest(BaseModel): + model_config = ConfigDict(extra="forbid", populate_by_name=True) + + template_id: NonBlankStr = Field(alias="templateID") + timeout: int = Field(gt=0, strict=True) + metadata: dict[str, str] + secure: bool | None = None + allow_internet_access: bool | None = None + env_vars: dict[str, str] = Field(default_factory=dict, alias="envVars") + auto_pause: bool | None = Field(default=None, alias="autoPause") + auto_resume: dict[str, Any] | None = Field(default=None, alias="autoResume") + + class SandboxStartRequest(BaseModel): image: NonBlankStr """image""" @@ -190,9 +204,12 @@ def __init__( x_user_id: str | None = Header(default="default", alias="X-User-Id"), x_experiment_id: str | None = Header(default="default", alias="X-Experiment-Id"), rock_authorization: str | None = Header(default="default", alias="X-Key"), + x_api_key: str | None = Header(default=None, alias="X-API-Key"), x_namespace: str | None = Header(default="default", alias="X-Namespace"), x_cluster: str | None = Header(default="default", alias="X-Cluster"), ): + if x_api_key is not None: + rock_authorization = f"{BEARER_AUTHORIZATION_PREFIX}{x_api_key}" self.user_info: UserInfo = { "user_id": x_user_id, "experiment_id": x_experiment_id, diff --git a/rock/admin/proto/response.py b/rock/admin/proto/response.py index 9d0401ec0e..c954d558a1 100644 --- a/rock/admin/proto/response.py +++ b/rock/admin/proto/response.py @@ -1,10 +1,25 @@ -from pydantic import BaseModel, Field +import datetime +import math +from ipaddress import ip_address +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field from rock.actions import SandboxResponse 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): + sandbox_id: str = Field(alias="sandboxID") + envd_version: str = Field(alias="envdVersion") + client_id: str = Field(alias="clientID") + template_id: str = Field(alias="templateID") class SandboxStartResponse(SandboxResponse): @@ -27,6 +42,7 @@ class SandboxStatusResponse(BaseModel): host_ip: str | None = None is_alive: bool = True image: str | None = None + metadata: dict[str, str] | None = None gateway_version: str | None = None swe_rex_version: str | None = None user_id: str | None = None @@ -62,6 +78,7 @@ def from_sandbox_info(cls, sandbox_info: "SandboxInfo") -> "SandboxStatusRespons host_ip=sandbox_info.get("host_ip"), host_name=sandbox_info.get("host_name"), image=sandbox_info.get("image"), + metadata=sandbox_info.get("metadata"), user_id=sandbox_info.get("user_id"), experiment_id=sandbox_info.get("experiment_id"), namespace=sandbox_info.get("namespace"), @@ -83,6 +100,84 @@ def from_sandbox_info(cls, sandbox_info: "SandboxInfo") -> "SandboxStatusRespons ) +class E2BSandboxDetail(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + sandbox_id: str = Field(alias="sandboxID") + metadata: dict[str, str] + state: Literal["running", "paused"] + client_id: str = Field(alias="clientID") + template_id: str = Field(alias="templateID") + envd_version: str = Field(alias="envdVersion") + cpu_count: int = Field(alias="cpuCount") + memory_mb: int = Field(alias="memoryMB") + disk_size_mb: int = Field(alias="diskSizeMB") + 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 diff --git a/rock/common/constants.py b/rock/common/constants.py index b835de8988..c256611780 100644 --- a/rock/common/constants.py +++ b/rock/common/constants.py @@ -1,4 +1,5 @@ from enum import Enum +from typing import Literal KATA_RUNTIME_SWITCH = "use_kata_enabled" SUPPORT_KATA_SWITCH = "support_kata_enabled" @@ -11,6 +12,15 @@ PID_PREFIX = "PIDSTART" PID_SUFFIX = "PIDEND" SCHEDULER_LOG_NAME = "scheduler.log" +BEARER_AUTHORIZATION_PREFIX = "Bearer " +AP_SANDBOX_ID_METADATA_KEY = "ap-sandbox-id" +E2B_CLIENT_ID = "rock" +E2B_ENVD_VERSION = "0.1.0" +E2B_SANDBOX_IP_METADATA_KEY = "e2b.agents.kruise.io/sandbox-ip" +E2B_STATE_BY_ROCK_STATE: dict[str, Literal["running", "paused"]] = { + "running": "running", + "archived": "paused", +} class DeploymentHookStep(str, Enum): diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index 8bdae60792..49ce48f568 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -282,8 +282,13 @@ async def restart_async(self, sandbox_id: str) -> SandboxStartResponse: ) @monitor_sandbox_operation() - async def start(self, config: DeploymentConfig) -> SandboxStartResponse: - response = await self.start_async(config) + async def start( + self, + config: DeploymentConfig, + user_info: UserInfo = {}, + cluster_info: ClusterInfo = {}, + ) -> SandboxStartResponse: + response = await self.start_async(config, user_info=user_info, cluster_info=cluster_info) sandbox_id = response.sandbox_id deadline = time.time() + REQUEST_TIMEOUT_SECONDS with StageTimer("startup_timing", f"[{sandbox_id}] Wait sandbox running", logger): @@ -472,6 +477,7 @@ async def get_status(self, sandbox_id, include_all_states: bool = False) -> Sand host_ip=sandbox_info.get("host_ip"), is_alive=is_alive, image=sandbox_info.get("image"), + metadata=sandbox_info.get("metadata"), swe_rex_version=swe_version, gateway_version=gateway_version, user_id=sandbox_info.get("user_id"), diff --git a/rock/sandbox/service/opensandbox_proxy_service.py b/rock/sandbox/service/opensandbox_proxy_service.py index 20a304e1cd..f85e819479 100644 --- a/rock/sandbox/service/opensandbox_proxy_service.py +++ b/rock/sandbox/service/opensandbox_proxy_service.py @@ -34,6 +34,7 @@ from rock.sandbox.sandbox_meta_store import SandboxMetaStore from rock.sandbox.service.backends.opensandbox import OpenSandboxBackend from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService +from rock.sandbox.utils.timeout import SandboxTimeoutHelper from rock.sdk.common.exceptions import BadRequestRockError OPENSANDBOX_BACKEND = "opensandbox" @@ -287,6 +288,12 @@ async def get_status(self, sandbox_id: str, include_all_states: bool = False) -> sandbox_info = sm.sandbox_info info = operator_sandbox_info if operator_sandbox_info is not None else sandbox_info + timeout_info = await self._meta_store.get_timeout(sandbox_id) + auto_stop_time, auto_archive_time, auto_delete_time = SandboxTimeoutHelper.auto_transition_times_for_status( + info.get("state"), + info, + timeout_info, + ) return SandboxStatusResponse( sandbox_id=sandbox_id, status=info.get("phases"), @@ -296,6 +303,7 @@ async def get_status(self, sandbox_id: str, include_all_states: bool = False) -> host_ip=info.get("host_ip"), is_alive=is_alive, image=info.get("image"), + metadata=info.get("metadata"), swe_rex_version=None, gateway_version=gateway_version, user_id=info.get("user_id"), @@ -310,6 +318,10 @@ async def get_status(self, sandbox_id: str, include_all_states: bool = False) -> start_time=info.get("start_time"), stop_time=info.get("stop_time"), create_time=info.get("create_time"), + archive_time=info.get("archive_time"), delete_time=info.get("delete_time"), + auto_stop_time=auto_stop_time, + auto_archive_time=auto_archive_time, + auto_delete_time=auto_delete_time, state_history=sandbox_info.get("state_history", []), ) diff --git a/rock/sandbox/service/sandbox_proxy_service.py b/rock/sandbox/service/sandbox_proxy_service.py index 851111dcec..a7b1ba839e 100644 --- a/rock/sandbox/service/sandbox_proxy_service.py +++ b/rock/sandbox/service/sandbox_proxy_service.py @@ -1183,6 +1183,7 @@ async def get_status(self, sandbox_id: str, include_all_states: bool = False) -> host_ip=host_ip, is_alive=is_alive, image=info.get("image"), + metadata=info.get("metadata"), swe_rex_version=swe_version, gateway_version=gateway_version, user_id=info.get("user_id"), diff --git a/rock/sdk/common/exceptions.py b/rock/sdk/common/exceptions.py index 20a6c98579..2618ef2903 100644 --- a/rock/sdk/common/exceptions.py +++ b/rock/sdk/common/exceptions.py @@ -9,6 +9,10 @@ def __init__(self, code: CommitErrorCode, message: str): self.code = code +class E2BSandboxNotFoundError(Exception): + pass + + class RockException(Exception): _code: codes = None diff --git a/sql/migrations/20260810_add_sandbox_labels.sql b/sql/migrations/20260810_add_sandbox_labels.sql deleted file mode 100644 index d9a8a799ce..0000000000 --- a/sql/migrations/20260810_add_sandbox_labels.sql +++ /dev/null @@ -1,6 +0,0 @@ -ALTER TABLE sandbox_record -ADD COLUMN IF NOT EXISTS labels JSONB NOT NULL DEFAULT '{}'::jsonb; - -CREATE INDEX CONCURRENTLY IF NOT EXISTS ix_sandbox_record_labels_gin -ON sandbox_record -USING GIN (labels jsonb_path_ops); diff --git a/sql/migrations/20260810_create_template.sql b/sql/migrations/20260810_create_template.sql deleted file mode 100644 index a16b6c2fd9..0000000000 --- a/sql/migrations/20260810_create_template.sql +++ /dev/null @@ -1,17 +0,0 @@ -CREATE TABLE IF NOT EXISTS template ( - template_id VARCHAR(128) NOT NULL, - os_type VARCHAR(32) NOT NULL, - spec JSONB, - status VARCHAR(32) NOT NULL, - current_step VARCHAR(32), - artifact_uri VARCHAR(1024), - fiber_pool_id VARCHAR(128), - execution_context JSONB, - error_code VARCHAR(128), - error_message TEXT, - created_at TIMESTAMPTZ NOT NULL, - updated_at TIMESTAMPTZ NOT NULL, - PRIMARY KEY (template_id) -); - -CREATE INDEX IF NOT EXISTS ix_template_status ON template (status); diff --git a/tests/integration/admin/test_e2b_create_sandbox.py b/tests/integration/admin/test_e2b_create_sandbox.py index 896e5db83c..ed29e21618 100644 --- a/tests/integration/admin/test_e2b_create_sandbox.py +++ b/tests/integration/admin/test_e2b_create_sandbox.py @@ -2,8 +2,7 @@ import os import pytest -from ap_sandbox import SandboxConfig -from ap_sandbox.client import create_sandbox +from ap_sandbox import SandboxConfig, SandboxManager TEMPLATE_ID = "st_5efa0210685646998525" @@ -27,10 +26,11 @@ def test_create_sandbox_with_ap_sandbox_sdk(caplog: pytest.LogCaptureFixture): ap_sandbox_metadata={"ap-job-id": "e2e-create-sandbox-test"}, ) - sandbox = create_sandbox( - config, + sandbox_manager = SandboxManager(config) + sandbox_manager.create( api_url=f"https://{domain}", validate_api_key=False, ) - assert sandbox.sandbox_id + assert sandbox_manager.sandbox_id not in {None, "", ""} + assert sandbox_manager.sandbox_ip diff --git a/tests/unit/admin/core/test_sandbox_table.py b/tests/unit/admin/core/test_sandbox_table.py index 75891141ed..07a463e3e4 100644 --- a/tests/unit/admin/core/test_sandbox_table.py +++ b/tests/unit/admin/core/test_sandbox_table.py @@ -57,6 +57,17 @@ async def test_metadata_is_persisted_as_labels(self, db): assert record["labels"] == metadata + async def test_labels_are_restored_as_metadata_on_read(self, db): + labels = {"legacy-key": "legacy-value"} + await db.create( + "sqlite-sbx-labels", + {"create_time": "2025-01-01T00:00:00Z", "labels": labels}, + ) + + record = await db.get("sqlite-sbx-labels") + + assert record["metadata"] == labels + async def test_insert_duplicate_raises(self, db): sandbox_id = "sqlite-sbx-002" data = {"state": "pending", "create_time": "2025-01-01T00:00:00Z"} diff --git a/tests/unit/admin/entrypoints/test_e2b_api.py b/tests/unit/admin/entrypoints/test_e2b_api.py index 04699641c3..963869c2c5 100644 --- a/tests/unit/admin/entrypoints/test_e2b_api.py +++ b/tests/unit/admin/entrypoints/test_e2b_api.py @@ -33,6 +33,7 @@ async def test_create_sandbox_returns_e2b_response_and_maps_request(e2b_app): "templateID": "linux-dind", "timeout": 3601, "metadata": { + "ap-sandbox-id": "ap-sandbox-123", "ap-job-id": "job-123", "ap-template": "swe-bench", "e2b.agents.kruise.io/return-sandbox-ip": "true", @@ -45,7 +46,18 @@ async def test_create_sandbox_returns_e2b_response_and_maps_request(e2b_app): } async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: - response = await client.post("/sandboxes", json=request_body) + response = await client.post( + "/sandboxes", + json=request_body, + headers={ + "X-User-Id": "user-123", + "X-Experiment-Id": "experiment-123", + "X-Namespace": "namespace-123", + "X-Cluster": "cluster-123", + "X-Key": "legacy-key", + "X-API-Key": "e2b-key", + }, + ) assert response.status_code == 201 assert response.json() == { @@ -60,7 +72,16 @@ async def test_create_sandbox_returns_e2b_response_and_maps_request(e2b_app): assert config.auto_clear_time_minutes == 61 assert config.metadata == request_body["metadata"] assert config.env_vars == request_body["envVars"] - assert manager.start.await_args.kwargs == {} + assert config.container_name == "ap-sandbox-123" + assert manager.start.await_args.kwargs == { + "user_info": { + "user_id": "user-123", + "experiment_id": "experiment-123", + "namespace": "namespace-123", + "rock_authorization": "Bearer e2b-key", + }, + "cluster_info": {"cluster_name": "cluster-123"}, + } @pytest.mark.parametrize("timeout", [0, True, "3600"]) diff --git a/tests/unit/admin/entrypoints/test_e2b_proxy_api.py b/tests/unit/admin/entrypoints/test_e2b_proxy_api.py new file mode 100644 index 0000000000..08cdf2dd45 --- /dev/null +++ b/tests/unit/admin/entrypoints/test_e2b_proxy_api.py @@ -0,0 +1,103 @@ +from unittest.mock import AsyncMock + +import pytest +from fastapi import FastAPI +from httpx import ASGITransport, AsyncClient + +from rock.actions.sandbox.response import State +from rock.admin.entrypoints.e2b_proxy_api import e2b_proxy_router, set_e2b_proxy_service +from rock.sandbox.sandbox_meta_store import SandboxMetaStore +from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService + + +@pytest.fixture +async def e2b_proxy_app(redis_provider, _memory_sandbox_table, rock_config, monkeypatch): + meta_store = SandboxMetaStore( + redis_provider=redis_provider, + sandbox_table=_memory_sandbox_table, + rock_config=rock_config, + ) + proxy_service = SandboxProxyService(rock_config=rock_config, meta_store=meta_store) + proxy_service._rpc_client = AsyncMock() + proxy_service._rpc_client.post.side_effect = RuntimeError("rocklet unavailable") + proxy_service._rpc_client.get.side_effect = RuntimeError("rocklet unavailable") + set_e2b_proxy_service(proxy_service) + monkeypatch.setattr("rock.sandbox.utils.timeout.time.time", lambda: 1767222000) + + app = FastAPI() + app.include_router(e2b_proxy_router) + return app, meta_store + + +async def _seed_sandbox(meta_store: SandboxMetaStore) -> None: + await meta_store.create( + "sandbox-123", + { + "sandbox_id": "sandbox-123", + "host_ip": "10.0.1.23", + "image": "linux-dind", + "metadata": { + "ap-job-id": "job-123", + "ap-template": "swe-bench", + "e2b.agents.kruise.io/return-sandbox-ip": "true", + }, + "state": State.RUNNING, + "cpus": 4, + "memory": "8g", + "disk": "20g", + "create_time": "2026-01-01T06:59:00+08:00", + "start_time": "2026-01-01T07:00:00+08:00", + }, + timeout_info={"auto_clear_time": "60", "expire_time": "1767225600"}, + ) + + +@pytest.mark.asyncio +async def test_get_sandbox_returns_e2b_detail_without_requiring_headers(e2b_proxy_app): + app, meta_store = e2b_proxy_app + await _seed_sandbox(meta_store) + + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + response = await client.get("/sandboxes/sandbox-123") + + assert response.status_code == 200 + assert response.json() == { + "sandboxID": "sandbox-123", + "metadata": { + "ap-job-id": "job-123", + "ap-template": "swe-bench", + "e2b.agents.kruise.io/return-sandbox-ip": "true", + "e2b.agents.kruise.io/sandbox-ip": "10.0.1.23", + }, + "state": "running", + "clientID": "rock", + "templateID": "linux-dind", + "envdVersion": "0.1.0", + "cpuCount": 4, + "memoryMB": 8192, + "diskSizeMB": 20480, + "startedAt": "2026-01-01T07:00:00+08:00", + "endAt": "2026-01-01T08:00:00+08:00", + } + + +@pytest.mark.asyncio +async def test_get_sandbox_maps_archived_to_paused(e2b_proxy_app): + app, meta_store = e2b_proxy_app + await _seed_sandbox(meta_store) + await meta_store.archive( + "sandbox-123", + { + "state": State.ARCHIVED, + "archive_time": "2026-01-01T09:00:00+08:00", + }, + ) + + assert await meta_store.get_timeout("sandbox-123") is None + + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + response = await client.get("/sandboxes/sandbox-123") + + assert response.status_code == 200 + assert response.json()["state"] == "paused" + assert response.json()["endAt"] == "2026-01-01T09:00:00+08:00" diff --git a/tests/unit/admin/test_create_app_routes.py b/tests/unit/admin/test_create_app_routes.py index 99259e8389..585b5fbc56 100644 --- a/tests/unit/admin/test_create_app_routes.py +++ b/tests/unit/admin/test_create_app_routes.py @@ -21,6 +21,7 @@ def test_proxy_role_mounts_proxy_router(): assert commit_path in paths assert commit_status_path in paths assert "/sandboxes" not in paths + assert "/sandboxes/{sandboxID}" in paths def test_proxy_role_excludes_admin_router(): @@ -40,6 +41,7 @@ def test_admin_role_mounts_admin_routers(): assert commit_path in paths assert commit_status_path not in paths assert "/sandboxes" in paths + assert "/sandboxes/{sandboxID}" not in paths @pytest.mark.skipif(not hasattr(socket, "SO_REUSEPORT"), reason="SO_REUSEPORT unavailable") diff --git a/tests/unit/sandbox/service/test_opensandbox_service_proxy.py b/tests/unit/sandbox/service/test_opensandbox_service_proxy.py index 23a82db756..530ad5288f 100644 --- a/tests/unit/sandbox/service/test_opensandbox_service_proxy.py +++ b/tests/unit/sandbox/service/test_opensandbox_service_proxy.py @@ -34,7 +34,11 @@ def service(): result._opensandbox_backend = backend result._session_registry = registry result._opensandbox_protocol = "https" - result._meta_store = SimpleNamespace(get=AsyncMock(return_value=_info()), update=AsyncMock()) + result._meta_store = SimpleNamespace( + get=AsyncMock(return_value=_info()), + get_timeout=AsyncMock(return_value=None), + update=AsyncMock(), + ) result._update_expire_time = AsyncMock() return result, backend diff --git a/tests/unit/sandbox/service/test_runtime_backend_routing.py b/tests/unit/sandbox/service/test_runtime_backend_routing.py index ace4090332..a139ba1217 100644 --- a/tests/unit/sandbox/service/test_runtime_backend_routing.py +++ b/tests/unit/sandbox/service/test_runtime_backend_routing.py @@ -27,7 +27,12 @@ def _info(*, backend=None, state=State.RUNNING, opensandbox_id=None): extended_params["backend"] = backend if opensandbox_id is not None: extended_params["opensandbox_id"] = opensandbox_id - return {"sandbox_id": "sbx-1", "state": state, "extended_params": extended_params} + return { + "sandbox_id": "sbx-1", + "state": state, + "metadata": {"ap-job-id": "job-123"}, + "extended_params": extended_params, + } @pytest.mark.asyncio @@ -103,6 +108,7 @@ async def test_opensandbox_get_status_does_not_probe_rocklet(opensandbox_proxy_s assert result.host_ip is None assert result.port_mapping is None assert result.swe_rex_version is None + assert result.metadata == {"ap-job-id": "job-123"} backend.get_state.assert_awaited_once() rocklet_probe.assert_not_awaited() diff --git a/tests/unit/sandbox/test_proxy_get_status.py b/tests/unit/sandbox/test_proxy_get_status.py index 1a12ad1dc2..604be9b293 100644 --- a/tests/unit/sandbox/test_proxy_get_status.py +++ b/tests/unit/sandbox/test_proxy_get_status.py @@ -27,6 +27,7 @@ def _make_meta_info(state: State = State.RUNNING, host_ip: str = "10.0.0.1") -> "host_ip": host_ip, "host_name": "node-1", "image": "python:3.11", + "metadata": {"ap-job-id": "job-123"}, "user_id": "u1", "experiment_id": "e1", "namespace": "ns1", @@ -133,6 +134,7 @@ async def test_running_sandbox_returns_phases_and_port_mapping( "docker_run": {"status": "running", "message": "running"}, } assert result.port_mapping == {22555: 22555, 8080: 8080} + assert result.metadata == {"ap-job-id": "job-123"} async def test_get_status_returns_gpu_info(self, proxy_service, mock_meta_store, mock_rpc_client): info = _make_meta_info(state=State.RUNNING)