fix(docker): return sandbox extensions in lifecycle responses

Docker provider was missing extensions support added in PR #1112.
Store opensandbox.extensions.* keys as container labels and return
them in CreateSandboxResponse, get_sandbox, and list_sandboxes.

Also rename apply_extensions_to_annotations/extract_extensions_from_annotations
to apply_extensions_to_mapping/extract_extensions_from_mapping since these
functions operate on generic mappings, not just k8s annotations.

Author: Pangjiping

server/opensandbox_server/extensions/__init__.py

@@ -18,8 +18,8 @@
 
 from opensandbox_server.extensions.codec import (
     apply_access_renew_extend_seconds_to_mapping,
-    apply_extensions_to_annotations,
-    extract_extensions_from_annotations,
+    apply_extensions_to_mapping,
+    extract_extensions_from_mapping,
 )
 from opensandbox_server.extensions.keys import (
     ACCESS_RENEW_EXTEND_SECONDS_KEY,
@@ -42,8 +42,8 @@
     "ACCESS_RENEW_EXTEND_SECONDS_MAX",
     "validate_extensions",
     "apply_access_renew_extend_seconds_to_mapping",
-    "apply_extensions_to_annotations",
-    "extract_extensions_from_annotations",
+    "apply_extensions_to_mapping",
+    "extract_extensions_from_mapping",
     "EXTENSIONS_ANNOTATION_PREFIX",
     "ANNOTATION_METADATA_PREFIX",
     "BOOTSTRAP_EXECD_ISOLATION_KEY",

server/opensandbox_server/extensions/codec.py

@@ -46,7 +46,7 @@ def apply_access_renew_extend_seconds_to_mapping(
     mapping[metadata_key] = s
 
 
-def apply_extensions_to_annotations(
+def apply_extensions_to_mapping(
     annotations: MutableMapping[str, str],
     extensions: Optional[Dict[str, str]],
 ) -> None:
@@ -70,7 +70,7 @@ def apply_extensions_to_annotations(
             annotations[annotation_key] = value
 
 
-def extract_extensions_from_annotations(
+def extract_extensions_from_mapping(
     annotations: Optional[Mapping[str, str]],
 ) -> Optional[Dict[str, str]]:
     """

server/opensandbox_server/services/docker/container_ops.py

@@ -31,6 +31,7 @@
 
 from opensandbox_server.extensions import (
     apply_access_renew_extend_seconds_to_mapping,
+    apply_extensions_to_mapping,
 )
 
 from opensandbox_server.api.schema import (
@@ -304,6 +305,7 @@ def _build_labels_and_env(
             labels[SANDBOX_SNAPSHOT_ID_LABEL] = request.snapshot_id
 
         apply_access_renew_extend_seconds_to_mapping(labels, request.extensions)
+        apply_extensions_to_mapping(labels, request.extensions)
 
         env_dict = request.env or {}
         environment = []

server/opensandbox_server/services/docker/docker_service.py

@@ -41,7 +41,9 @@
     ACCESS_RENEW_EXTEND_SECONDS_METADATA_KEY,
     BOOTSTRAP_EXECD_ISOLATION_KEY,
     ISOLATION_UPPER_MOUNT_PATH,
+    extract_extensions_from_mapping,
 )
+from opensandbox_server.extensions.keys import EXTENSIONS_ANNOTATION_PREFIX
 from opensandbox_server.api.schema import (
     CreateSandboxRequest,
     CreateSandboxResponse,
@@ -594,6 +596,7 @@ def _container_to_sandbox(self, container, sandbox_id: Optional[str] = None) ->
             platform=platform_spec,
             status=status_info,
             metadata=metadata,
+            extensions=extract_extensions_from_mapping(labels),
             entrypoint=entrypoint,
             expiresAt=expires_at,
             createdAt=created_at,
@@ -762,13 +765,18 @@ def _pending_to_sandbox(sandbox_id: str, pending: PendingSandbox) -> Sandbox:
         snapshot_id = getattr(pending.request, "snapshot_id", None)
         if not isinstance(snapshot_id, str) or not snapshot_id:
             snapshot_id = None
+        extensions = {
+            k: v for k, v in (pending.request.extensions or {}).items()
+            if k.startswith(EXTENSIONS_ANNOTATION_PREFIX)
+        } or None
         return Sandbox(
             id=sandbox_id,
             image=None if snapshot_id else pending.request.image,
             snapshotId=snapshot_id,
             platform=pending.request.platform,
             status=pending.status,
             metadata=pending.request.metadata,
+            extensions=extensions,
             entrypoint=pending.request.entrypoint,
             expiresAt=pending.expires_at,
             createdAt=pending.created_at,
@@ -1029,6 +1037,7 @@ def _provision_sandbox(
             id=sandbox_id,
             status=status_info,
             metadata=request.metadata,
+            extensions=extract_extensions_from_mapping(labels),
             platform=effective_platform or request.platform,
             expiresAt=expires_at,
             createdAt=created_at,

server/opensandbox_server/services/k8s/kubernetes_service.py

@@ -30,8 +30,8 @@
 
 from opensandbox_server.extensions import (
     apply_access_renew_extend_seconds_to_mapping,
-    apply_extensions_to_annotations,
-    extract_extensions_from_annotations,
+    apply_extensions_to_mapping,
+    extract_extensions_from_mapping,
 )
 from opensandbox_server.extensions.keys import ACCESS_RENEW_EXTEND_SECONDS_METADATA_KEY
 from opensandbox_server.api.schema import (
@@ -677,7 +677,7 @@ async def create_sandbox(self, request: CreateSandboxRequest) -> CreateSandboxRe
                 secure_access_token_factory=generate_secure_access_token,
             )
             apply_access_renew_extend_seconds_to_mapping(context.annotations, request.extensions)
-            apply_extensions_to_annotations(context.annotations, request.extensions)
+            apply_extensions_to_mapping(context.annotations, request.extensions)
 
             ensure_volumes_valid(
                 request.volumes,
@@ -832,7 +832,7 @@ async def create_sandbox(self, request: CreateSandboxRequest) -> CreateSandboxRe
                     created_at=created_at,
                     expires_at=context.expires_at,
                     metadata=request.metadata,
-                    extensions=extract_extensions_from_annotations(annotations),
+                    extensions=extract_extensions_from_mapping(annotations),
                     entrypoint=request.entrypoint,
                     platform=effective_platform or request.platform,
                 )

server/opensandbox_server/services/k8s/workload_mapper.py

@@ -17,7 +17,7 @@
 from typing import Any, Optional
 
 from opensandbox_server.api.schema import ImageSpec, PlatformSpec, Sandbox, SandboxStatus
-from opensandbox_server.extensions import extract_extensions_from_annotations
+from opensandbox_server.extensions import extract_extensions_from_mapping
 from opensandbox_server.services.constants import SANDBOX_ID_LABEL, SANDBOX_SNAPSHOT_ID_LABEL
 
 
@@ -78,7 +78,7 @@ def _build_sandbox_from_workload(workload: Any, workload_provider: Any) -> Sandb
         created_at=creation_timestamp,
         expires_at=expires_at,
         metadata=user_metadata if user_metadata else None,
-        extensions=extract_extensions_from_annotations(annotations),
+        extensions=extract_extensions_from_mapping(annotations),
         image=image_spec,
         snapshotId=snapshot_id,
         entrypoint=entrypoint,

server/tests/test_docker_service.py

@@ -1348,6 +1348,119 @@ def test_build_labels_stores_extensions_json():
 
     assert labels[ACCESS_RENEW_EXTEND_SECONDS_METADATA_KEY] == "3600"
 
+
+def test_build_labels_stores_opensandbox_extensions():
+    service = DockerSandboxService(config=_app_config())
+    request = CreateSandboxRequest(
+        image=ImageSpec(uri="python:3.11"),
+        resourceLimits=ResourceLimits(root={}),
+        env={},
+        entrypoint=["python"],
+        extensions={
+            "opensandbox.extensions.pool-ref": "my-pool",
+            "opensandbox.extensions.custom": "value",
+            "access.renew.extend.seconds": "1800",
+        },
+    )
+
+    labels, _ = service._build_labels_and_env("sandbox-ext2", request, None)
+
+    assert labels["opensandbox.io/extensions.pool-ref"] == "my-pool"
+    assert labels["opensandbox.io/extensions.custom"] == "value"
+    assert "opensandbox.io/extensions.access.renew.extend.seconds" not in labels
+
+
+@patch("opensandbox_server.services.docker.docker_service.docker")
+def test_container_to_sandbox_returns_extensions(mock_docker):
+    mock_client = MagicMock()
+    mock_client.containers.list.return_value = []
+    mock_docker.from_env.return_value = mock_client
+
+    service = DockerSandboxService(config=_app_config())
+    container = MagicMock()
+    container.attrs = {
+        "Config": {
+            "Labels": {
+                SANDBOX_ID_LABEL: "sandbox-ext",
+                "opensandbox.io/extensions.pool-ref": "my-pool",
+                "opensandbox.io/extensions.custom": "value",
+                "opensandbox.io/access-renew-extend-seconds": "1800",
+            },
+            "Cmd": ["python"],
+        },
+        "Created": "2025-01-01T00:00:00Z",
+        "State": {
+            "Status": "running",
+            "Running": True,
+            "FinishedAt": "0001-01-01T00:00:00Z",
+            "ExitCode": 0,
+        },
+    }
+    container.image = MagicMock(tags=["python:3.11"], short_id="sha-img")
+
+    sandbox = service._container_to_sandbox(container)
+
+    assert sandbox.extensions == {
+        "opensandbox.extensions.pool-ref": "my-pool",
+        "opensandbox.extensions.custom": "value",
+    }
+    assert sandbox.metadata is None
+
+
+@pytest.mark.asyncio
+@patch("opensandbox_server.services.docker.docker_service.docker")
+async def test_create_sandbox_response_includes_extensions(mock_docker):
+    mock_client = MagicMock()
+    mock_client.containers.list.return_value = []
+    mock_docker.from_env.return_value = mock_client
+
+    service = DockerSandboxService(config=_app_config())
+    request = CreateSandboxRequest(
+        image=ImageSpec(uri="python:3.11"),
+        resourceLimits=ResourceLimits(root={}),
+        env={},
+        entrypoint=["python"],
+        extensions={"opensandbox.extensions.test-key": "test-value"},
+    )
+
+    with patch.object(service, "_create_and_start_container") as mock_create:
+        mock_container = MagicMock()
+        mock_container.image = MagicMock(tags=["python:3.11"])
+        mock_create.return_value = mock_container
+        response = await service.create_sandbox(request)
+
+    assert response.extensions == {"opensandbox.extensions.test-key": "test-value"}
+
+
+@patch("opensandbox_server.services.docker.docker_service.docker")
+def test_pending_sandbox_includes_extensions(mock_docker):
+    mock_client = MagicMock()
+    mock_client.containers.list.return_value = []
+    mock_docker.from_env.return_value = mock_client
+
+    service = DockerSandboxService(config=_app_config())
+    pending = PendingSandbox(
+        request=MagicMock(
+            metadata=None,
+            entrypoint=["python"],
+            image=ImageSpec(uri="python:3.11"),
+            platform=None,
+            snapshot_id=None,
+            extensions={
+                "opensandbox.extensions.pool-ref": "my-pool",
+                "access.renew.extend.seconds": "1800",
+            },
+        ),
+        created_at=datetime.now(timezone.utc),
+        expires_at=datetime.now(timezone.utc),
+        status=SandboxStatus(state="Pending"),
+    )
+
+    sandbox = service._pending_to_sandbox("sandbox-pending-ext", pending)
+
+    assert sandbox.extensions == {"opensandbox.extensions.pool-ref": "my-pool"}
+
+
 def test_build_labels_store_platform_constraints():
     service = DockerSandboxService(config=_app_config())
     request = CreateSandboxRequest(

server/tests/test_extensions.py

@@ -21,8 +21,8 @@
     ACCESS_RENEW_EXTEND_SECONDS_METADATA_KEY,
     ACCESS_RENEW_EXTEND_SECONDS_MIN,
     apply_access_renew_extend_seconds_to_mapping,
-    apply_extensions_to_annotations,
-    extract_extensions_from_annotations,
+    apply_extensions_to_mapping,
+    extract_extensions_from_mapping,
     validate_extensions
 )
 
@@ -105,7 +105,7 @@ class TestExtensionsToAnnotations:
     def test_single_extension_propagated(self):
         annotations: dict[str, str] = {}
         extensions = {"opensandbox.extensions.pool-ref": "my-pool"}
-        apply_extensions_to_annotations(annotations, extensions)
+        apply_extensions_to_mapping(annotations, extensions)
         assert annotations == {"opensandbox.io/extensions.pool-ref": "my-pool"}
 
     def test_multiple_extensions_propagated(self):
@@ -114,7 +114,7 @@ def test_multiple_extensions_propagated(self):
             "opensandbox.extensions.pool-ref": "my-pool",
             "opensandbox.extensions.custom-key": "custom-value",
         }
-        apply_extensions_to_annotations(annotations, extensions)
+        apply_extensions_to_mapping(annotations, extensions)
         assert annotations == {
             "opensandbox.io/extensions.pool-ref": "my-pool",
             "opensandbox.io/extensions.custom-key": "custom-value",
@@ -126,7 +126,7 @@ def test_non_prefix_extension_not_propagated(self):
             "poolRef": "my-pool",
             "other.key": "value",
         }
-        apply_extensions_to_annotations(annotations, extensions)
+        apply_extensions_to_mapping(annotations, extensions)
         assert annotations == {}
 
     def test_mixed_extensions_propagated_only_prefix_keys(self):
@@ -136,23 +136,23 @@ def test_mixed_extensions_propagated_only_prefix_keys(self):
             "poolRef": "ignored",
             "access.renew.extend.seconds": "1800",
         }
-        apply_extensions_to_annotations(annotations, extensions)
+        apply_extensions_to_mapping(annotations, extensions)
         assert annotations == {"opensandbox.io/extensions.pool-ref": "my-pool"}
 
     def test_empty_extensions_noop(self):
         annotations: dict[str, str] = {"existing": "value"}
-        apply_extensions_to_annotations(annotations, None)
+        apply_extensions_to_mapping(annotations, None)
         assert annotations == {"existing": "value"}
 
     def test_empty_extensions_dict_noop(self):
         annotations: dict[str, str] = {"existing": "value"}
-        apply_extensions_to_annotations(annotations, {})
+        apply_extensions_to_mapping(annotations, {})
         assert annotations == {"existing": "value"}
 
     def test_preserves_existing_annotations(self):
         annotations: dict[str, str] = {"existing": "value"}
         extensions = {"opensandbox.extensions.new-key": "new-value"}
-        apply_extensions_to_annotations(annotations, extensions)
+        apply_extensions_to_mapping(annotations, extensions)
         assert annotations == {
             "existing": "value",
             "opensandbox.io/extensions.new-key": "new-value",
@@ -166,7 +166,7 @@ def test_restores_prefixed_annotations(self):
             "opensandbox.io/extensions.localized": "中文数据",
         }
 
-        assert extract_extensions_from_annotations(annotations) == {
+        assert extract_extensions_from_mapping(annotations) == {
             "opensandbox.extensions.custom-key": "custom-value",
             "opensandbox.extensions.localized": "中文数据",
         }
@@ -177,8 +177,8 @@ def test_ignores_non_extension_annotations(self):
             "other": "value",
         }
 
-        assert extract_extensions_from_annotations(annotations) is None
+        assert extract_extensions_from_mapping(annotations) is None
 
     def test_empty_annotations_return_none(self):
-        assert extract_extensions_from_annotations({}) is None
-        assert extract_extensions_from_annotations(None) is None
+        assert extract_extensions_from_mapping({}) is None
+        assert extract_extensions_from_mapping(None) is None

tests/python/tests/test_sandbox_e2e.py

@@ -357,6 +357,36 @@ async def test_01b_manual_cleanup(self):
 
         logger.info("TEST 1 PASSED: Sandbox lifecycle and health test completed successfully")
 
+    @pytest.mark.timeout(120)
+    @pytest.mark.order(1)
+    async def test_01_extensions_round_trip(self):
+        """Verify extensions are returned in create response, get_info, and list."""
+        cfg = create_connection_config()
+        ext_sandbox = await Sandbox.create(
+            image=SandboxImageSpec(get_sandbox_image()),
+            resource=get_e2e_sandbox_resource(),
+            connection_config=cfg,
+            timeout=timedelta(minutes=2),
+            ready_timeout=timedelta(seconds=30),
+            metadata={"tag": "e2e-extensions"},
+            extensions={
+                "opensandbox.extensions.test-key": "test-value",
+                "opensandbox.extensions.second": "second-value",
+            },
+            health_check_polling_interval=timedelta(milliseconds=500),
+        )
+        try:
+            info = await ext_sandbox.get_info()
+            assert info.extensions is not None, "extensions missing from get_info"
+            assert info.extensions.get("opensandbox.extensions.test-key") == "test-value"
+            assert info.extensions.get("opensandbox.extensions.second") == "second-value"
+            logger.info("extensions round-trip OK: %s", info.extensions)
+        finally:
+            try:
+                await ext_sandbox.kill()
+            except Exception as e:
+                logger.warning("extensions test teardown kill failed: %s", e)
+            await ext_sandbox.close()
 
     @pytest.mark.timeout(120)
     @pytest.mark.order(1)