feat(box): add session workspace quota enforcement and SDK quota metadata

Author: huanghuoguoguo

src/langbot_plugin/box/backend.py

@@ -130,7 +130,7 @@ async def start_session(self, spec: BoxSpec) -> BoxSessionInfo:
             f'image={spec.image} network={spec.network.value} '
             f'host_path={spec.host_path} host_path_mode={spec.host_path_mode.value} mount_path={spec.mount_path} '
             f'cpus={spec.cpus} memory_mb={spec.memory_mb} pids_limit={spec.pids_limit} '
-            f'read_only_rootfs={spec.read_only_rootfs}'
+            f'read_only_rootfs={spec.read_only_rootfs} workspace_quota_mb={spec.workspace_quota_mb}'
         )
 
         await self._run_command(args, timeout_sec=30, check=True)
@@ -148,6 +148,7 @@ async def start_session(self, spec: BoxSpec) -> BoxSessionInfo:
             memory_mb=spec.memory_mb,
             pids_limit=spec.pids_limit,
             read_only_rootfs=spec.read_only_rootfs,
+            workspace_quota_mb=spec.workspace_quota_mb,
             created_at=now,
             last_used_at=now,
         )

src/langbot_plugin/box/models.py

@@ -47,6 +47,7 @@ class BoxSpec(pydantic.BaseModel):
     memory_mb: int = 512
     pids_limit: int = 128
     read_only_rootfs: bool = True
+    workspace_quota_mb: int = 0
 
     @pydantic.model_validator(mode='before')
     @classmethod
@@ -102,6 +103,13 @@ def validate_pids_limit(cls, value: int) -> int:
             raise ValueError('pids_limit must be at least 1')
         return value
 
+    @pydantic.field_validator('workspace_quota_mb')
+    @classmethod
+    def validate_workspace_quota_mb(cls, value: int) -> int:
+        if value < 0:
+            raise ValueError('workspace_quota_mb must be greater than or equal to 0')
+        return value
+
     @pydantic.field_validator('session_id')
     @classmethod
     def validate_session_id(cls, value: str) -> str:
@@ -162,6 +170,7 @@ class BoxProfile(pydantic.BaseModel):
     memory_mb: int = 512
     pids_limit: int = 128
     read_only_rootfs: bool = True
+    workspace_quota_mb: int = 0
     locked: frozenset[str] = frozenset()
 
     model_config = pydantic.ConfigDict(frozen=True)
@@ -225,6 +234,7 @@ class BoxSessionInfo(pydantic.BaseModel):
     memory_mb: int = 512
     pids_limit: int = 128
     read_only_rootfs: bool = True
+    workspace_quota_mb: int = 0
     created_at: dt.datetime
     last_used_at: dt.datetime
 

src/langbot_plugin/box/nsjail_backend.py

@@ -136,7 +136,8 @@ async def start_session(self, spec: BoxSpec) -> BoxSessionInfo:
             f'session_id={spec.session_id} session_dir={session_dir} '
             f'network={spec.network.value} '
             f'host_path={spec.host_path} host_path_mode={spec.host_path_mode.value} mount_path={spec.mount_path} '
-            f'cpus={spec.cpus} memory_mb={spec.memory_mb} pids_limit={spec.pids_limit}'
+            f'cpus={spec.cpus} memory_mb={spec.memory_mb} pids_limit={spec.pids_limit} '
+            f'workspace_quota_mb={spec.workspace_quota_mb}'
         )
 
         return BoxSessionInfo(
@@ -152,6 +153,7 @@ async def start_session(self, spec: BoxSpec) -> BoxSessionInfo:
             memory_mb=spec.memory_mb,
             pids_limit=spec.pids_limit,
             read_only_rootfs=spec.read_only_rootfs,
+            workspace_quota_mb=spec.workspace_quota_mb,
             created_at=now,
             last_used_at=now,
         )

src/langbot_plugin/box/runtime.py

@@ -220,7 +220,8 @@ async def _get_or_create_session(self, spec: BoxSpec) -> _RuntimeSession:
                 f'network={info.network.value} '
                 f'host_path={info.host_path} '
                 f'host_path_mode={info.host_path_mode.value} '
-                f'mount_path={info.mount_path}'
+                f'mount_path={info.mount_path} '
+                f'workspace_quota_mb={info.workspace_quota_mb}'
             )
             return runtime_session
 
@@ -290,6 +291,7 @@ def _assert_session_compatible(self, session: BoxSessionInfo, spec: BoxSpec):
             'memory_mb',
             'pids_limit',
             'read_only_rootfs',
+            'workspace_quota_mb',
         )
         for field in _COMPAT_FIELDS:
             session_val = getattr(session, field)