Merge pull request #36010 from kshitijk4poor/fix/terminal-cwd-acp-aware

fix(tools): preserve live session cwd in terminal_tool, keep ACP update_cwd authoritative

tests/tools/test_terminal_task_cwd.py

@@ -1,6 +1,7 @@
 """Regression tests for task/session cwd propagation in terminal_tool."""
 
 import json
+from types import SimpleNamespace
 
 import tools.terminal_tool as terminal_tool
 
@@ -10,6 +11,7 @@ def _minimal_terminal_config(cwd="/default"):
         "env_type": "local",
         "cwd": cwd,
         "timeout": 60,
+        "lifetime_seconds": 3600,
     }
 
 
@@ -72,3 +74,147 @@ def test_explicit_workdir_still_wins_over_registered_task_cwd(monkeypatch):
 
     assert result["exit_code"] == 0
     assert calls == [{"timeout": 60, "cwd": "/explicit/workdir"}]
+
+
+def test_foreground_command_prefers_live_env_cwd_over_init_time_cwd(monkeypatch):
+    """A prior `cd` updates env.cwd; terminal_tool must honor that live cwd."""
+    calls = []
+
+    class FakeEnv:
+        env = {}
+        cwd = "/workspace/live"
+
+        def execute(self, command, **kwargs):
+            calls.append((command, kwargs))
+            return {"output": "ok", "returncode": 0}
+
+    task_id = "session-live-cwd"
+    monkeypatch.setattr(terminal_tool, "_active_environments", {task_id: FakeEnv()})
+    monkeypatch.setattr(terminal_tool, "_last_activity", {})
+    monkeypatch.setattr(terminal_tool, "_task_env_overrides", {task_id: {"cwd": "/workspace/init"}})
+    monkeypatch.setattr(terminal_tool, "_get_env_config", lambda: _minimal_terminal_config(cwd="/workspace/init"))
+    monkeypatch.setattr(terminal_tool, "_start_cleanup_thread", lambda: None)
+    monkeypatch.setattr(terminal_tool, "_resolve_container_task_id", lambda value: value or "default")
+    monkeypatch.setattr(
+        terminal_tool,
+        "_check_all_guards",
+        lambda command, env_type: {"approved": True},
+    )
+
+    result = json.loads(terminal_tool.terminal_tool(command="pwd", task_id=task_id))
+
+    assert result["exit_code"] == 0
+    assert calls == [("pwd", {"timeout": 60, "cwd": "/workspace/live"})]
+
+
+def test_background_command_prefers_live_env_cwd_over_init_time_cwd(monkeypatch):
+    """Background process launches must also use the live session cwd."""
+
+    class FakeEnv:
+        env = {}
+        cwd = "/workspace/live"
+
+    class FakeRegistry:
+        def __init__(self):
+            self.calls = []
+            self.pending_watchers = []
+
+        def spawn_local(self, **kwargs):
+            self.calls.append(kwargs)
+            return SimpleNamespace(id="proc_test", pid=1234)
+
+    import tools.process_registry as process_registry_mod
+
+    registry = FakeRegistry()
+    task_id = "session-live-cwd-bg"
+    monkeypatch.setattr(terminal_tool, "_active_environments", {task_id: FakeEnv()})
+    monkeypatch.setattr(terminal_tool, "_last_activity", {})
+    monkeypatch.setattr(terminal_tool, "_task_env_overrides", {task_id: {"cwd": "/workspace/init"}})
+    monkeypatch.setattr(terminal_tool, "_get_env_config", lambda: _minimal_terminal_config(cwd="/workspace/init"))
+    monkeypatch.setattr(terminal_tool, "_start_cleanup_thread", lambda: None)
+    monkeypatch.setattr(terminal_tool, "_resolve_container_task_id", lambda value: value or "default")
+    monkeypatch.setattr(
+        terminal_tool,
+        "_check_all_guards",
+        lambda command, env_type: {"approved": True},
+    )
+    monkeypatch.setattr(process_registry_mod, "process_registry", registry)
+
+    result = json.loads(
+        terminal_tool.terminal_tool(
+            command="sleep 1",
+            task_id=task_id,
+            background=True,
+        )
+    )
+
+    assert result["exit_code"] == 0
+    assert registry.calls == [{
+        "command": "sleep 1",
+        "cwd": "/workspace/live",
+        "task_id": task_id,
+        "session_key": "",
+        "env_vars": {},
+        "use_pty": False,
+    }]
+
+
+def test_registering_cwd_override_updates_live_env_cwd(monkeypatch):
+    """An ACP ``update_cwd`` (re-)registered mid-session must win over a
+    previously ``cd``-ed live ``env.cwd``.
+
+    Preferring live ``env.cwd`` (so session-local ``cd`` survives) means a
+    freshly registered ``cwd`` override would otherwise sit *below* the
+    already-set ``env.cwd`` and be silently ignored. ``register_task_env_overrides``
+    syncs the new cwd onto the live cached env so an explicit ACP project-root
+    change takes effect, as the editor client expects.
+    """
+
+    class FakeEnv:
+        env = {}
+        cwd = "/workspace/old"
+
+    task_id = "acp-session-update"
+    fake_env = FakeEnv()
+    monkeypatch.setattr(terminal_tool, "_active_environments", {task_id: fake_env})
+    monkeypatch.setattr(terminal_tool, "_task_env_overrides", {})
+
+    terminal_tool.register_task_env_overrides(task_id, {"cwd": "/workspace/new"})
+
+    # The live env now reflects the editor's new project root.
+    assert fake_env.cwd == "/workspace/new"
+
+    # A subsequent command resolves to the new cwd (env.cwd precedence).
+    assert terminal_tool._resolve_command_cwd(
+        workdir=None, env=fake_env, default_cwd="/workspace/config"
+    ) == "/workspace/new"
+
+
+def test_registering_cwd_override_noop_when_no_live_env(monkeypatch):
+    """Registering an override before the env exists must not crash; the cwd
+    is applied at env creation time instead."""
+    monkeypatch.setattr(terminal_tool, "_active_environments", {})
+    monkeypatch.setattr(terminal_tool, "_task_env_overrides", {})
+
+    # Should not raise even though no env is cached yet.
+    terminal_tool.register_task_env_overrides("acp-session-pending", {"cwd": "/workspace/new"})
+
+    assert terminal_tool._task_env_overrides["acp-session-pending"] == {"cwd": "/workspace/new"}
+
+
+def test_registering_non_cwd_override_leaves_live_env_cwd_untouched(monkeypatch):
+    """A non-cwd override (e.g. a per-task Modal image) must not disturb the
+    live env's cwd."""
+
+    class FakeEnv:
+        env = {}
+        cwd = "/workspace/keep"
+
+    task_id = "rl-rollout-1"
+    fake_env = FakeEnv()
+    monkeypatch.setattr(terminal_tool, "_active_environments", {task_id: fake_env})
+    monkeypatch.setattr(terminal_tool, "_task_env_overrides", {})
+
+    terminal_tool.register_task_env_overrides(task_id, {"modal_image": "custom:latest"})
+
+    assert fake_env.cwd == "/workspace/keep"

