Separate DockerWorkspace and DockerDevWorkspace to fix import issue

examples/02_remote_agent_server/02_convo_with_docker_sandboxed_server.py

@@ -11,7 +11,7 @@
     get_logger,
 )
 from openhands.tools.preset.default import get_default_agent
-from openhands.workspace import DockerWorkspace
+from openhands.workspace import DockerDevWorkspace
 
 
 logger = get_logger(__name__)
@@ -38,7 +38,7 @@ def detect_platform():
 
 # 2) Create a Docker-based remote workspace that will set up and manage
 #    the Docker container automatically
-with DockerWorkspace(
+with DockerDevWorkspace(
     # dynamically build agent-server image
     base_image="nikolaik/python-nodejs:python3.12-nodejs22",
     # use pre-built image for faster startup

examples/02_remote_agent_server/03_browser_use_with_docker_sandboxed_server.py

@@ -7,7 +7,7 @@
 from openhands.sdk import LLM, Conversation, get_logger
 from openhands.sdk.conversation.impl.remote_conversation import RemoteConversation
 from openhands.tools.preset.default import get_default_agent
-from openhands.workspace import DockerWorkspace
+from openhands.workspace import DockerDevWorkspace
 
 
 logger = get_logger(__name__)
@@ -32,7 +32,7 @@ def detect_platform():
 
 
 # Create a Docker-based remote workspace with extra ports for browser access
-with DockerWorkspace(
+with DockerDevWorkspace(
     base_image="nikolaik/python-nodejs:python3.12-nodejs22",
     host_port=8011,
     platform=detect_platform(),
@@ -85,7 +85,7 @@ def event_callback(event) -> None:
         y = None
         while y != "y":
             y = input(
-                "Because you've enabled extra_ports=True in DockerWorkspace, "
+                "Because you've enabled extra_ports=True in DockerDevWorkspace, "
                 "you can open a browser tab to see the *actual* browser OpenHands "
                 "is interacting with via VNC.\n\n"
                 "Link: http://localhost:8012/vnc.html?autoconnect=1&resize=remote\n\n"

examples/02_remote_agent_server/04_vscode_with_docker_sandboxed_server.py

@@ -7,7 +7,7 @@
 from openhands.sdk import LLM, Conversation, get_logger
 from openhands.sdk.conversation.impl.remote_conversation import RemoteConversation
 from openhands.tools.preset.default import get_default_agent
-from openhands.workspace import DockerWorkspace
+from openhands.workspace import DockerDevWorkspace
 
 
 logger = get_logger(__name__)
@@ -35,7 +35,7 @@ def detect_platform():
     return "linux/amd64"
 
 
-with DockerWorkspace(
+with DockerDevWorkspace(
     base_image="nikolaik/python-nodejs:python3.12-nodejs22",
     host_port=18010,
     platform=detect_platform(),
@@ -97,7 +97,7 @@ def event_callback(event) -> None:
     while y != "y":
         y = input(
             "\n"
-            "Because you've enabled extra_ports=True in DockerWorkspace, "
+            "Because you've enabled extra_ports=True in DockerDevWorkspace, "
             "you can open VSCode Web to see the workspace.\n\n"
             f"VSCode URL: {vscode_url}\n\n"
             "The VSCode should have the OpenHands settings extension installed:\n"

openhands-workspace/openhands/workspace/__init__.py

@@ -1,10 +1,25 @@
 """OpenHands Workspace - Docker and container-based workspace implementations."""
 
+from typing import TYPE_CHECKING
+
 from .docker import DockerWorkspace
 from .remote_api import APIRemoteWorkspace
 
 
+if TYPE_CHECKING:
+    from .docker import DockerDevWorkspace
+
 __all__ = [
     "DockerWorkspace",
+    "DockerDevWorkspace",
     "APIRemoteWorkspace",
 ]
+
+
+def __getattr__(name: str):
+    """Lazy import DockerDevWorkspace to avoid build module imports."""
+    if name == "DockerDevWorkspace":
+        from .docker import DockerDevWorkspace
+
+        return DockerDevWorkspace
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

openhands-workspace/openhands/workspace/docker/__init__.py

@@ -1,6 +1,20 @@
 """Docker workspace implementation."""
 
+from typing import TYPE_CHECKING
+
 from .workspace import DockerWorkspace
 
 
-__all__ = ["DockerWorkspace"]
+if TYPE_CHECKING:
+    from .dev_workspace import DockerDevWorkspace
+
+__all__ = ["DockerWorkspace", "DockerDevWorkspace"]
+
+
+def __getattr__(name: str):
+    """Lazy import DockerDevWorkspace to avoid build module imports."""
+    if name == "DockerDevWorkspace":
+        from .dev_workspace import DockerDevWorkspace
+
+        return DockerDevWorkspace
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

openhands-workspace/openhands/workspace/docker/dev_workspace.py

@@ -0,0 +1,99 @@
+"""Docker development workspace with on-the-fly image building capability."""
+
+from typing import Literal
+
+from pydantic import Field, model_validator
+
+from openhands.agent_server.docker.build import (
+    BuildOptions,
+    TargetType,
+    build,
+)
+
+from .workspace import DockerWorkspace
+
+
+class DockerDevWorkspace(DockerWorkspace):
+    """Docker workspace with on-the-fly image building capability.
+
+    This workspace extends DockerWorkspace to support building Docker images
+    on-the-fly from a base image. This is useful for development and testing
+    scenarios where you need to customize the agent server environment.
+
+    Note: This class requires the OpenHands SDK workspace structure and should
+    only be used within the OpenHands development environment or when you have
+    the full SDK source code available.
+
+    For production use cases with pre-built images, use DockerWorkspace instead.
+
+    Example:
+        with DockerDevWorkspace(
+            base_image="python:3.12",
+            target="source"
+        ) as workspace:
+            result = workspace.execute_command("ls -la")
+    """
+
+    # Add base_image support
+    base_image: str | None = Field(
+        default=None,
+        description=(
+            "Base Docker image to build the agent server from. "
+            "Mutually exclusive with server_image."
+        ),
+    )
+
+    # Add build-specific options
+    target: TargetType = Field(
+        default="source", description="Build target for the Docker image."
+    )
+
+    # Override platform with stricter type for building
+    platform: Literal["linux/amd64", "linux/arm64"] = Field(  # type: ignore[assignment]
+        default="linux/amd64",
+        description=(
+            "Platform for the Docker image. "
+            "Only linux/amd64 and linux/arm64 are supported for building."
+        ),
+    )
+
+    @model_validator(mode="after")
+    def _validate_images(self):
+        """Ensure exactly one of base_image or server_image is provided."""
+        if (self.base_image is None) == (self.server_image is None):
+            raise ValueError(
+                "Exactly one of 'base_image' or 'server_image' must be set."
+            )
+        if self.base_image and "ghcr.io/openhands/agent-server" in self.base_image:
+            raise ValueError(
+                "base_image cannot be a pre-built agent-server image. "
+                "Use server_image=... instead."
+            )
+        return self
+
+    def _get_image(self) -> str:
+        """Build the image if base_image is provided, otherwise use server_image.
+
+        This overrides the parent method to add on-the-fly image building
+        capability.
+
+        Returns:
+            The Docker image tag to use.
+        """
+        if self.base_image:
+            # Build the image from base_image
+            build_opts = BuildOptions(
+                base_image=self.base_image,
+                target=self.target,
+                platforms=[self.platform],
+                push=False,
+            )
+            tags = build(opts=build_opts)
+            if not tags or len(tags) == 0:
+                raise RuntimeError("Build failed, no image tags returned")
+            return tags[0]
+        elif self.server_image:
+            # Use pre-built image
+            return self.server_image
+        else:
+            raise ValueError("Either base_image or server_image must be set")

openhands-workspace/openhands/workspace/docker/workspace.py

@@ -11,12 +11,6 @@
 
 from pydantic import Field, PrivateAttr, model_validator
 
-from openhands.agent_server.docker.build import (
-    BuildOptions,
-    PlatformType,
-    TargetType,
-    build,
-)
 from openhands.sdk.logger import get_logger
 from openhands.sdk.utils.command import execute_command
 from openhands.sdk.workspace import RemoteWorkspace
@@ -59,12 +53,17 @@ def find_available_tcp_port(
 class DockerWorkspace(RemoteWorkspace):
     """Remote workspace that sets up and manages a Docker container.
 
-    This workspace creates a Docker container running the OpenHands agent server,
-    waits for it to become healthy, and then provides remote workspace operations
-    through the container's HTTP API.
+    This workspace creates a Docker container running a pre-built OpenHands agent
+    server image, waits for it to become healthy, and then provides remote workspace
+    operations through the container's HTTP API.
+
+    Note: This class only works with pre-built images. To build images on-the-fly
+    from a base image, use DockerDevWorkspace instead.
 
     Example:
-        with DockerWorkspace(base_image="python:3.12") as workspace:
+        with DockerWorkspace(
+            server_image="ghcr.io/openhands/agent-server:latest"
+        ) as workspace:
             result = workspace.execute_command("ls -la")
     """
 
@@ -79,17 +78,9 @@ class DockerWorkspace(RemoteWorkspace):
     )
 
     # Docker-specific configuration
-    base_image: str | None = Field(
-        default=None,
-        description="Base Docker image to use for the agent server container. "
-        "Mutually exclusive with server_image.",
-    )
     server_image: str | None = Field(
         default=None,
-        description=(
-            "Pre-built agent server image to use. If None, builds from base_image."
-            "Mutually exclusive with base_image."
-        ),
+        description="Pre-built agent server image to use.",
     )
     host_port: int | None = Field(
         default=None,
@@ -106,10 +97,7 @@ class DockerWorkspace(RemoteWorkspace):
     detach_logs: bool = Field(
         default=True, description="Whether to stream Docker logs in background."
     )
-    target: TargetType = Field(
-        default="source", description="Build target for the Docker image."
-    )
-    platform: PlatformType = Field(
+    platform: str = Field(
         default="linux/amd64", description="Platform for the Docker image."
     )
     extra_ports: bool = Field(
@@ -124,19 +112,55 @@ class DockerWorkspace(RemoteWorkspace):
     _container_id: str | None = PrivateAttr(default=None)
     _logs_thread: threading.Thread | None = PrivateAttr(default=None)
     _stop_logs: threading.Event = PrivateAttr(default_factory=threading.Event)
-    _image: str = PrivateAttr()
 
     @model_validator(mode="after")
-    def _validate_images(self):
-        """Ensure exactly one of base_image or server_image is provided; cache it."""
-        if (self.base_image is None) == (self.server_image is None):
-            raise ValueError(
-                "Exactly one of 'base_image' or 'server_image' must be set."
-            )
+    def _validate_server_image(self):
+        """Ensure server_image is provided, unless subclass will provide it."""
+        # Allow subclasses to defer server_image validation if they have base_image
+        # Check if this is a subclass with base_image attribute
+        if (
+            type(self).__name__ != "DockerWorkspace"
+            and hasattr(self, "base_image")
+            and getattr(self, "base_image") is not None
+        ):
+            # This is a subclass (e.g., DockerDevWorkspace) that will provide
+            # server_image in model_post_init
+            return self
+        if self.server_image is None:
+            raise ValueError("server_image must be provided")
         return self
 
     def model_post_init(self, context: Any) -> None:
         """Set up the Docker container and initialize the remote workspace."""
+        # Subclasses should call _get_image() to get the image to use
+        # This allows them to build or prepare the image before container startup
+        image = self._get_image()
+        self._start_container(image, context)
+
+    def _get_image(self) -> str:
+        """Get the Docker image to use for the container.
+
+        Subclasses can override this to provide custom image resolution logic
+        (e.g., building images on-the-fly).
+
+        Returns:
+            The Docker image tag to use.
+        """
+        if self.server_image is None:
+            raise ValueError("server_image must be set")
+        return self.server_image
+
+    def _start_container(self, image: str, context: Any) -> None:
+        """Start the Docker container with the given image.
+
+        This method handles all container lifecycle: port allocation, Docker
+        validation, container creation, health checks, and RemoteWorkspace
+        initialization.
+
+        Args:
+            image: The Docker image tag to use.
+            context: The Pydantic context from model_post_init.
+        """
         # Determine port
         if self.host_port is None:
             self.host_port = find_available_tcp_port()
@@ -164,29 +188,6 @@ def model_post_init(self, context: Any) -> None:
                 "Docker Desktop/daemon."
             )
 
-        # Build image if needed
-
-        if self.base_image:
-            if "ghcr.io/openhands/agent-server" in self.base_image:
-                raise RuntimeError(
-                    "base_image cannot be a pre-built agent-server image. "
-                    "Use server_image=... instead."
-                )
-            build_opts = BuildOptions(
-                base_image=self.base_image,
-                target=self.target,
-                platforms=[self.platform],
-                push=False,
-            )
-            tags = build(opts=build_opts)
-            assert tags and len(tags) > 0, "Build failed, no image tags returned"
-            self._image = tags[0]
-
-        elif self.server_image:
-            self._image = self.server_image
-        else:
-            raise RuntimeError("Unreachable: one of base_image or server_image is set")
-
         # Prepare Docker run flags
         flags: list[str] = []
         for key in self.forward_env:
@@ -227,7 +228,7 @@ def model_post_init(self, context: Any) -> None:
             "--name",
             f"agent-server-{uuid.uuid4()}",
             *flags,
-            self._image,
+            image,
             "--host",
             "0.0.0.0",
             "--port",