diff --git a/coding-agent/postmortem/index.yaml b/coding-agent/postmortem/index.yaml index d4b850c06..29f75355f 100644 --- a/coding-agent/postmortem/index.yaml +++ b/coding-agent/postmortem/index.yaml @@ -590,3 +590,28 @@ patterns: - and - teardown - idempotent +- id: PM-0024 + title: Preserve cloud workspaces until cleanup is verified + path: patterns/PM-0024-preserve-cloud-workspaces-until-cleanup-is-verified.md + severity: medium + confidence: medium + subsystems: + - environment + - http + related_files: + - src/coding_agent/environment/docker_workspace_provider.py + - src/coding_agent/ui/http_server.py + - tests/coding_agent/environment/test_docker_workspace_provider.py + - tests/ui/test_http_server.py + related_commits: [] + keywords: + - preserve + - cloud + - workspaces + - until + - cleanup + - verified + - rollback + - cancellation + - docker + - inspect diff --git a/coding-agent/postmortem/patterns/PM-0024-preserve-cloud-workspaces-until-cleanup-is-verified.md b/coding-agent/postmortem/patterns/PM-0024-preserve-cloud-workspaces-until-cleanup-is-verified.md new file mode 100644 index 000000000..2d8dc4d56 --- /dev/null +++ b/coding-agent/postmortem/patterns/PM-0024-preserve-cloud-workspaces-until-cleanup-is-verified.md @@ -0,0 +1,40 @@ +--- +id: PM-0024 +title: Preserve cloud workspaces until cleanup is verified +status: active +severity: medium +confidence: medium +subsystems: +- environment +- http +related_commits: [] +related_files: +- src/coding_agent/environment/docker_workspace_provider.py +- src/coding_agent/ui/http_server.py +- tests/coding_agent/environment/test_docker_workspace_provider.py +- tests/ui/test_http_server.py +release_checks: +- Run focused cloud workspace rollback regressions in `tests/coding_agent/environment/test_docker_workspace_provider.py` and `tests/ui/test_http_server.py` before release. +- Review create-then-cleanup, cancellation, and external-subprocess verification paths before deleting host workspace state. +--- + +# Summary + +Cloud workspace teardown must preserve host workspace state until container cleanup is known to have succeeded or the container is known to be absent. + +# Trigger Conditions + +- Changes in cloud workspace provisioning, rollback, Docker cleanup, HTTP session creation, or cancellation handling. +- Review feedback mentions leaked provisioned workspaces, skipped rollback on cancellation, or treating Docker inspect failures as missing containers. + +# Known Fix Signals + +- `src/coding_agent/environment/docker_workspace_provider.py` +- `src/coding_agent/ui/http_server.py` +- `tests/coding_agent/environment/test_docker_workspace_provider.py` +- `tests/ui/test_http_server.py` + +# Release Review Checklist + +- Run focused cloud workspace rollback regressions in `tests/coding_agent/environment/test_docker_workspace_provider.py` and `tests/ui/test_http_server.py` before release. +- Review create-then-cleanup, cancellation, and external-subprocess verification paths before deleting host workspace state. diff --git a/coding-agent/src/coding_agent/environment/__init__.py b/coding-agent/src/coding_agent/environment/__init__.py index af9cb531a..1bd346beb 100644 --- a/coding-agent/src/coding_agent/environment/__init__.py +++ b/coding-agent/src/coding_agent/environment/__init__.py @@ -6,9 +6,12 @@ from .docker_workspace_provider import DockerCloudWorkspaceClient from .local import LocalEnvironment from .workspace_provider import ( + cleanup_cloud_binding_from_config, CloudWorkspaceClientFactory, + CloudWorkspaceSource, WorkspaceProvider, cloud_client_factory_from_config, + provision_cloud_binding_from_config, register_workspace_provider, ) @@ -18,8 +21,11 @@ "CloudEnvironment", "CloudWorkspaceClient", "CloudWorkspaceClientFactory", + "CloudWorkspaceSource", + "cleanup_cloud_binding_from_config", "LocalEnvironment", "WorkspaceProvider", "cloud_client_factory_from_config", + "provision_cloud_binding_from_config", "register_workspace_provider", ] diff --git a/coding-agent/src/coding_agent/environment/docker_workspace_provider.py b/coding-agent/src/coding_agent/environment/docker_workspace_provider.py index 3bc42b023..4260bf2aa 100644 --- a/coding-agent/src/coding_agent/environment/docker_workspace_provider.py +++ b/coding-agent/src/coding_agent/environment/docker_workspace_provider.py @@ -4,14 +4,22 @@ import posixpath import re import secrets +import shutil import subprocess +import time +import uuid from collections.abc import Iterable from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, cast, override from .cloud import CloudCommandResult -from .workspace_provider import CloudWorkspaceClientFactory, WorkspaceProvider, register_workspace_provider +from .workspace_provider import ( + CloudWorkspaceClientFactory, + CloudWorkspaceSource, + WorkspaceProvider, + register_workspace_provider, +) if TYPE_CHECKING: from .cloud import CloudWorkspaceClient @@ -24,6 +32,8 @@ _DOCKER_TIMEOUT_MARGIN_SECONDS = 1 _DOCKER_TIMEOUT_CLEANUP_SECONDS = 3 _DOCKER_TIMEOUT_SENTINEL_PREFIX = "__CODING_AGENT_DOCKER_TIMEOUT__:" +_DOCKER_CONTAINER_REMOVAL_WAIT_SECONDS = 5.0 +_DOCKER_CONTAINER_REMOVAL_POLL_INTERVAL_SECONDS = 0.1 @dataclass(frozen=True) @@ -34,6 +44,7 @@ class _DockerWorkspaceProviderConfig: docker_binary: str env_allowlist: tuple[str, ...] exec_user: str | None + image: str class DockerCloudWorkspaceClient: @@ -94,7 +105,9 @@ def grep_search(self, pattern: str, directory: str, include: str) -> list[str]: if not resolved_path.is_file(): continue for line_number, line in enumerate( - resolved_path.read_text(encoding="utf-8", errors="replace").splitlines(), + resolved_path.read_text( + encoding="utf-8", errors="replace" + ).splitlines(), start=1, ): if matcher.search(line): @@ -158,7 +171,7 @@ def run_command( "}", "trap _coding_agent_cleanup EXIT", "trap _coding_agent_timeout TERM", - 'if command -v setsid >/dev/null 2>&1; then', + "if command -v setsid >/dev/null 2>&1; then", ' setsid /bin/sh -c "$1" sh "$pidfile" "$2" &', "else", ' /bin/sh -c "$1" sh "$pidfile" "$2" &', @@ -202,7 +215,9 @@ def run_command( ) except subprocess.TimeoutExpired as exc: self._cleanup_timed_out_command(timeout_pidfile) - raise TimeoutError(f"docker exec command timed out after {timeout}s") from exc + raise TimeoutError( + f"docker exec command timed out after {timeout}s" + ) from exc if _contains_timeout_sentinel(result.stderr, timeout_sentinel): raise TimeoutError(f"docker exec command timed out after {timeout}s") @@ -227,9 +242,7 @@ def _resolve_directory_path(self, directory: str) -> tuple[str, Path]: ) return remote_path, self._validated_workspace_path(host_path, remote_path) - def _filtered_env_items( - self, env: dict[str, str] | None - ) -> list[tuple[str, str]]: + def _filtered_env_items(self, env: dict[str, str] | None) -> list[tuple[str, str]]: if env is None: return [] items: list[tuple[str, str]] = [] @@ -287,7 +300,9 @@ def _validated_workspace_path(self, host_path: Path, remote_path: str) -> Path: try: _ = resolved_path.relative_to(self._workspace_root) except ValueError as exc: - raise ValueError(f"Path is outside docker workspace: {remote_path}") from exc + raise ValueError( + f"Path is outside docker workspace: {remote_path}" + ) from exc return resolved_path def _workspace_entry_remote_path(self, host_path: Path) -> str: @@ -309,21 +324,85 @@ def build_client(binding: CloudWorkspaceBinding) -> CloudWorkspaceClient: return build_client + @override + def provision_cloud_workspace_binding( + self, + config: dict[str, object], + source: CloudWorkspaceSource, + ) -> CloudWorkspaceBinding: + from ..ui.execution_binding import CloudWorkspaceBinding + + kind = source.get("kind") + if kind != "docker": + raise ValueError( + "cloud workspace source kind is not supported for provider=docker" + ) + + provider_config = _docker_workspace_provider_config(config) + workspace_id = f"ws-{uuid.uuid4().hex}" + workspace_root = _workspace_root_for_id( + provider_config.workspace_root, workspace_id + ) + workspace_root.mkdir(parents=True, exist_ok=False) + binding = CloudWorkspaceBinding( + workspace_url=( + f"docker://{_container_name(provider_config, workspace_id)}{provider_config.container_workspace_root}" + ), + workspace_id=workspace_id, + ) + try: + _start_docker_workspace_container(provider_config, binding) + except Exception as exc: + cleanup_failed = False + try: + _remove_docker_workspace_container(provider_config, workspace_id) + except Exception as cleanup_exc: + cleanup_failed = True + note = ( + "failed to clean up docker workspace container after start failure: " + + str(cleanup_exc) + ) + exc.add_note(note) + if not cleanup_failed and workspace_root.exists(): + shutil.rmtree(workspace_root) + raise + return binding + + @override + def cleanup_cloud_workspace_binding( + self, + config: dict[str, object], + binding: CloudWorkspaceBinding, + ) -> None: + provider_config = _docker_workspace_provider_config(config) + _remove_docker_workspace_container(provider_config, binding.workspace_id) + workspace_root = _workspace_root_for_id( + provider_config.workspace_root, binding.workspace_id + ) + if workspace_root.exists(): + shutil.rmtree(workspace_root) + def _docker_workspace_provider_config( config: dict[str, object], ) -> _DockerWorkspaceProviderConfig: workspace_root_raw = config.get("workspace_root") if not isinstance(workspace_root_raw, str) or not workspace_root_raw.strip(): - raise ValueError("cloud_workspace.workspace_root is required for provider=docker") + raise ValueError( + "cloud_workspace.workspace_root is required for provider=docker" + ) container_workspace_root = _container_workspace_root(config) container_name_prefix = _optional_string(config.get("container_name_prefix"), "") docker_binary = _optional_string(config.get("docker_binary"), "docker") exec_user = _optional_string(config.get("exec_user"), None) - env_allowlist = tuple(_string_list(config.get("env_allowlist"), key="env_allowlist")) + env_allowlist = tuple( + _string_list(config.get("env_allowlist"), key="env_allowlist") + ) + image = _optional_string(config.get("image"), "python:3.11-slim") assert container_name_prefix is not None assert docker_binary is not None + assert image is not None return _DockerWorkspaceProviderConfig( workspace_root=Path(workspace_root_raw).expanduser().resolve(), @@ -332,6 +411,7 @@ def _docker_workspace_provider_config( docker_binary=docker_binary, env_allowlist=env_allowlist, exec_user=exec_user, + image=image, ) @@ -342,7 +422,9 @@ def _container_workspace_root(config: dict[str, object]) -> str: if not normalized.startswith("/"): raise ValueError("cloud_workspace.container_workspace_root must be absolute") if not normalized.lstrip("/"): - raise ValueError("cloud_workspace.container_workspace_root must not resolve to /") + raise ValueError( + "cloud_workspace.container_workspace_root must not resolve to /" + ) return normalized @@ -388,7 +470,9 @@ def _contains_timeout_sentinel(stderr: str, sentinel: str) -> bool: def _validate_workspace_id(workspace_id: str) -> None: if not _WORKSPACE_ID_RE.fullmatch(workspace_id): - raise ValueError(f"unsupported workspace id for docker provider: {workspace_id}") + raise ValueError( + f"unsupported workspace id for docker provider: {workspace_id}" + ) def _workspace_root_for_id(workspace_root: Path, workspace_id: str) -> Path: @@ -396,14 +480,131 @@ def _workspace_root_for_id(workspace_root: Path, workspace_id: str) -> Path: try: _ = candidate.relative_to(workspace_root) except ValueError as exc: - raise ValueError(f"workspace id escapes configured workspace root: {workspace_id}") from exc + raise ValueError( + f"workspace id escapes configured workspace root: {workspace_id}" + ) from exc return candidate +def _container_name( + provider_config: _DockerWorkspaceProviderConfig, workspace_id: str +) -> str: + return f"{provider_config.container_name_prefix}{workspace_id}" + + +def _start_docker_workspace_container( + provider_config: _DockerWorkspaceProviderConfig, + binding: CloudWorkspaceBinding, +) -> None: + workspace_root = _workspace_root_for_id( + provider_config.workspace_root, binding.workspace_id + ) + command = [ + provider_config.docker_binary, + "run", + "-d", + "--name", + _container_name(provider_config, binding.workspace_id), + "-v", + f"{workspace_root}:{provider_config.container_workspace_root}", + "-w", + provider_config.container_workspace_root, + provider_config.image, + "sleep", + "infinity", + ] + _ = subprocess.run( + command, + shell=False, + capture_output=True, + text=True, + timeout=30, + env=None, + check=True, + ) + + +def _remove_docker_workspace_container( + provider_config: _DockerWorkspaceProviderConfig, + workspace_id: str, +) -> None: + container_name = _container_name(provider_config, workspace_id) + command = [ + provider_config.docker_binary, + "rm", + "-f", + container_name, + ] + try: + result = subprocess.run( + command, + shell=False, + capture_output=True, + text=True, + timeout=10, + env=None, + check=False, + ) + except (OSError, subprocess.SubprocessError) as exc: + raise RuntimeError( + f"failed to remove docker workspace container: {container_name}" + ) from exc + + if result.returncode != 0 and "No such container" not in result.stderr: + raise RuntimeError( + f"failed to remove docker workspace container: {container_name}" + ) + + deadline = time.monotonic() + _DOCKER_CONTAINER_REMOVAL_WAIT_SECONDS + while time.monotonic() < deadline: + if not _docker_container_exists(provider_config, container_name): + return + time.sleep(_DOCKER_CONTAINER_REMOVAL_POLL_INTERVAL_SECONDS) + + raise TimeoutError( + f"docker workspace container still exists after cleanup: {container_name}" + ) + + +def _docker_container_exists( + provider_config: _DockerWorkspaceProviderConfig, + container_name: str, +) -> bool: + command = [ + provider_config.docker_binary, + "container", + "inspect", + container_name, + ] + try: + result = subprocess.run( + command, + shell=False, + capture_output=True, + text=True, + timeout=5, + env=None, + check=False, + ) + except (OSError, subprocess.SubprocessError) as exc: + raise RuntimeError( + f"failed to inspect docker workspace container: {container_name}" + ) from exc + if result.returncode == 0: + return True + if "No such container" in result.stderr: + return False + raise RuntimeError( + f"failed to inspect docker workspace container: {container_name}" + ) + + def _normalize_remote_path(path: str, workspace_root: str) -> str: normalized_root = posixpath.normpath(workspace_root) if not normalized_root.lstrip("/"): - raise ValueError("cloud_workspace.container_workspace_root must not resolve to /") + raise ValueError( + "cloud_workspace.container_workspace_root must not resolve to /" + ) resolved = path if path.startswith("/") else posixpath.join(workspace_root, path) normalized = posixpath.normpath(resolved) workspace_prefix = normalized_root.rstrip("/") + "/" diff --git a/coding-agent/src/coding_agent/environment/workspace_provider.py b/coding-agent/src/coding_agent/environment/workspace_provider.py index 2684bfc61..5a7750f6d 100644 --- a/coding-agent/src/coding_agent/environment/workspace_provider.py +++ b/coding-agent/src/coding_agent/environment/workspace_provider.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Protocol, TypeAlias from .cloud import CloudWorkspaceClient @@ -12,6 +12,7 @@ CloudWorkspaceClientFactory: TypeAlias = Callable[ ["CloudWorkspaceBinding"], CloudWorkspaceClient ] +CloudWorkspaceSource: TypeAlias = Mapping[str, object] class WorkspaceProvider(Protocol): @@ -19,6 +20,18 @@ def build_cloud_client_factory( self, config: dict[str, object] ) -> CloudWorkspaceClientFactory: ... + def provision_cloud_workspace_binding( + self, + config: dict[str, object], + source: CloudWorkspaceSource, + ) -> "CloudWorkspaceBinding": ... + + def cleanup_cloud_workspace_binding( + self, + config: dict[str, object], + binding: "CloudWorkspaceBinding", + ) -> None: ... + _WORKSPACE_PROVIDERS: dict[str, WorkspaceProvider] = {} @@ -33,6 +46,27 @@ def register_workspace_provider(name: str, provider: WorkspaceProvider) -> None: def cloud_client_factory_from_config( config: dict[str, object], ) -> CloudWorkspaceClientFactory: + provider = _provider_from_config(config) + return provider.build_cloud_client_factory(config) + + +def provision_cloud_binding_from_config( + config: dict[str, object], + source: CloudWorkspaceSource, +) -> "CloudWorkspaceBinding": + provider = _provider_from_config(config) + return provider.provision_cloud_workspace_binding(config, source) + + +def cleanup_cloud_binding_from_config( + config: dict[str, object], + binding: "CloudWorkspaceBinding", +) -> None: + provider = _provider_from_config(config) + provider.cleanup_cloud_workspace_binding(config, binding) + + +def _provider_from_config(config: dict[str, object]) -> WorkspaceProvider: provider_name = config.get("provider") if not isinstance(provider_name, str) or not provider_name.strip(): raise ValueError( @@ -42,4 +76,4 @@ def cloud_client_factory_from_config( provider = _WORKSPACE_PROVIDERS.get(provider_name.strip().lower()) if provider is None: raise ValueError(f"unsupported cloud workspace provider: {provider_name}") - return provider.build_cloud_client_factory(config) + return provider diff --git a/coding-agent/src/coding_agent/ui/http_server.py b/coding-agent/src/coding_agent/ui/http_server.py index b21f0ca31..63852d78d 100644 --- a/coding-agent/src/coding_agent/ui/http_server.py +++ b/coding-agent/src/coding_agent/ui/http_server.py @@ -20,7 +20,11 @@ from agentkit.config.loader import load_config as load_agent_toml from agentkit.errors import ConfigError from coding_agent.approval import ApprovalPolicy -from coding_agent.environment import cloud_client_factory_from_config +from coding_agent.environment import ( + cleanup_cloud_binding_from_config, + cloud_client_factory_from_config, + provision_cloud_binding_from_config, +) from coding_agent.ui.binding_resolver import DefaultBindingResolver from coding_agent.ui.session_manager import Session, SessionManager from coding_agent.ui.session_owner_store import ( @@ -97,7 +101,8 @@ def _load_cloud_workspace_config() -> dict[str, Any]: config_path = Path(__file__).resolve().parent.parent / "agent.toml" try: return cast( - dict[str, Any], load_agent_toml(config_path).extra.get("cloud_workspace", {}) + dict[str, Any], + load_agent_toml(config_path).extra.get("cloud_workspace", {}), ) except (ConfigError, OSError) as exc: if isinstance(exc, ConfigError): @@ -121,6 +126,14 @@ def _build_binding_resolver() -> DefaultBindingResolver: ) +def _cleanup_provisioned_cloud_binding(binding: CloudWorkspaceBinding) -> None: + cloud_workspace_config = _load_cloud_workspace_config() + provider = cloud_workspace_config.get("provider") + if not isinstance(provider, str) or not provider.strip(): + return + cleanup_cloud_binding_from_config(cloud_workspace_config, binding) + + def _storage_uses_pg_http_sessions(storage_config: dict[str, Any]) -> bool: http_backend = storage_config.get("http_session_backend") if http_backend is not None: @@ -129,8 +142,7 @@ def _storage_uses_pg_http_sessions(storage_config: dict[str, Any]) -> bool: session_backend = storage_config.get("session_backend") if session_backend is not None: return ( - isinstance(session_backend, str) - and session_backend.strip().lower() == "pg" + isinstance(session_backend, str) and session_backend.strip().lower() == "pg" ) tape_backend = str(storage_config.get("tape_backend", "jsonl")).strip().lower() @@ -167,6 +179,7 @@ def _build_session_manager() -> SessionManager: manager = SessionManager( storage_config=storage_config, binding_resolver=_build_binding_resolver(), + provisioned_cloud_binding_cleanup=_cleanup_provisioned_cloud_binding, ) if not _storage_uses_pg_http_sessions(storage_config): return manager @@ -304,6 +317,45 @@ def _execution_binding_from_request( ) +def _provisioned_execution_binding_from_request( + body: CreateSessionRequest | None, +) -> ExecutionBinding | None: + if body is None: + return None + if body.execution_binding is not None and body.workspace_source is not None: + raise ValueError( + "execution_binding and workspace_source cannot be set together" + ) + explicit_binding = _execution_binding_from_request(body) + if explicit_binding is not None: + return explicit_binding + if body.workspace_source is None: + return None + + cloud_workspace_config = _load_cloud_workspace_config() + if cloud_workspace_config.get("enabled") is not True: + raise ValueError( + "cloud workspace provisioning requires cloud_workspace.enabled=true" + ) + return provision_cloud_binding_from_config( + cloud_workspace_config, + body.workspace_source.model_dump(mode="python"), + ) + + +def _session_origin_from_request( + body: CreateSessionRequest | None, + binding: ExecutionBinding | None, +) -> dict[str, str]: + origin = { + "channel": "http", + "binding_kind": "local" if binding is None else binding.kind, + } + if body is not None and body.workspace_source is not None: + origin["workspace_source_kind"] = body.workspace_source.kind + return origin + + def _http_safe_tool_result_payload(msg: ToolResultDelta) -> dict[str, Any]: return { "session_id": msg.session_id, @@ -670,17 +722,48 @@ async def create_session( "auto": ApprovalPolicy.AUTO, } approval_policy = approval_policy_map.get(approval_policy_str, ApprovalPolicy.AUTO) + provisioned_binding: CloudWorkspaceBinding | None = None - # Create session using SessionManager - session_id = await session_manager.create_session( - repo_path=repo_path, - execution_binding=_execution_binding_from_request(body), - approval_policy=approval_policy, - provider_name=body.provider if body else None, - model_name=body.model if body else None, - base_url=body.base_url if body else None, - max_steps=body.max_steps if body and body.max_steps is not None else 30, - ) + try: + execution_binding = _provisioned_execution_binding_from_request(body) + if ( + body is not None + and body.workspace_source is not None + and isinstance(execution_binding, CloudWorkspaceBinding) + ): + provisioned_binding = execution_binding + session_id = await session_manager.create_session( + repo_path=repo_path, + origin=_session_origin_from_request(body, execution_binding), + execution_binding=execution_binding, + approval_policy=approval_policy, + provider_name=body.provider if body else None, + model_name=body.model if body else None, + base_url=body.base_url if body else None, + max_steps=body.max_steps if body and body.max_steps is not None else 30, + ) + except (asyncio.CancelledError, KeyboardInterrupt, SystemExit): + if provisioned_binding is not None: + try: + await asyncio.to_thread( + _cleanup_provisioned_cloud_binding, provisioned_binding + ) + except Exception: + logger.exception("Failed to roll back provisioned cloud workspace") + raise + except Exception as exc: + if provisioned_binding is not None: + try: + await asyncio.to_thread( + _cleanup_provisioned_cloud_binding, provisioned_binding + ) + except Exception: + logger.exception("Failed to roll back provisioned cloud workspace") + if isinstance(exc, RuntimeError): + raise HTTPException(status_code=500, detail=str(exc)) from exc + if isinstance(exc, (ValueError, TypeError)): + raise HTTPException(status_code=400, detail=str(exc)) from exc + raise HTTPException(status_code=500, detail=str(exc)) from exc logger.info(f"Created session: {session_id}") return SessionResponse(session_id=session_id) @@ -713,7 +796,9 @@ async def send_prompt( raise _owner_conflict_http_exception(exc, session_id=session_id) from exc except RuntimeError as exc: if str(exc) == "turn already in progress": - raise HTTPException(status_code=409, detail="Turn already in progress") from exc + raise HTTPException( + status_code=409, detail="Turn already in progress" + ) from exc raise session.turn_in_progress = True diff --git a/coding-agent/src/coding_agent/ui/schemas.py b/coding-agent/src/coding_agent/ui/schemas.py index d1f25d40d..7ca1d9cce 100644 --- a/coding-agent/src/coding_agent/ui/schemas.py +++ b/coding-agent/src/coding_agent/ui/schemas.py @@ -19,7 +19,12 @@ class CloudWorkspaceBindingRequest(BaseModel): workspace_id: str = Field(..., min_length=1, max_length=200) +class DockerWorkspaceSourceRequest(BaseModel): + kind: Literal["docker"] + + ExecutionBindingRequest = LocalExecutionBindingRequest | CloudWorkspaceBindingRequest +WorkspaceSourceRequest = DockerWorkspaceSourceRequest class PromptRequest(BaseModel): @@ -33,6 +38,7 @@ class CreateSessionRequest(BaseModel): repo_path: str | None = Field(None, max_length=500) execution_binding: ExecutionBindingRequest | None = None + workspace_source: WorkspaceSourceRequest | None = None approval_policy: str = Field("auto", pattern="^(yolo|interactive|auto)$") provider: Literal[ "openai", diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index 9ccb5a32b..22d4517f4 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -65,7 +65,11 @@ from coding_agent.ui.session_owner_store import SessionOwnershipConflictError from coding_agent.ui.session_owner_store import SessionOwnershipConflictReason from coding_agent.ui.binding_resolver import BindingResolver, DefaultBindingResolver -from coding_agent.ui.execution_binding import ExecutionBinding, LocalExecutionBinding +from coding_agent.ui.execution_binding import ( + CloudWorkspaceBinding, + ExecutionBinding, + LocalExecutionBinding, +) logger = logging.getLogger(__name__) @@ -149,6 +153,7 @@ class Session: wire: LocalWire = field(init=False) approval_store: ApprovalStore = field(default_factory=ApprovalStore) repo_path: Path | None = None # legacy metadata; seeds default local binding + origin: dict[str, str] | None = None execution_binding: ExecutionBinding = field( # type: ignore[assignment] default=cast(ExecutionBinding, _DEFAULT_EXECUTION_BINDING), ) @@ -200,6 +205,7 @@ def as_dict(self) -> dict[str, Any]: "model_name": self.model_name, "base_url": self.base_url, "max_steps": self.max_steps, + "origin": None if self.origin is None else dict(self.origin), } def to_store_data(self) -> dict[str, Any]: @@ -210,6 +216,7 @@ def to_store_data(self) -> dict[str, Any]: # repo_path remains backward-compatible metadata and seeds the # default local binding when execution_binding is omitted. "repo_path": None if self.repo_path is None else str(self.repo_path), + "origin": None if self.origin is None else dict(self.origin), "execution_binding": self.execution_binding.to_dict(), "approval_policy": self.approval_policy.value, "provider_name": self.provider_name, @@ -224,6 +231,16 @@ def from_store_data(cls, data: dict[str, Any]) -> Session: repo_path_raw = data.get("repo_path") if repo_path_raw is not None and not isinstance(repo_path_raw, str): raise TypeError("session metadata has invalid repo_path") + origin_raw = data.get("origin") + if origin_raw is not None: + if not isinstance(origin_raw, dict): + raise TypeError("session metadata has invalid origin") + origin = { + key: _required_session_str(cast(dict[str, Any], origin_raw), key) + for key in cast(dict[str, Any], origin_raw) + } + else: + origin = None approval_policy_raw = data.get("approval_policy") if not isinstance(approval_policy_raw, str): raise TypeError("session metadata is missing approval_policy") @@ -261,6 +278,7 @@ def from_store_data(cls, data: dict[str, Any]) -> Session: ), approval_store=ApprovalStore(), repo_path=None if repo_path_raw is None else Path(repo_path_raw), + origin=origin, execution_binding=execution_binding, approval_policy=ApprovalPolicy(approval_policy_raw), provider_name=provider_name_raw, @@ -421,6 +439,9 @@ def __init__( checkpoint_service: CheckpointService | None = None, create_agent_fn: Callable[..., tuple[Any, Any]] | None = None, binding_resolver: BindingResolver | None = None, + provisioned_cloud_binding_cleanup: ( + Callable[[CloudWorkspaceBinding], None] | None + ) = None, owner_store: SessionOwnerStoreProtocol | None = None, owner_id: str | None = None, fencing_token: int | None = None, @@ -445,6 +466,7 @@ def __init__( ) self._create_agent = create_agent_fn self._binding_resolver = binding_resolver or DefaultBindingResolver() + self._provisioned_cloud_binding_cleanup = provisioned_cloud_binding_cleanup self.configure_owner_leases( owner_store=owner_store, owner_id=owner_id, @@ -574,6 +596,34 @@ def _close_runtime_sync_safe(self, session: Session) -> None: return _ = loop.create_task(self._close_runtime_adapter(adapter)) + def _session_uses_provisioned_cloud_workspace(self, session: Session) -> bool: + origin = session.origin + return ( + isinstance(session.execution_binding, CloudWorkspaceBinding) + and origin is not None + and origin.get("binding_kind") == "cloud" + and origin.get("workspace_source_kind") is not None + ) + + def _cleanup_provisioned_cloud_binding(self, session: Session) -> None: + if self._provisioned_cloud_binding_cleanup is None: + return + if not self._session_uses_provisioned_cloud_workspace(session): + return + binding = session.execution_binding + if not isinstance(binding, CloudWorkspaceBinding): + return + self._provisioned_cloud_binding_cleanup(binding) + + async def _cleanup_provisioned_cloud_binding_async(self, session: Session) -> None: + try: + await asyncio.to_thread(self._cleanup_provisioned_cloud_binding, session) + except Exception: + logger.exception( + "Failed to clean up provisioned cloud workspace for session %s", + session.id, + ) + def _create_agent_for_session(self, **kwargs: Any) -> tuple[Any, Any]: factory = self._create_agent if factory is None: @@ -869,6 +919,7 @@ async def _close_resource_async(self, resource: object) -> None: async def _remove_session_async_no_lock(self, session_id: str) -> None: session = await self.get_session_async(session_id) await self._close_runtime(session) + await self._cleanup_provisioned_cloud_binding_async(session) self._session_cache.pop(session_id, None) await self._run_store_io(self._store.delete, session_id) await self._release_owner_lease_for_session(session_id) @@ -1080,6 +1131,7 @@ def _hydrate_session(self, session: Session) -> Session: async def create_session( self, repo_path: Path | None = None, + origin: dict[str, str] | None = None, approval_policy: ApprovalPolicy = ApprovalPolicy.AUTO, provider: Any | None = None, provider_name: str | None = None, @@ -1141,6 +1193,7 @@ async def create_session( created_at=now, last_activity=now, repo_path=resolved_repo_path, + origin=None if origin is None else dict(origin), execution_binding=binding, approval_policy=approval_policy, provider=provider, @@ -1217,15 +1270,23 @@ def remove_session(self, session_id: str) -> None: raise KeyError(f"Session not found: {session_id}") session = self.get_session(session_id) self._close_runtime_sync_safe(session) + self._cleanup_provisioned_cloud_binding(session) self._session_cache.pop(session_id, None) self._store.delete(session_id) self._approval_stores.pop(session_id, None) self._session_turn_locks.pop(session_id, None) def clear_sessions(self) -> None: + cleared_session_ids = set(self._session_cache) for session in list(self._session_cache.values()): self._close_runtime_sync_safe(session) + self._cleanup_provisioned_cloud_binding(session) for session_id in list(self._store.list_sessions()): + if session_id not in cleared_session_ids: + session = self.get_session(session_id) + self._close_runtime_sync_safe(session) + self._cleanup_provisioned_cloud_binding(session) + self._session_cache.pop(session_id, None) self._store.delete(session_id) self._session_cache.clear() self._approval_stores.clear() diff --git a/coding-agent/tests/coding_agent/environment/test_docker_workspace_provider.py b/coding-agent/tests/coding_agent/environment/test_docker_workspace_provider.py index 2283aca4b..e7f674454 100644 --- a/coding-agent/tests/coding_agent/environment/test_docker_workspace_provider.py +++ b/coding-agent/tests/coding_agent/environment/test_docker_workspace_provider.py @@ -7,7 +7,12 @@ import pytest -from coding_agent.environment import DockerCloudWorkspaceClient, cloud_client_factory_from_config +from coding_agent.environment import ( + DockerCloudWorkspaceClient, + cleanup_cloud_binding_from_config, + cloud_client_factory_from_config, + provision_cloud_binding_from_config, +) from coding_agent.ui.execution_binding import CloudWorkspaceBinding @@ -50,7 +55,9 @@ def test_docker_workspace_provider_builds_client_from_config(tmp_path: Path) -> assert client.default_cwd == "/workspace" -def test_docker_cloud_client_maps_file_tools_to_remote_workspace(tmp_path: Path) -> None: +def test_docker_cloud_client_maps_file_tools_to_remote_workspace( + tmp_path: Path, +) -> None: workspace_root = tmp_path / "workspaces" target_root = workspace_root / "ws-123" nested = target_root / "pkg" @@ -184,7 +191,7 @@ def fake_run( timeout_wrapper = captured_command[17] child_wrapper = captured_command[19] assert TIMEOUT_SENTINEL_PREFIX in timeout_wrapper - assert 'trap _coding_agent_timeout TERM' in timeout_wrapper + assert "trap _coding_agent_timeout TERM" in timeout_wrapper assert 'setsid /bin/sh -c "$1" sh "$pidfile" "$2" &' in timeout_wrapper assert 'printf "%s\\n" "$$" > "$1"' in child_wrapper assert 'exec /bin/sh -c "$2"' in child_wrapper @@ -474,7 +481,9 @@ def test_docker_cloud_client_rejects_symlink_escape_paths(tmp_path: Path) -> Non ValueError, match=r"Path is outside docker workspace: /workspace/escape", ): - _ = client.run_command("python -V", cwd="/workspace/escape", env=None, timeout=5) + _ = client.run_command( + "python -V", cwd="/workspace/escape", env=None, timeout=5 + ) def test_docker_workspace_provider_requires_workspace_root() -> None: @@ -486,7 +495,9 @@ def test_docker_workspace_provider_requires_workspace_root() -> None: @pytest.mark.parametrize("root", ["/", "//", "/./", "/workspace/.."]) -def test_docker_workspace_provider_rejects_root_container_workspace_root(root: str) -> None: +def test_docker_workspace_provider_rejects_root_container_workspace_root( + root: str, +) -> None: with pytest.raises( ValueError, match=r"cloud_workspace\.container_workspace_root must not resolve to /", @@ -498,3 +509,388 @@ def test_docker_workspace_provider_rejects_root_container_workspace_root(root: s "container_workspace_root": root, } ) + + +def test_docker_workspace_provider_provisions_runnable_container( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + workspace_root = tmp_path / "workspaces" + workspace_root.mkdir() + captured_command: list[str] = [] + captured_kwargs: dict[str, object] = {} + + def fake_run( + command: list[str], **kwargs: object + ) -> subprocess.CompletedProcess[str]: + captured_command[:] = command + captured_kwargs.clear() + captured_kwargs.update(kwargs) + return subprocess.CompletedProcess(command, 0, stdout="", stderr="") + + monkeypatch.setattr(subprocess, "run", fake_run) + + binding = provision_cloud_binding_from_config( + { + "provider": "docker", + "workspace_root": str(workspace_root), + "container_name_prefix": "agent-", + "docker_binary": "/usr/bin/docker", + }, + {"kind": "docker"}, + ) + + assert binding.workspace_id.startswith("ws-") + assert (workspace_root / binding.workspace_id).is_dir() + assert captured_command[:8] == [ + "/usr/bin/docker", + "run", + "-d", + "--name", + f"agent-{binding.workspace_id}", + "-v", + f"{workspace_root / binding.workspace_id}:/workspace", + "-w", + ] + assert captured_command[8:] == [ + "/workspace", + "python:3.11-slim", + "sleep", + "infinity", + ] + assert captured_kwargs == { + "shell": False, + "capture_output": True, + "text": True, + "timeout": 30, + "env": None, + "check": True, + } + + +def test_docker_workspace_provider_provision_cleans_up_container_when_start_times_out( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + workspace_root = tmp_path / "workspaces" + workspace_root.mkdir() + commands: list[list[str]] = [] + + def fake_run( + command: list[str], **kwargs: object + ) -> subprocess.CompletedProcess[str]: + commands.append(command) + if command[1:3] == ["run", "-d"]: + timeout = kwargs.get("timeout") + assert isinstance(timeout, int) + raise subprocess.TimeoutExpired(command, timeout=timeout) + if command[1:3] == ["rm", "-f"]: + return subprocess.CompletedProcess(command, 0, stdout="", stderr="") + if command[1:3] == ["container", "inspect"]: + return subprocess.CompletedProcess( + command, + 1, + stdout="", + stderr=f"Error: No such container: {command[-1]}\n", + ) + raise AssertionError(f"unexpected docker command: {command}") + + monkeypatch.setattr(subprocess, "run", fake_run) + + with pytest.raises(subprocess.TimeoutExpired): + _ = provision_cloud_binding_from_config( + { + "provider": "docker", + "workspace_root": str(workspace_root), + "container_name_prefix": "agent-", + "docker_binary": "/usr/bin/docker", + }, + {"kind": "docker"}, + ) + + container_name = commands[0][4] + assert commands[1] == ["/usr/bin/docker", "rm", "-f", container_name] + assert commands[2] == ["/usr/bin/docker", "container", "inspect", container_name] + assert not any(workspace_root.iterdir()) + + +def test_docker_workspace_provider_cleanup_removes_nonempty_workspace( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + workspace_root = tmp_path / "workspaces" + workspace_root.mkdir() + removed_commands: list[list[str]] = [] + + def fake_run( + command: list[str], **kwargs: object + ) -> subprocess.CompletedProcess[str]: + del kwargs + removed_commands.append(command) + if command[1:3] == ["container", "inspect"]: + return subprocess.CompletedProcess( + command, + 1, + stdout="", + stderr=f"Error: No such container: {command[-1]}\n", + ) + return subprocess.CompletedProcess(command, 0, stdout="", stderr="") + + monkeypatch.setattr(subprocess, "run", fake_run) + + binding = provision_cloud_binding_from_config( + { + "provider": "docker", + "workspace_root": str(workspace_root), + "container_name_prefix": "agent-", + "docker_binary": "/usr/bin/docker", + }, + {"kind": "docker"}, + ) + proof_file = workspace_root / binding.workspace_id / "qa-proof.txt" + _ = proof_file.write_text("qa-from-cloud\n", encoding="utf-8") + + cleanup_cloud_binding_from_config( + { + "provider": "docker", + "workspace_root": str(workspace_root), + "container_name_prefix": "agent-", + "docker_binary": "/usr/bin/docker", + }, + binding, + ) + + assert removed_commands[1] == [ + "/usr/bin/docker", + "rm", + "-f", + f"agent-{binding.workspace_id}", + ] + assert not (workspace_root / binding.workspace_id).exists() + + +def test_docker_workspace_provider_cleanup_waits_for_container_removal( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + workspace_root = tmp_path / "workspaces" + workspace_root.mkdir() + commands: list[list[str]] = [] + inspect_calls = 0 + sleep_calls: list[float] = [] + + def fake_run( + command: list[str], **kwargs: object + ) -> subprocess.CompletedProcess[str]: + nonlocal inspect_calls + del kwargs + commands.append(command) + if command[1:3] == ["container", "inspect"]: + inspect_calls += 1 + if inspect_calls == 1: + return subprocess.CompletedProcess(command, 0, stdout="[]", stderr="") + return subprocess.CompletedProcess( + command, + 1, + stdout="", + stderr=f"Error: No such container: {command[-1]}\n", + ) + return subprocess.CompletedProcess(command, 0, stdout="", stderr="") + + monkeypatch.setattr(subprocess, "run", fake_run) + monkeypatch.setattr( + "coding_agent.environment.docker_workspace_provider.time.sleep", + sleep_calls.append, + ) + + binding = provision_cloud_binding_from_config( + { + "provider": "docker", + "workspace_root": str(workspace_root), + "container_name_prefix": "agent-", + "docker_binary": "/usr/bin/docker", + }, + {"kind": "docker"}, + ) + + cleanup_cloud_binding_from_config( + { + "provider": "docker", + "workspace_root": str(workspace_root), + "container_name_prefix": "agent-", + "docker_binary": "/usr/bin/docker", + }, + binding, + ) + + assert [command[1:3] for command in commands].count(["container", "inspect"]) == 2 + assert sleep_calls == [0.1] + + +def test_docker_workspace_provider_cleanup_preserves_workspace_when_container_remove_fails( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + workspace_root = tmp_path / "workspaces" + workspace_root.mkdir() + + def fake_run( + command: list[str], **kwargs: object + ) -> subprocess.CompletedProcess[str]: + del kwargs + if command[1:3] == ["container", "inspect"]: + return subprocess.CompletedProcess( + command, + 1, + stdout="", + stderr=f"Error: No such container: {command[-1]}\n", + ) + if command[1:3] == ["run", "-d"]: + return subprocess.CompletedProcess(command, 0, stdout="", stderr="") + return subprocess.CompletedProcess( + command, + 1, + stdout="", + stderr="permission denied", + ) + + monkeypatch.setattr(subprocess, "run", fake_run) + + binding = provision_cloud_binding_from_config( + { + "provider": "docker", + "workspace_root": str(workspace_root), + "container_name_prefix": "agent-", + "docker_binary": "/usr/bin/docker", + }, + {"kind": "docker"}, + ) + proof_file = workspace_root / binding.workspace_id / "qa-proof.txt" + _ = proof_file.write_text("qa-from-cloud\n", encoding="utf-8") + + with pytest.raises( + RuntimeError, match="failed to remove docker workspace container" + ): + cleanup_cloud_binding_from_config( + { + "provider": "docker", + "workspace_root": str(workspace_root), + "container_name_prefix": "agent-", + "docker_binary": "/usr/bin/docker", + }, + binding, + ) + + assert proof_file.exists() + + +@pytest.mark.parametrize( + "inspect_failure", + [ + subprocess.CompletedProcess( + ["/usr/bin/docker", "container", "inspect", "agent-ws-broken"], + 1, + stdout="", + stderr="Cannot connect to the Docker daemon\n", + ), + OSError("docker unavailable"), + ], +) +def test_docker_workspace_provider_cleanup_preserves_workspace_when_container_inspect_fails( + inspect_failure: subprocess.CompletedProcess[str] | OSError, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + workspace_root = tmp_path / "workspaces" + workspace_root.mkdir() + + def fake_run( + command: list[str], **kwargs: object + ) -> subprocess.CompletedProcess[str]: + del kwargs + if command[1:3] == ["run", "-d"]: + return subprocess.CompletedProcess(command, 0, stdout="", stderr="") + if command[1:3] == ["rm", "-f"]: + return subprocess.CompletedProcess(command, 0, stdout="", stderr="") + if command[1:3] == ["container", "inspect"]: + if isinstance(inspect_failure, OSError): + raise inspect_failure + return inspect_failure + raise AssertionError(f"unexpected docker command: {command}") + + monkeypatch.setattr(subprocess, "run", fake_run) + + binding = provision_cloud_binding_from_config( + { + "provider": "docker", + "workspace_root": str(workspace_root), + "container_name_prefix": "agent-", + "docker_binary": "/usr/bin/docker", + }, + {"kind": "docker"}, + ) + proof_file = workspace_root / binding.workspace_id / "qa-proof.txt" + _ = proof_file.write_text("qa-from-cloud\n", encoding="utf-8") + + with pytest.raises( + RuntimeError, match="failed to inspect docker workspace container" + ): + cleanup_cloud_binding_from_config( + { + "provider": "docker", + "workspace_root": str(workspace_root), + "container_name_prefix": "agent-", + "docker_binary": "/usr/bin/docker", + }, + binding, + ) + + assert proof_file.exists() + + +def test_docker_workspace_provider_cleanup_preserves_workspace_when_container_removal_times_out( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + workspace_root = tmp_path / "workspaces" + workspace_root.mkdir() + monotonic_values = iter([0.0, 0.0, 5.1]) + + def fake_run( + command: list[str], **kwargs: object + ) -> subprocess.CompletedProcess[str]: + del kwargs + return subprocess.CompletedProcess(command, 0, stdout="[]", stderr="") + + monkeypatch.setattr(subprocess, "run", fake_run) + monkeypatch.setattr( + "coding_agent.environment.docker_workspace_provider.time.monotonic", + lambda: next(monotonic_values), + ) + + def ignore_sleep(delay: float) -> None: + _ = delay + + monkeypatch.setattr( + "coding_agent.environment.docker_workspace_provider.time.sleep", + ignore_sleep, + ) + + binding = provision_cloud_binding_from_config( + { + "provider": "docker", + "workspace_root": str(workspace_root), + "container_name_prefix": "agent-", + "docker_binary": "/usr/bin/docker", + }, + {"kind": "docker"}, + ) + proof_file = workspace_root / binding.workspace_id / "qa-proof.txt" + _ = proof_file.write_text("qa-from-cloud\n", encoding="utf-8") + + with pytest.raises(TimeoutError, match="docker workspace container still exists"): + cleanup_cloud_binding_from_config( + { + "provider": "docker", + "workspace_root": str(workspace_root), + "container_name_prefix": "agent-", + "docker_binary": "/usr/bin/docker", + }, + binding, + ) + + assert proof_file.exists() diff --git a/coding-agent/tests/ui/test_execution_binding.py b/coding-agent/tests/ui/test_execution_binding.py index 24cdd1e03..9ca70fcde 100644 --- a/coding-agent/tests/ui/test_execution_binding.py +++ b/coding-agent/tests/ui/test_execution_binding.py @@ -12,7 +12,9 @@ from coding_agent.environment import CloudCommandResult, CloudEnvironment, LocalEnvironment from coding_agent.environment import workspace_provider as workspace_provider_module from coding_agent.environment.workspace_provider import ( + CloudWorkspaceSource, cloud_client_factory_from_config, + provision_cloud_binding_from_config, register_workspace_provider, ) from coding_agent.ui.execution_binding import ( @@ -210,6 +212,26 @@ def build_client(binding: CloudWorkspaceBinding) -> FakeCloudClient: return build_client + def provision_cloud_workspace_binding( + self, + config: dict[str, object], + source: CloudWorkspaceSource, + ) -> CloudWorkspaceBinding: + self.seen_configs.append(dict(config)) + assert source == {"kind": "docker"} + return CloudWorkspaceBinding( + workspace_url="https://workspace.example.com", + workspace_id="ws-123", + ) + + def cleanup_cloud_workspace_binding( + self, + config: dict[str, object], + binding: CloudWorkspaceBinding, + ) -> None: + self.seen_configs.append(dict(config)) + assert binding.workspace_id == "ws-123" + def test_cloud_workspace_provider_registry_builds_factory( monkeypatch: pytest.MonkeyPatch, @@ -244,3 +266,22 @@ def test_cloud_workspace_provider_config_requires_provider() -> None: match=r"cloud_workspace\.provider is required when cloud_workspace\.enabled=true", ): _ = cloud_client_factory_from_config({"enabled": True}) + + +def test_workspace_provider_registry_provisions_binding( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(workspace_provider_module, "_WORKSPACE_PROVIDERS", {}) + provider = FakeWorkspaceProvider() + register_workspace_provider("fake-provider", provider) + + binding = provision_cloud_binding_from_config( + { + "provider": "fake-provider", + "workspace_root": "/srv/workspaces", + }, + {"kind": "docker"}, + ) + + assert isinstance(binding, CloudWorkspaceBinding) + assert binding.workspace_id == "ws-123" diff --git a/coding-agent/tests/ui/test_http_server.py b/coding-agent/tests/ui/test_http_server.py index 72a62ef8c..6570a836a 100644 --- a/coding-agent/tests/ui/test_http_server.py +++ b/coding-agent/tests/ui/test_http_server.py @@ -28,7 +28,10 @@ from coding_agent.approval.store import ApprovalStore from coding_agent.environment import CloudCommandResult, CloudEnvironment from coding_agent.ui.binding_resolver import CloudBindingNotImplementedError -from coding_agent.ui.execution_binding import CloudWorkspaceBinding, LocalExecutionBinding +from coding_agent.ui.execution_binding import ( + CloudWorkspaceBinding, + LocalExecutionBinding, +) from coding_agent.wire.local import LocalWire from coding_agent.ui.session_manager import Session from coding_agent.ui.session_owner_store import SessionOwnerRecord @@ -36,6 +39,7 @@ from coding_agent.ui.http_server import ( APPROVAL_TIMEOUT_SECONDS, SESSION_IDLE_TIMEOUT_MINUTES, + _build_binding_resolver, _build_session_manager, _renew_owner_leases, _cleanup_event_queue_on_disconnect, @@ -293,9 +297,338 @@ async def test_http_create_session_stores_cloud_execution_binding(self, client): assert response.status_code == 200 session = session_manager.get_session(response.json()["session_id"]) assert isinstance(session.execution_binding, CloudWorkspaceBinding) - assert session.execution_binding.workspace_url == "https://workspace.example.com" + assert ( + session.execution_binding.workspace_url == "https://workspace.example.com" + ) assert session.execution_binding.workspace_id == "ws-123" + async def test_http_create_session_provisions_docker_cloud_workspace( + self, client, monkeypatch, tmp_path + ): + monkeypatch.setattr( + "coding_agent.environment.docker_workspace_provider._start_docker_workspace_container", + lambda provider_config, binding: None, + ) + monkeypatch.setattr( + "coding_agent.ui.http_server._load_cloud_workspace_config", + lambda: { + "enabled": True, + "provider": "docker", + "workspace_root": str(tmp_path), + "container_name_prefix": "agent-", + }, + ) + + response = await client.post( + "/sessions", + json={"workspace_source": {"kind": "docker"}}, + ) + + assert response.status_code == 200 + session = session_manager.get_session(response.json()["session_id"]) + assert isinstance(session.execution_binding, CloudWorkspaceBinding) + assert session.origin == { + "channel": "http", + "binding_kind": "cloud", + "workspace_source_kind": "docker", + } + assert (tmp_path / session.execution_binding.workspace_id).is_dir() + + info_response = await client.get(f"/sessions/{response.json()['session_id']}") + assert info_response.status_code == 200 + assert info_response.json()["origin"] == { + "channel": "http", + "binding_kind": "cloud", + "workspace_source_kind": "docker", + } + + async def test_create_session_rejects_conflicting_workspace_binding_inputs( + self, client + ): + response = await client.post( + "/sessions", + json={ + "execution_binding": { + "kind": "cloud", + "workspace_url": "https://workspace.example.com", + "workspace_id": "ws-123", + }, + "workspace_source": {"kind": "docker"}, + }, + ) + + assert response.status_code == 400 + assert response.json()["detail"] == ( + "execution_binding and workspace_source cannot be set together" + ) + + async def test_create_session_rejects_workspace_provisioning_when_disabled( + self, client, monkeypatch + ): + monkeypatch.setattr( + "coding_agent.ui.http_server._load_cloud_workspace_config", + lambda: {}, + ) + + response = await client.post( + "/sessions", + json={"workspace_source": {"kind": "docker"}}, + ) + + assert response.status_code == 400 + assert response.json()["detail"] == ( + "cloud workspace provisioning requires cloud_workspace.enabled=true" + ) + + async def test_create_session_rolls_back_provisioned_workspace_on_failure( + self, client, monkeypatch + ): + binding = CloudWorkspaceBinding( + workspace_url="docker://agent-ws-orphan/workspace", + workspace_id="ws-orphan", + ) + cleaned: list[CloudWorkspaceBinding] = [] + + monkeypatch.setattr( + "coding_agent.ui.http_server._load_cloud_workspace_config", + lambda: { + "enabled": True, + "provider": "docker", + "workspace_root": "/tmp/unused", + }, + ) + monkeypatch.setattr( + "coding_agent.ui.http_server.provision_cloud_binding_from_config", + lambda config, source: binding, + ) + + async def fail_create_session(**kwargs): + del kwargs + raise RuntimeError("session store unavailable") + + monkeypatch.setattr(session_manager, "create_session", fail_create_session) + monkeypatch.setattr( + "coding_agent.ui.http_server.cleanup_cloud_binding_from_config", + lambda config, target_binding: cleaned.append(target_binding), + ) + + response = await client.post( + "/sessions", + json={"workspace_source": {"kind": "docker"}}, + ) + + assert response.status_code == 500 + assert cleaned == [binding] + + async def test_create_session_rolls_back_provisioned_workspace_on_non_runtime_failure( + self, client, monkeypatch + ): + binding = CloudWorkspaceBinding( + workspace_url="docker://agent-ws-orphan/workspace", + workspace_id="ws-orphan-nonruntime", + ) + cleaned: list[CloudWorkspaceBinding] = [] + + monkeypatch.setattr( + "coding_agent.ui.http_server._load_cloud_workspace_config", + lambda: { + "enabled": True, + "provider": "docker", + "workspace_root": "/tmp/unused", + }, + ) + monkeypatch.setattr( + "coding_agent.ui.http_server.provision_cloud_binding_from_config", + lambda config, source: binding, + ) + + async def fail_create_session(**kwargs): + del kwargs + raise KeyError("owner store unavailable") + + monkeypatch.setattr(session_manager, "create_session", fail_create_session) + monkeypatch.setattr( + "coding_agent.ui.http_server.cleanup_cloud_binding_from_config", + lambda config, target_binding: cleaned.append(target_binding), + ) + + response = await client.post( + "/sessions", + json={"workspace_source": {"kind": "docker"}}, + ) + + assert response.status_code == 500 + assert cleaned == [binding] + + async def test_create_session_keeps_original_failure_when_rollback_cleanup_fails( + self, client, monkeypatch + ): + binding = CloudWorkspaceBinding( + workspace_url="docker://agent-ws-orphan/workspace", + workspace_id="ws-orphan-cleanup-fails", + ) + + monkeypatch.setattr( + "coding_agent.ui.http_server._load_cloud_workspace_config", + lambda: { + "enabled": True, + "provider": "docker", + "workspace_root": "/tmp/unused", + }, + ) + monkeypatch.setattr( + "coding_agent.ui.http_server.provision_cloud_binding_from_config", + lambda config, source: binding, + ) + + async def fail_create_session(**kwargs): + del kwargs + raise RuntimeError("session store unavailable") + + def fail_cleanup(target_binding): + del target_binding + raise RuntimeError("cleanup unavailable") + + monkeypatch.setattr(session_manager, "create_session", fail_create_session) + monkeypatch.setattr( + "coding_agent.ui.http_server._cleanup_provisioned_cloud_binding", + fail_cleanup, + ) + + response = await client.post( + "/sessions", + json={"workspace_source": {"kind": "docker"}}, + ) + + assert response.status_code == 500 + assert response.json()["detail"] == "session store unavailable" + + async def test_create_session_rolls_back_provisioned_workspace_on_cancellation( + self, client, monkeypatch + ): + binding = CloudWorkspaceBinding( + workspace_url="docker://agent-ws-cancelled/workspace", + workspace_id="ws-cancelled", + ) + cleaned: list[CloudWorkspaceBinding] = [] + + monkeypatch.setattr( + "coding_agent.ui.http_server._load_cloud_workspace_config", + lambda: { + "enabled": True, + "provider": "docker", + "workspace_root": "/tmp/unused", + }, + ) + monkeypatch.setattr( + "coding_agent.ui.http_server.provision_cloud_binding_from_config", + lambda config, source: binding, + ) + + async def cancel_create_session(**kwargs): + del kwargs + raise asyncio.CancelledError + + monkeypatch.setattr(session_manager, "create_session", cancel_create_session) + monkeypatch.setattr( + "coding_agent.ui.http_server.cleanup_cloud_binding_from_config", + lambda config, target_binding: cleaned.append(target_binding), + ) + + with pytest.raises(asyncio.CancelledError): + await client.post( + "/sessions", + json={"workspace_source": {"kind": "docker"}}, + ) + + assert cleaned == [binding] + + async def test_close_session_cleans_up_provisioned_workspace_on_delete( + self, client, monkeypatch, tmp_path + ): + cleaned: list[CloudWorkspaceBinding] = [] + + monkeypatch.setattr( + "coding_agent.environment.docker_workspace_provider._start_docker_workspace_container", + lambda provider_config, binding: None, + ) + monkeypatch.setattr( + "coding_agent.ui.http_server._load_cloud_workspace_config", + lambda: { + "enabled": True, + "provider": "docker", + "workspace_root": str(tmp_path), + "container_name_prefix": "agent-", + }, + ) + monkeypatch.setattr( + "coding_agent.ui.http_server.cleanup_cloud_binding_from_config", + lambda config, target_binding: cleaned.append(target_binding), + ) + + create_response = await client.post( + "/sessions", + json={"workspace_source": {"kind": "docker"}}, + ) + + assert create_response.status_code == 200 + binding = session_manager.get_session( + create_response.json()["session_id"] + ).execution_binding + assert isinstance(binding, CloudWorkspaceBinding) + + close_response = await client.delete( + f"/sessions/{create_response.json()['session_id']}" + ) + + assert close_response.status_code == 200 + assert cleaned == [binding] + + async def test_close_session_cleans_up_when_new_provisioning_is_disabled( + self, client, monkeypatch, tmp_path + ): + cleaned: list[CloudWorkspaceBinding] = [] + config_enabled = True + + monkeypatch.setattr( + "coding_agent.environment.docker_workspace_provider._start_docker_workspace_container", + lambda provider_config, binding: None, + ) + + def cloud_workspace_config(): + return { + "enabled": config_enabled, + "provider": "docker", + "workspace_root": str(tmp_path), + "container_name_prefix": "agent-", + } + + monkeypatch.setattr( + "coding_agent.ui.http_server._load_cloud_workspace_config", + cloud_workspace_config, + ) + monkeypatch.setattr( + "coding_agent.ui.http_server.cleanup_cloud_binding_from_config", + lambda config, target_binding: cleaned.append(target_binding), + ) + + create_response = await client.post( + "/sessions", + json={"workspace_source": {"kind": "docker"}}, + ) + binding = session_manager.get_session( + create_response.json()["session_id"] + ).execution_binding + assert isinstance(binding, CloudWorkspaceBinding) + config_enabled = False + + close_response = await client.delete( + f"/sessions/{create_response.json()['session_id']}" + ) + + assert close_response.status_code == 200 + assert cleaned == [binding] + async def test_create_session_accepts_runtime_provider_metadata(self, client): response = await client.post( "/sessions", @@ -332,6 +665,103 @@ async def test_create_session_rejects_invalid_runtime_provider(self, client): assert response.status_code == 422 + async def test_send_prompt_uses_cloud_environment_from_provisioned_workspace( + self, client, monkeypatch, tmp_path + ): + monkeypatch.setattr( + "coding_agent.environment.docker_workspace_provider._start_docker_workspace_container", + lambda provider_config, binding: None, + ) + + class _CreateAgentCapture: + def __init__(self) -> None: + self.environment: CloudEnvironment | None = None + self.session_id: str | None = None + self.workspace_root: object | None = None + + monkeypatch.setattr( + "coding_agent.ui.http_server._load_cloud_workspace_config", + lambda: { + "enabled": True, + "provider": "docker", + "workspace_root": str(tmp_path), + "container_name_prefix": "agent-", + }, + ) + monkeypatch.setattr( + session_manager, + "_binding_resolver", + _build_binding_resolver(), + ) + + captured = _CreateAgentCapture() + + class FakeAdapter: + def __init__(self, pipeline, ctx, consumer) -> None: + del pipeline, ctx + self._consumer = consumer + + async def run_turn(self, prompt: str) -> None: + del prompt + assert captured.session_id is not None + await self._consumer.emit( + TurnEnd( + session_id=captured.session_id, + agent_id="", + turn_id="turn-provisioned", + completion_status=CompletionStatus.COMPLETED, + ) + ) + + fake_pipeline = types.SimpleNamespace( + _registry=types.SimpleNamespace( + get=lambda _: types.SimpleNamespace(_instance=None) + ), + _directive_executor=None, + ) + + def fake_create_agent(**kwargs): + environment = kwargs.get("environment") + assert isinstance(environment, CloudEnvironment) + session_id_override = kwargs.get("session_id_override") + assert isinstance(session_id_override, str) + captured.environment = environment + captured.session_id = session_id_override + captured.workspace_root = kwargs.get("workspace_root") + return fake_pipeline, types.SimpleNamespace(config={}, tape=Tape()) + + monkeypatch.setattr(session_manager, "_create_agent", fake_create_agent) + monkeypatch.setattr( + "coding_agent.ui.session_manager.PipelineAdapter", FakeAdapter + ) + + create_response = await client.post( + "/sessions", + json={"workspace_source": {"kind": "docker"}}, + ) + session_id = create_response.json()["session_id"] + + events = [] + async with aconnect_sse( + client, + "POST", + f"/sessions/{session_id}/prompt", + json={"prompt": "Hello"}, + ) as event_source: + async for sse in event_source.aiter_sse(): + events.append(sse.event) + if sse.event == "TurnEnd": + break + + assert captured.environment is not None + binding = session_manager.get_session(session_id).execution_binding + assert isinstance(binding, CloudWorkspaceBinding) + assert ( + captured.environment.tool_config()["workspace_id"] == binding.workspace_id + ) + assert captured.workspace_root is None + assert events[-1] == "TurnEnd" + def test_build_session_manager_enables_owner_store_for_pg_http_sessions( self, monkeypatch ): @@ -628,7 +1058,13 @@ async def renew( new_fencing_token: int = 2, current_fencing_token: int = 1, ) -> bool: - del session_id, owner_id, lease_seconds, new_fencing_token, current_fencing_token + del ( + session_id, + owner_id, + lease_seconds, + new_fencing_token, + current_fencing_token, + ) return True async def release( @@ -676,7 +1112,13 @@ async def renew( new_fencing_token: int = 2, current_fencing_token: int = 1, ) -> bool: - del session_id, owner_id, lease_seconds, new_fencing_token, current_fencing_token + del ( + session_id, + owner_id, + lease_seconds, + new_fencing_token, + current_fencing_token, + ) return True async def release( @@ -832,7 +1274,9 @@ async def close(self) -> None: types.SimpleNamespace(config={}, tape=kwargs.get("tape") or Tape()), ), ) - monkeypatch.setattr("coding_agent.ui.session_manager.PipelineAdapter", FakeAdapter) + monkeypatch.setattr( + "coding_agent.ui.session_manager.PipelineAdapter", FakeAdapter + ) events = [] async with aconnect_sse( @@ -1826,9 +2270,7 @@ async def fake_close() -> None: assert close_calls == ["closed"] assert "Owner lease renewal task failed during shutdown" in caplog.text - async def test_lifespan_backfills_owner_leases_before_renewal( - self, monkeypatch - ): + async def test_lifespan_backfills_owner_leases_before_renewal(self, monkeypatch): events: list[str] = [] async def fake_cleanup_idle_sessions() -> None: @@ -2567,8 +3009,9 @@ async def test_wait_for_approval_request_can_be_approved_via_http_endpoint( wait_task = asyncio.create_task(wait_for_approval(session_id, req)) for _ in range(20): if ( - session_manager.get_session(session_id) - .approval_coordinator.get_request("req-http-wait") + session_manager.get_session( + session_id + ).approval_coordinator.get_request("req-http-wait") is not None ): break diff --git a/coding-agent/tests/ui/test_session_manager_public_api.py b/coding-agent/tests/ui/test_session_manager_public_api.py index fa49850aa..891f742e1 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -21,7 +21,7 @@ from coding_agent.approval.store import ApprovalStore from coding_agent.approval import ApprovalPolicy from agentkit.config.loader import ConfigError -from coding_agent.ui.execution_binding import LocalExecutionBinding +from coding_agent.ui.execution_binding import CloudWorkspaceBinding, LocalExecutionBinding from coding_agent.wire.protocol import ( ApprovalRequest, CompletionStatus, @@ -149,6 +149,38 @@ def test_clear_sessions_uses_public_api() -> None: assert manager.list_sessions() == [] +def test_clear_sessions_cleans_up_persisted_provisioned_cloud_sessions() -> None: + store = InMemorySessionStore() + cleaned: list[CloudWorkspaceBinding] = [] + manager = SessionManager( + store=store, + provisioned_cloud_binding_cleanup=cleaned.append, + ) + binding = CloudWorkspaceBinding( + workspace_url="docker://agent-ws-persisted/workspace", + workspace_id="ws-persisted", + ) + session = Session( + id="persisted-provisioned-session", + created_at=datetime.now(), + last_activity=datetime.now(), + approval_store=ApprovalStore(), + execution_binding=binding, + origin={ + "channel": "http", + "binding_kind": "cloud", + "workspace_source_kind": "docker", + }, + ) + manager.register_session(session) + manager._session_cache.clear() + + manager.clear_sessions() + + assert cleaned == [binding] + assert store.list_sessions() == [] + + @pytest.mark.asyncio async def test_create_session_persists_to_store_backing() -> None: store = InMemorySessionStore() @@ -225,6 +257,24 @@ def test_session_metadata_round_trips_local_execution_binding(tmp_path: Path) -> assert reloaded.execution_binding.workspace_root == str(bound_repo) +def test_session_metadata_round_trips_origin() -> None: + store = InMemorySessionStore() + manager = SessionManager(store=store) + session = Session( + id="origin-session", + created_at=datetime.now(), + last_activity=datetime.now(), + approval_store=ApprovalStore(), + origin={"channel": "http", "binding_kind": "cloud"}, + ) + + manager.register_session(session) + + reloaded = SessionManager(store=store).get_session("origin-session") + + assert reloaded.origin == {"channel": "http", "binding_kind": "cloud"} + + def test_session_defaults_local_execution_binding_from_repo_path( tmp_path: Path, ) -> None: @@ -270,6 +320,7 @@ def test_session_as_dict_includes_restart_metadata() -> None: session.model_name = "claude-test" session.base_url = "http://llm.local" session.max_steps = 12 + session.origin = {"channel": "http", "binding_kind": "local"} payload = session.as_dict() @@ -277,6 +328,7 @@ def test_session_as_dict_includes_restart_metadata() -> None: assert payload["model_name"] == "claude-test" assert payload["base_url"] == "http://llm.local" assert payload["max_steps"] == 12 + assert payload["origin"] == {"channel": "http", "binding_kind": "local"} def test_session_metadata_defaults_missing_binding_to_local_from_repo_path( diff --git a/coding-agent/tests/ui/test_session_manager_runtime.py b/coding-agent/tests/ui/test_session_manager_runtime.py index e79e8693e..0976ae08b 100644 --- a/coding-agent/tests/ui/test_session_manager_runtime.py +++ b/coding-agent/tests/ui/test_session_manager_runtime.py @@ -1526,6 +1526,116 @@ async def initialize(self) -> None: assert session.execution_binding.workspace_id == "ws-123" +@pytest.mark.asyncio +async def test_close_session_cleans_up_server_provisioned_cloud_binding() -> None: + cleaned: list[CloudWorkspaceBinding] = [] + manager = SessionManager( + store=InMemorySessionStore(), + provisioned_cloud_binding_cleanup=cleaned.append, + ) + binding = CloudWorkspaceBinding( + workspace_url="docker://agent-ws-123/workspace", + workspace_id="ws-123", + ) + session_id = await manager.create_session( + execution_binding=binding, + origin={ + "channel": "http", + "binding_kind": "cloud", + "workspace_source_kind": "docker", + }, + ) + + await manager.close_session(session_id) + + assert cleaned == [binding] + + +@pytest.mark.asyncio +async def test_close_session_offloads_provisioned_cloud_cleanup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + manager = SessionManager( + store=InMemorySessionStore(), + provisioned_cloud_binding_cleanup=lambda binding: None, + ) + binding = CloudWorkspaceBinding( + workspace_url="docker://agent-ws-threaded/workspace", + workspace_id="ws-threaded", + ) + session_id = await manager.create_session( + execution_binding=binding, + origin={ + "channel": "http", + "binding_kind": "cloud", + "workspace_source_kind": "docker", + }, + ) + session = manager.get_session(session_id) + to_thread_calls: list[tuple[object, tuple[object, ...]]] = [] + + async def fake_to_thread(func, *args): + to_thread_calls.append((func, args)) + return func(*args) + + monkeypatch.setattr("coding_agent.ui.session_manager.asyncio.to_thread", fake_to_thread) + + await manager.close_session(session_id) + + assert to_thread_calls == [(manager._cleanup_provisioned_cloud_binding, (session,))] + + +@pytest.mark.asyncio +async def test_close_session_logs_provisioned_cloud_cleanup_failure( + caplog: pytest.LogCaptureFixture, +) -> None: + def fail_cleanup(binding: CloudWorkspaceBinding) -> None: + del binding + raise RuntimeError("cleanup failed") + + manager = SessionManager( + store=InMemorySessionStore(), + provisioned_cloud_binding_cleanup=fail_cleanup, + ) + session_id = await manager.create_session( + execution_binding=CloudWorkspaceBinding( + workspace_url="docker://agent-ws-fails/workspace", + workspace_id="ws-fails", + ), + origin={ + "channel": "http", + "binding_kind": "cloud", + "workspace_source_kind": "docker", + }, + ) + + with caplog.at_level("ERROR"): + await manager.close_session(session_id) + + assert not manager.has_session(session_id) + assert "Failed to clean up provisioned cloud workspace" in caplog.text + + +@pytest.mark.asyncio +async def test_close_session_keeps_explicit_cloud_binding_untouched() -> None: + cleaned: list[CloudWorkspaceBinding] = [] + manager = SessionManager( + store=InMemorySessionStore(), + provisioned_cloud_binding_cleanup=cleaned.append, + ) + session_id = await manager.create_session( + execution_binding=CloudWorkspaceBinding( + workspace_url="https://workspace.example.com", + workspace_id="ws-explicit", + ), + origin={"channel": "http", "binding_kind": "cloud"}, + ) + + await manager.close_session(session_id) + + assert cleaned == [] + + @pytest.mark.asyncio async def test_restore_legacy_checkpoint_without_session_config_uses_current_session_metadata() -> ( None