feat: add async sandbox commit

Closes #1279

Author: zhongwen666

rock/actions/__init__.py

@@ -23,6 +23,9 @@
     CloseResponse,
     CloseSessionResponse,
     CommandResponse,
+    CommitErrorCode,
+    CommitPhase,
+    CommitStatusResponse,
     CreateBashSessionResponse,
     CreateSessionResponse,
     ExecuteBashSessionResponse,
@@ -69,6 +72,9 @@
     "IsAliveResponse",
     "SandboxStatusResponse",
     "CommandResponse",
+    "CommitPhase",
+    "CommitErrorCode",
+    "CommitStatusResponse",
     "WriteFileResponse",
     "OssSetupResponse",
     "ExecuteBashSessionResponse",

rock/actions/sandbox/response.py

@@ -67,6 +67,40 @@ class SandboxStatusResponse(BaseModel):
     state_history: list[StateTransitionRecord] = []
 
 
+class CommitPhase(str, Enum):
+    RUNNING = "RUNNING"
+    SUCCEEDED = "SUCCEEDED"
+    FAILED = "FAILED"
+
+
+class CommitErrorCode(str, Enum):
+    SANDBOX_NOT_FOUND = "SANDBOX_NOT_FOUND"
+    WORKER_NOT_FOUND = "WORKER_NOT_FOUND"
+    WORKER_UNREACHABLE = "WORKER_UNREACHABLE"
+    SANDBOX_CONTAINER_NOT_FOUND = "SANDBOX_CONTAINER_NOT_FOUND"
+    COMMIT_CONFLICT = "COMMIT_CONFLICT"
+    DISPATCH_FAILED = "DISPATCH_FAILED"
+    LOGIN_FAILED = "LOGIN_FAILED"
+    COMMIT_FAILED = "COMMIT_FAILED"
+    PUSH_FAILED = "PUSH_FAILED"
+    TIMEOUT = "TIMEOUT"
+    PROCESS_LOST = "PROCESS_LOST"
+    STATUS_NOT_FOUND = "STATUS_NOT_FOUND"
+    STATUS_CORRUPTED = "STATUS_CORRUPTED"
+
+
+class CommitStatusResponse(BaseModel):
+    sandbox_id: str
+    image_tag: str
+    phase: CommitPhase
+    started_at: str
+    completed_at: str | None = None
+    exit_code: int | None = None
+    failed_stage: str | None = None
+    error_code: CommitErrorCode | None = None
+    error_message: str | None = None
+
+
 class CommandResponse(BaseModel):
     stdout: str = ""
     stderr: str = ""

rock/admin/entrypoints/sandbox_proxy_api.py

@@ -8,15 +8,18 @@
     BashObservation,
     CloseBashSessionResponse,
     CommandResponse,
+    CommitStatusResponse,
     CreateBashSessionResponse,
     ReadFileResponse,
     ResponseStatus,
     RockResponse,
+    SandboxResponse,
     UploadResponse,
     WriteFileResponse,
 )
 from rock.admin.proto.request import (
     BatchSandboxStatusRequest,
+    CommitRequest,
     SandboxBashAction,
     SandboxCloseBashSessionRequest,
     SandboxCommand,
@@ -28,7 +31,7 @@
 from rock.admin.proto.response import BatchSandboxStatusResponse, SandboxListResponse, SandboxStatusResponse
 from rock.common.exception import handle_exceptions
 from rock.common.port_validation import validate_port_forward_port
-from rock.common.validation import NonBlankStr
+from rock.common.validation import NonBlankStr, SandboxContainerName
 from rock.logger import init_logger
 from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService
 from rock.sdk.common.exceptions import BadRequestRockError
@@ -115,6 +118,18 @@ async def execute(command: SandboxCommand) -> RockResponse[CommandResponse]:
     return RockResponse(result=await sandbox_proxy_service.execute(command))
 
 
+@sandbox_proxy_router.post("/commit")
+@handle_exceptions(error_message="commit sandbox failed")
+async def commit(request: CommitRequest) -> RockResponse[CommitStatusResponse | SandboxResponse]:
+    return RockResponse(result=await sandbox_proxy_service.commit(request))
+
+
+@sandbox_proxy_router.get("/commit/{sandbox_id}")
+@handle_exceptions(error_message="get commit status failed")
+async def get_commit_status(sandbox_id: SandboxContainerName) -> RockResponse[CommitStatusResponse | SandboxResponse]:
+    return RockResponse(result=await sandbox_proxy_service.get_commit_status(sandbox_id))
+
+
 @sandbox_proxy_router.post("/create_session")
 @handle_exceptions(error_message="create session failed")
 async def create_session(request: SandboxCreateBashSessionRequest) -> RockResponse[CreateBashSessionResponse]:

rock/admin/proto/request.py

@@ -1,7 +1,7 @@
 from typing import Annotated, Literal, TypedDict
 
 from fastapi import Header
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, field_validator
 
 from rock import env_vars
 from rock.actions import (
@@ -12,7 +12,7 @@
     ReadFileRequest,
     WriteFileRequest,
 )
-from rock.common.validation import NonBlankStr
+from rock.common.validation import NonBlankStr, SandboxContainerName
 
 
 class SandboxStartRequest(BaseModel):
@@ -50,6 +50,23 @@ class SandboxStartRequest(BaseModel):
     """GPU accelerator type (e.g. 'A100', 'V100'). If not specified, any available GPU will be used."""
 
 
+class CommitRequest(BaseModel):
+    sandbox_id: SandboxContainerName
+    image_tag: NonBlankStr
+    username: str
+    password: str
+
+    @field_validator("image_tag", mode="before")
+    @classmethod
+    def validate_image_tag_for_status_file(cls, value: object) -> object:
+        if isinstance(value, str) and any(
+            character in value
+            for character in ("\x00", "\n", "\r", "\v", "\f", "\x1c", "\x1d", "\x1e", "\x85", "\u2028", "\u2029")
+        ):
+            raise ValueError("image_tag cannot contain NUL or line separators")
+        return value
+
+
 class SandboxCommand(Command):
     timeout: float | None = 1200
     """The timeout for the command. None means no timeout."""

rock/common/validation.py

@@ -16,3 +16,16 @@
 handler in ``rock.common.exception.request_validation_exception_handler`` maps
 back to the project's ``RockResponse(status=Failed, error=...)`` envelope.
 """
+
+
+SANDBOX_CONTAINER_NAME_PATTERN = r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$"
+
+SandboxContainerName = Annotated[
+    str,
+    StringConstraints(
+        strip_whitespace=True,
+        min_length=1,
+        max_length=128,
+        pattern=SANDBOX_CONTAINER_NAME_PATTERN,
+    ),
+]