feat(cloud): harden workspace shell parity

coding-agent/src/coding_agent/app.py

@@ -52,6 +52,25 @@ def _legacy_config_agent_id(run_context: AgentRunContext) -> str:
     return legacy_agent_id_str(run_context.agent_id)
 
 
+def _run_trace_metadata(
+    environment: Environment,
+    trace_metadata: Mapping[str, object] | None,
+) -> dict[str, object]:
+    metadata = dict(trace_metadata or {})
+    for key in list(metadata):
+        if key.startswith("cloud."):
+            del metadata[key]
+
+    if environment.kind != "cloud":
+        return metadata
+
+    workspace_id = environment.tool_config().get("workspace_id")
+    if not isinstance(workspace_id, str) or not workspace_id:
+        raise ValueError("cloud environment must expose string workspace_id")
+    metadata["cloud.workspace_id"] = workspace_id
+    return metadata
+
+
 def _child_system_prompt_suffix(tool_filter: ToolFilter) -> str:
     if tool_filter is None:
         return ""
@@ -335,7 +354,7 @@ async def _memory_handler(directive: Any) -> None:
         parent_run_id=parent_run_id_override,
         environment=environment,
         context_budget=ContextBudget() if context_budget is None else context_budget,
-        trace_metadata={} if trace_metadata is None else trace_metadata,
+        trace_metadata=_run_trace_metadata(environment, trace_metadata),
     )
 
     ctx_config: dict[str, Any] = {

coding-agent/src/coding_agent/environment/cloud.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import json
+import posixpath
+import shlex
 from dataclasses import dataclass
 from typing import Any, Protocol
 
@@ -46,7 +48,7 @@ class CloudEnvironment:
     """Provider-neutral cloud workspace environment backed by a client."""
 
     def __init__(self, client: CloudWorkspaceClient) -> None:
-        self._client = client
+        self._client: CloudWorkspaceClient = client
 
     @property
     def kind(self) -> str:
@@ -113,6 +115,41 @@ def file_patch(path: str, patch: str) -> str:
 
         return file_patch
 
+    def _resolve_workspace_path(self, target: str, *, cwd: str | None = None) -> str:
+        base = cwd or self._client.default_cwd
+        resolved = target if target.startswith("/") else posixpath.join(base, target)
+        normalized = posixpath.normpath(resolved)
+        workspace_root = posixpath.normpath(self._client.default_cwd)
+        if normalized != workspace_root and not normalized.startswith(
+            workspace_root.rstrip("/") + "/"
+        ):
+            raise ValueError(f"Directory is outside cloud workspace: {normalized}")
+        return normalized
+
+    def _validate_command_cwd(self, cwd: str | None) -> str | None:
+        if cwd is None:
+            return None
+        return self._resolve_workspace_path(cwd)
+
+    def _apply_cd(self, command: str, cwd: str | None) -> str | None:
+        args = shlex.split(command)
+        if not args or args[0] != "cd":
+            return None
+        if len(args) != 2:
+            raise ValueError("cd requires exactly one target directory")
+        return self._resolve_workspace_path(args[1], cwd=cwd)
+
+    def _apply_export(self, command: str) -> tuple[str, str] | None:
+        args = shlex.split(command)
+        if not args or args[0] != "export":
+            return None
+        if len(args) != 2 or "=" not in args[1]:
+            raise ValueError("export requires KEY=VALUE")
+        key, value = args[1].split("=", 1)
+        if not key:
+            raise ValueError("export requires a non-empty variable name")
+        return key, value
+
     def build_shell_tool(self):
         @tool(description="Run a shell command and return stdout/stderr.")
         def bash_run(
@@ -121,12 +158,29 @@ def bash_run(
             cwd: str | None = None,
             env: dict[str, str] | None = None,
         ) -> str:
-            result = self._client.run_command(
-                command,
-                cwd=cwd,
-                env=env,
-                timeout=timeout,
-            )
+            try:
+                changed_dir = self._apply_cd(command, cwd)
+                if changed_dir is not None:
+                    return f"Changed directory to {changed_dir}"
+                exported = self._apply_export(command)
+                if exported is not None:
+                    key, value = exported
+                    return f"Exported {key}={value}"
+            except ValueError as exc:
+                return f"Error: {exc}"
+
+            try:
+                execution_cwd = self._validate_command_cwd(cwd)
+                result = self._client.run_command(
+                    command,
+                    cwd=execution_cwd,
+                    env=env,
+                    timeout=timeout,
+                )
+            except ValueError as exc:
+                return f"Error: {exc}"
+            except TimeoutError:
+                return f"Error: command timed out after {timeout}s"
             output = ""
             if result.stdout:
                 output += result.stdout

coding-agent/src/coding_agent/ui/binding_resolver.py

@@ -17,8 +17,6 @@
 
 
 class BindingResolver(Protocol):
-    def resolve_workspace_root(self, binding: ExecutionBinding) -> Path: ...
-
     def resolve_tool_config(self, binding: ExecutionBinding) -> dict[str, object]: ...
 
     def resolve_environment(self, binding: ExecutionBinding) -> Environment: ...

coding-agent/src/coding_agent/ui/session_manager.py

@@ -978,7 +978,8 @@ async def _restore_checkpoint(self, session: Session, checkpoint_id: str) -> Non
         previous_provider_name = session.provider_name
         previous_model_name = session.model_name
         previous_base_url = session.base_url
-        workspace_root, environment = self._resolve_runtime_environment(session)
+        environment = self._resolve_environment(session)
+        workspace_root = self._environment_workspace_root(environment)
 
         approval_mode_map = {
             ApprovalPolicy.YOLO: "yolo",
@@ -1051,9 +1052,6 @@ def _persist_session(self, session: Session) -> None:
         self._session_cache[session.id] = session
         self._store.save(session.id, cast(dict[str, Any], session.to_store_data()))
 
-    def _resolve_workspace_root(self, session: Session) -> Path:
-        return self._binding_resolver.resolve_workspace_root(session.execution_binding)
-
     def _resolve_environment(self, session: Session) -> Environment:
         return self._binding_resolver.resolve_environment(session.execution_binding)
 
@@ -1063,13 +1061,6 @@ def _environment_workspace_root(self, environment: Environment) -> Path | None:
             return None
         return Path(local_root).expanduser().resolve()
 
-    def _resolve_runtime_environment(
-        self,
-        session: Session,
-    ) -> tuple[Path | None, Environment]:
-        environment = self._resolve_environment(session)
-        return self._environment_workspace_root(environment), environment
-
     def _invalidate_runtime(self, session: Session) -> None:
         session.runtime_pipeline = None
         session.runtime_ctx = None
@@ -1374,9 +1365,8 @@ async def run_agent(
                 adapter = session.runtime_adapter
 
                 if pipeline is None or ctx is None or adapter is None:
-                    workspace_root, environment = self._resolve_runtime_environment(
-                        session
-                    )
+                    environment = self._resolve_environment(session)
+                    workspace_root = self._environment_workspace_root(environment)
                     pipeline, ctx = self._create_agent_for_session(
                         workspace_root=workspace_root,
                         environment=environment,
@@ -1841,7 +1831,8 @@ async def _build_session_runtime(
         resolved_approval_policy = (
             session.approval_policy if approval_policy is None else approval_policy
         )
-        workspace_root, environment = self._resolve_runtime_environment(session)
+        environment = self._resolve_environment(session)
+        workspace_root = self._environment_workspace_root(environment)
 
         consumer = self._make_session_consumer(session)
         pipeline, ctx = self._create_agent_for_session(

coding-agent/tests/coding_agent/environment/test_cloud_environment.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import json
-from typing import Any
+from pathlib import Path
+from typing import Any, cast
 
 import pytest
 
@@ -10,17 +11,19 @@
 from agentkit.tools import ToolCallRequest, Toolset
 from coding_agent.environment import CloudCommandResult, CloudEnvironment
 from coding_agent.plugins.core_tools import CoreToolsPlugin
+from coding_agent.plugins.shell_session import ShellSessionPlugin
 
 
 class FakeCloudWorkspaceClient:
-    workspace_id = "ws-123"
-    workspace_url = "https://workspace.example.com"
-    default_cwd = "/workspace"
+    workspace_id: str = "ws-123"
+    workspace_url: str = "https://workspace.example.com"
+    default_cwd: str = "/workspace"
 
     def __init__(self) -> None:
         self.files: dict[str, str] = {"note.txt": "hello cloud"}
         self.shell_calls: list[dict[str, Any]] = []
-        self.fail_read = False
+        self.fail_read: bool = False
+        self.command_timeout: bool = False
 
     def read_file(self, path: str) -> str:
         if self.fail_read:
@@ -57,6 +60,8 @@ def run_command(
         env: dict[str, str] | None,
         timeout: int,
     ) -> CloudCommandResult:
+        if self.command_timeout:
+            raise TimeoutError("cloud command timed out")
         self.shell_calls.append(
             {"command": command, "cwd": cwd, "env": env, "timeout": timeout}
         )
@@ -117,9 +122,142 @@ def test_cloud_environment_shell_tool_preserves_cwd_and_env() -> None:
     ]
 
 
+def test_cloud_environment_shell_tool_updates_session_cwd_and_env() -> None:
+    client = FakeCloudWorkspaceClient()
+    shell_session = ShellSessionPlugin()
+    _ = shell_session.do_mount(
+        ctx=type("Ctx", (), {"config": {"environment": CloudEnvironment(client)}})()
+    )
+    plugin = CoreToolsPlugin(
+        environment=CloudEnvironment(client),
+        shell_session=shell_session,
+    )
+
+    cd_result = cast(
+        str,
+        plugin.execute_tool(
+            name="bash_run",
+            arguments={"command": "cd pkg"},
+        ),
+    )
+    export_result = cast(
+        str,
+        plugin.execute_tool(
+            name="bash_run",
+            arguments={"command": 'export TEST_VALUE="cloud ok"'},
+        ),
+    )
+    command_result = cast(
+        str,
+        plugin.execute_tool(
+            name="bash_run",
+            arguments={"command": "pytest -q", "timeout": 7},
+        ),
+    )
+
+    assert cd_result == "Changed directory to /workspace/pkg"
+    assert export_result == "Exported TEST_VALUE=cloud ok"
+    assert command_result == "cloud output"
+    assert shell_session.get_session_context() == {
+        "cwd": "/workspace/pkg",
+        "env_vars": {"TEST_VALUE": "cloud ok"},
+        "active": True,
+    }
+    assert client.shell_calls == [
+        {
+            "command": "pytest -q",
+            "cwd": "/workspace/pkg",
+            "env": {"TEST_VALUE": "cloud ok"},
+            "timeout": 7,
+        }
+    ]
+
+
+def test_cloud_environment_rejects_invalid_session_commands_without_mutation() -> None:
+    client = FakeCloudWorkspaceClient()
+    shell_session = ShellSessionPlugin()
+    _ = shell_session.do_mount(
+        ctx=type("Ctx", (), {"config": {"environment": CloudEnvironment(client)}})()
+    )
+    plugin = CoreToolsPlugin(
+        environment=CloudEnvironment(client),
+        shell_session=shell_session,
+    )
+
+    cd_result = cast(
+        str,
+        plugin.execute_tool(
+            name="bash_run",
+            arguments={"command": "cd one two"},
+        ),
+    )
+    export_result = cast(
+        str,
+        plugin.execute_tool(
+            name="bash_run",
+            arguments={"command": "export MISSING_VALUE"},
+        ),
+    )
+
+    assert cd_result == "Error: cd requires exactly one target directory"
+    assert export_result == "Error: export requires KEY=VALUE"
+    assert shell_session.get_session_context() == {
+        "cwd": "/workspace",
+        "env_vars": {},
+        "active": True,
+    }
+    assert client.shell_calls == []
+
+
+def test_cloud_environment_export_does_not_mutate_input_env() -> None:
+    client = FakeCloudWorkspaceClient()
+    shell_tool = CloudEnvironment(client).build_shell_tool()
+    caller_env = {"EXISTING": "1"}
+
+    result = shell_tool("export NEW_VALUE=2", env=caller_env)
+
+    assert result == "Exported NEW_VALUE=2"
+    assert caller_env == {"EXISTING": "1"}
+    assert client.shell_calls == []
+
+
+def test_cloud_environment_rejects_explicit_cwd_outside_workspace() -> None:
+    client = FakeCloudWorkspaceClient()
+    shell_tool = CloudEnvironment(client).build_shell_tool()
+
+    absolute_result = shell_tool("pwd", cwd="/etc")
+    relative_result = shell_tool("pwd", cwd="/workspace/../etc")
+
+    assert absolute_result == "Error: Directory is outside cloud workspace: /etc"
+    assert relative_result == "Error: Directory is outside cloud workspace: /etc"
+    assert client.shell_calls == []
+
+
+def test_cloud_environment_rejects_cd_and_explicit_cwd_with_same_prefix() -> None:
+    client = FakeCloudWorkspaceClient()
+    shell_tool = CloudEnvironment(client).build_shell_tool()
+
+    cd_result = shell_tool("cd /etc", cwd="/workspace")
+    cwd_result = shell_tool("pwd", cwd="/etc")
+
+    assert cd_result == "Error: Directory is outside cloud workspace: /etc"
+    assert cwd_result == "Error: Directory is outside cloud workspace: /etc"
+    assert client.shell_calls == []
+
+
+def test_cloud_environment_shell_timeout_is_model_visible() -> None:
+    client = FakeCloudWorkspaceClient()
+    client.command_timeout = True
+    shell_tool = CloudEnvironment(client).build_shell_tool()
+
+    result = shell_tool("sleep 10", timeout=7, cwd="/workspace", env={})
+
+    assert result == "Error: command timed out after 7s"
+
+
 @pytest.mark.asyncio
 async def test_cloud_environment_tool_errors_do_not_fallback_to_local_execution(
-    tmp_path,
+    tmp_path: Path,
 ) -> None:
     local_note = tmp_path / "note.txt"
     local_note.write_text("local content")