tools/terminal_tool.py

@@ -962,6 +962,23 @@ def register_task_env_overrides(task_id: str, overrides: Dict[str, Any]):
     """
     _task_env_overrides[task_id] = overrides
 
+    # If a live environment already exists for this task, a freshly registered
+    # ``cwd`` override (e.g. the ACP client switching the editor's project root
+    # mid-session via ``session/load`` / ``session/resume``) must take effect on
+    # the cached env too. ``terminal_tool`` resolves the per-command cwd as
+    # ``workdir > env.cwd > config/override cwd`` so that ordinary in-session
+    # ``cd`` state is preserved; without syncing here the override would sit
+    # below the (already-set) ``env.cwd`` and be silently ignored once any
+    # command has run. Pushing it onto the live env keeps ``cd`` tracking intact
+    # while letting an explicit ACP cwd change win, as the client expects.
+    new_cwd = overrides.get("cwd")
+    if isinstance(new_cwd, str) and new_cwd.strip():
+        container_id = _resolve_container_task_id(task_id)
+        with _env_lock:
+            env = _active_environments.get(container_id)
+        if env is not None and getattr(env, "cwd", None) is not None:
+            env.cwd = new_cwd
+
 
 def clear_task_env_overrides(task_id: str):
     """
@@ -1718,6 +1735,30 @@ def _resolve_notification_flag_conflict(
     return watch_patterns, ""
 
 
+def _resolve_command_cwd(
+    *,
+    workdir: Optional[str],
+    env: Any,
+    default_cwd: str,
+) -> str:
+    """Return the cwd for a command, preferring the live session cwd.
+
+    ``terminal_tool`` historically re-sent the init-time/config cwd on every
+    call. That broke session-local ``cd`` state: the environment tracked the
+    new directory in ``env.cwd``, but foreground/background calls kept forcing
+    the old cwd back through ``env.execute(..., cwd=...)``. Explicit
+    ``workdir=`` must still override everything.
+    """
+    if workdir:
+        return workdir
+
+    live_cwd = getattr(env, "cwd", None)
+    if isinstance(live_cwd, str) and live_cwd.strip():
+        return live_cwd
+
+    return default_cwd
+
+
 def terminal_tool(
     command: str,
     background: bool = False,
@@ -1990,7 +2031,11 @@ def terminal_tool(
             from tools.process_registry import process_registry
 
             session_key = get_current_session_key(default="")
-            effective_cwd = workdir or cwd
+            effective_cwd = _resolve_command_cwd(
+                workdir=workdir,
+                env=env,
+                default_cwd=cwd,
+            )
             try:
                 if env_type == "local":
                     proc_session = process_registry.spawn_local(
@@ -2207,7 +2252,11 @@ def terminal_tool(
                 try:
                     execute_kwargs = {
                         "timeout": effective_timeout,
-                        "cwd": workdir or cwd,
+                        "cwd": _resolve_command_cwd(
+                            workdir=workdir,
+                            env=env,
+                            default_cwd=cwd,
+                        ),
                     }
                     result = env.execute(command, **execute_kwargs)
                 except Exception as e: