Preserve Stirrup code-execution workspaces for scenario-suite runs

docs/stirrup-agent.md

@@ -186,6 +186,8 @@ export DOCKER_HOST=unix:///Users/<you>/.rd/docker.sock
 export STIRRUP_CODE_IMAGE=assetops-code
 
 uv run stirrup-agent --code-backend docker --show-trajectory \
+  --workspace-dir /tmp/assetopsbench-stirrup/smoke \
+  --preserve-workspace \
   --model-id watsonx/meta-llama/llama-4-maverick-17b-128e-instruct-fp8 \
   "Run python to compute the factorial of 12 and tell me the result"
 ```
@@ -195,25 +197,28 @@ call and the right answer (479001600).
 
 ---
 
-## Reading tool-produced files
+## Preserving code workspaces
 
-The Docker sandbox has its own filesystem: it can only see what is mounted, which
-is a private working directory — **not** wherever a tool downloaded a file on the
-host. So if a tool (e.g. `iot__history`) writes a file on the host and returns its
-path, code running in the Docker backend cannot open that path.
+By default, Stirrup creates a temporary execution directory for `code_exec` and
+removes it when the agent session exits. Pass `--workspace-dir` to choose the host
+base directory for that sandbox, and pass `--preserve-workspace` to copy the final
+code-execution files back into that directory before cleanup.
 
-Current guidance:
+For scenario-suite runs, use a root outside the repo:
 
-- For code scenarios that must read a tool-produced **host file**, use
-  `--code-backend local`. Local execution runs on the host, so the returned path
-  resolves directly.
-- Use the `docker` backend for computation that does not need to open a
-  tool-produced host file (the data the code needs is in the conversation), where
-  isolation matters more than direct file access.
+```bash
+uv run python -m benchmark.scenario_suite_runner \
+  --scenario-root /path/to/scenarios_data \
+  --agent_name stirrup_agent \
+  --stirrup-workspace-root /tmp/assetopsbench-stirrup-workspaces \
+  --preserve-workspaces
+```
 
-> Mounting a host directory into the Docker sandbox is not currently supported by
-> Stirrup's stock `DockerCodeExecToolProvider`. If that lands upstream
-> (an `extra_mounts` option), the Docker backend can read host files directly.
+For scenario `401`, the preserved files are available under a nested path such as:
+
+```text
+/tmp/assetopsbench-stirrup-workspaces/stirrup_agent/tokenrouter-MiniMax-M3/stirrup_agent_401/
+```
 
 ---
 
@@ -229,6 +234,8 @@ In addition to the [common flags](../INSTRUCTIONS.md#common-flags) (`--model-id`
 | `--code-backend`      | `docker` (default), `local`, or `e2b`.                                               |
 | `--max-turns N`       | Max agent turns (default: 30).                                                       |
 | `--max-tokens N`      | Max output tokens per model call; keep under the provider limit (default: 16384).    |
+| `--workspace-dir PATH` | Host base directory for Docker/local code-execution workspaces.                     |
+| `--preserve-workspace` | Copy final code-execution files into `--workspace-dir` before cleanup.              |
 
 Environment variable: `STIRRUP_CODE_IMAGE` (Docker image; default `python:3.12-slim`).
 
@@ -311,4 +318,4 @@ uv run stirrup-agent --no-code --run-id stirrup-smoke --scenario-id 101 \
 - `pyproject.toml` — `stirrup[mcp,litellm,docker]` dependency and the `stirrup-agent`
   entry point.
 
-No other parts of the repo are modified; the MCP servers are unchanged.
+No other parts of the repo are modified; the MCP servers are unchanged.

src/agent/stirrup_agent/cli.py

@@ -11,6 +11,7 @@
 from __future__ import annotations
 
 import argparse
+from pathlib import Path
 
 from .._cli_common import add_common_args, print_result, run_sdk_cli
 
@@ -83,6 +84,24 @@ def _build_parser() -> argparse.ArgumentParser:
         help="Max output tokens per model call; must stay under the provider "
         "limit (watsonx caps new tokens at 100k). Default: 16384.",
     )
+    parser.add_argument(
+        "--workspace-dir",
+        type=Path,
+        default=None,
+        metavar="PATH",
+        help=(
+            "Host directory used as the Stirrup code-execution workspace base. "
+            "Docker/local backends create their execution directory under this path."
+        ),
+    )
+    parser.add_argument(
+        "--preserve-workspace",
+        action="store_true",
+        help=(
+            "Copy final code-execution files into --workspace-dir before cleanup. "
+            "Supported with docker/local backends."
+        ),
+    )
     return parser
 
 
@@ -93,6 +112,8 @@ async def _run(args: argparse.Namespace) -> None:
         model=args.model_id,
         code_enabled=args.code_enabled,
         code_backend=args.code_backend,
+        workspace_dir=args.workspace_dir,
+        preserve_workspace=args.preserve_workspace,
         max_turns=args.max_turns,
         max_tokens=args.max_tokens,
     )

src/agent/stirrup_agent/runner.py

@@ -32,6 +32,7 @@
 import datetime as _dt
 import logging
 import os
+import shutil
 import time
 from pathlib import Path
 
@@ -51,6 +52,48 @@
 _DEFAULT_CODE_IMAGE = os.environ.get("STIRRUP_CODE_IMAGE", "python:3.12-slim")
 
 
+def _copy_workspace_contents(source: Path, destination: Path) -> None:
+    """Copy code-exec workspace contents out of Stirrup's temp child dir."""
+    destination.mkdir(parents=True, exist_ok=True)
+    for item in source.iterdir():
+        target = destination / item.name
+        if item.is_dir():
+            shutil.copytree(item, target, dirs_exist_ok=True)
+        else:
+            target.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy2(item, target)
+
+
+def _preserving_provider_class(provider_cls):
+    class _PreservingCodeExecToolProvider(provider_cls):
+        def __init__(self, *args, preserve_dir: Path, **kwargs) -> None:
+            super().__init__(*args, **kwargs)
+            self._assetops_preserve_dir = preserve_dir
+
+        async def __aexit__(self, exc_type, exc_val, exc_tb):
+            temp_dir = self.temp_dir
+            if temp_dir is not None and temp_dir.exists():
+                try:
+                    fix_ownership = getattr(self, "_fix_file_ownership", None)
+                    if fix_ownership is not None:
+                        await fix_ownership()
+                    _copy_workspace_contents(temp_dir, self._assetops_preserve_dir)
+                    _log.info(
+                        "Preserved Stirrup code workspace %s -> %s",
+                        temp_dir,
+                        self._assetops_preserve_dir,
+                    )
+                except Exception:
+                    _log.warning(
+                        "Failed to preserve Stirrup code workspace %s",
+                        temp_dir,
+                        exc_info=True,
+                    )
+            return await super().__aexit__(exc_type, exc_val, exc_tb)
+
+    return _PreservingCodeExecToolProvider
+
+
 class StirrupAgentRunner(AgentRunner):
     """Run a question through a Stirrup agent against the MCP servers.
 
@@ -60,6 +103,8 @@ class StirrupAgentRunner(AgentRunner):
         model: ``litellm_proxy/<provider>/<model>`` or native ``<provider>/<model>``.
         code_enabled: Add a sandboxed code-execution tool (the code track).
         code_backend: ``"docker"`` (sandboxed, default), ``"local"``, or ``"e2b"``.
+        workspace_dir: Optional host base directory for Docker/local code execution.
+        preserve_workspace: Copy final code-execution files back into ``workspace_dir``.
         max_turns: Stirrup agent loop bound.
         max_tokens: Context window hint passed to the client.
     """
@@ -71,13 +116,27 @@ def __init__(
         model: str = _DEFAULT_MODEL,
         code_enabled: bool = True,
         code_backend: str = "docker",
+        workspace_dir: Path | str | None = None,
+        preserve_workspace: bool = False,
         max_turns: int = 30,
         max_tokens: int = 16_384,
     ) -> None:
         super().__init__(llm, server_paths)
         self._model_id = model
         self._code_enabled = code_enabled
         self._code_backend = code_backend
+        self._workspace_dir = (
+            Path(workspace_dir).expanduser().resolve()
+            if workspace_dir is not None
+            else None
+        )
+        if preserve_workspace and self._workspace_dir is None:
+            raise ValueError("workspace_dir is required when preserve_workspace is enabled")
+        if preserve_workspace and code_backend not in {"docker", "local"}:
+            raise ValueError(
+                "preserve_workspace is only supported with docker or local code backends"
+            )
+        self._preserve_workspace = preserve_workspace
         self._max_turns = max_turns
         self._max_tokens = max_tokens
 
@@ -123,14 +182,30 @@ def _build_code_provider(self):
         if self._code_backend == "local":
             from stirrup.tools.code_backends.local import LocalCodeExecToolProvider
 
-            return LocalCodeExecToolProvider()
+            provider_cls = LocalCodeExecToolProvider
+            kwargs = {"temp_base_dir": self._workspace_dir}
+            if self._preserve_workspace:
+                provider_cls = _preserving_provider_class(provider_cls)
+                kwargs["preserve_dir"] = self._workspace_dir
+            return provider_cls(**kwargs)
         if self._code_backend == "e2b":
             from stirrup.tools.code_backends.e2b import E2BCodeExecToolProvider
 
             return E2BCodeExecToolProvider()
         from stirrup.tools.code_backends.docker import DockerCodeExecToolProvider
 
-        return DockerCodeExecToolProvider.from_image(_DEFAULT_CODE_IMAGE)
+        if self._preserve_workspace:
+            provider_cls = _preserving_provider_class(DockerCodeExecToolProvider)
+            return provider_cls(
+                _DEFAULT_CODE_IMAGE,
+                is_dockerfile=False,
+                temp_base_dir=self._workspace_dir,
+                preserve_dir=self._workspace_dir,
+            )
+        return DockerCodeExecToolProvider.from_image(
+            _DEFAULT_CODE_IMAGE,
+            temp_base_dir=self._workspace_dir,
+        )
 
     def _build_tools(self) -> list:
         tools: list = []
@@ -159,10 +234,12 @@ async def run(self, question: str) -> AgentResult:
             )
 
             _log.info(
-                "StirrupAgentRunner: starting (model=%s, code=%s, backend=%s)",
+                "StirrupAgentRunner: starting (model=%s, code=%s, backend=%s, workspace=%s, preserve=%s)",
                 self._model_id,
                 self._code_enabled,
                 self._code_backend,
+                self._workspace_dir,
+                self._preserve_workspace,
             )
 
             async with agent.session() as session:

src/agent/stirrup_agent/tests/test_runner.py

@@ -8,7 +8,11 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+from pathlib import Path
 
+import pytest
+
+from agent.stirrup_agent.runner import StirrupAgentRunner, _copy_workspace_contents
 from agent.stirrup_agent.trajectory import (
     build_trajectory,
     classify_tool,
@@ -69,6 +73,25 @@ def test_classify_tool():
     assert classify_tool("calculator", _DOMAIN) == "other"
 
 
+def test_copy_workspace_contents_preserves_nested_outputs(tmp_path: Path):
+    destination = tmp_path / "stirrup_agent_1001"
+    source = destination / "docker_exec_env_abc"
+    nested = source / "nested"
+    nested.mkdir(parents=True)
+    (source / "answer.txt").write_text("done", encoding="utf-8")
+    (nested / "plot.csv").write_text("x,y\n1,2\n", encoding="utf-8")
+
+    _copy_workspace_contents(source, destination)
+
+    assert (destination / "answer.txt").read_text(encoding="utf-8") == "done"
+    assert (destination / "nested" / "plot.csv").exists()
+
+
+def test_stirrup_runner_requires_workspace_when_preserving():
+    with pytest.raises(ValueError, match="workspace_dir"):
+        StirrupAgentRunner(preserve_workspace=True)
+
+
 def test_build_trajectory_maps_turns_calls_and_outputs():
     # history is list[list[ChatMessage]] (per-turn grouping)
     history = [

src/benchmark/scenario_suite_runner.py

@@ -49,6 +49,7 @@ class MethodConfig:
     model_id: str
     extra_args: tuple[str, ...] = ()
     workspace_root: Path | None = None
+    reset_workspace: bool = True
 
 
 def model_dir_name(model_id: str) -> str:
@@ -204,7 +205,7 @@ def run_agent_for_scenario(
         workspace_dir = method.workspace_root / run_id
         extra_args.extend(["--workspace-dir", str(workspace_dir)])
         if not dry_run:
-            if workspace_dir.exists():
+            if method.reset_workspace and workspace_dir.exists():
                 shutil.rmtree(workspace_dir)
             workspace_dir.mkdir(parents=True, exist_ok=True)
 
@@ -276,6 +277,13 @@ def run_evaluation(
 
 def build_methods(args: argparse.Namespace) -> dict[str, MethodConfig]:
     """Build available method configs from CLI args."""
+    stirrup_extra_args: list[str] = []
+    stirrup_extra_args.extend(["--max-tokens", str(args.stirrup_max_tokens)])
+    if getattr(args, "preserve_workspaces", False) and getattr(
+        args, "stirrup_workspace_root", None
+    ) is not None:
+        stirrup_extra_args.append("--preserve-workspace")
+
     opencode_extra_args: list[str] = []
     if args.opencode_allow_files:
         opencode_extra_args.append("--allow-files")
@@ -323,6 +331,8 @@ def build_methods(args: argparse.Namespace) -> dict[str, MethodConfig]:
             agent_name="stirrup_agent",
             command="stirrup-agent",
             model_id=args.model_id,
+            extra_args=tuple(stirrup_extra_args),
+            workspace_root=getattr(args, "stirrup_workspace_root", None),
         ),
         "opencode_agent": MethodConfig(
             agent_name="opencode_agent",
@@ -369,7 +379,12 @@ def _build_parser() -> argparse.ArgumentParser:
         prog="scenario_suite_runner",
         description="Run benchmark scenarios sequentially.",
     )
-
+    parser.add_argument(
+    "--stirrup-max-tokens",
+    type=int,
+    default=4096,
+    help="Max output tokens per Stirrup model call.",
+    )
     parser.add_argument(
         "--scenario-ids",
         type=Path,
@@ -412,6 +427,15 @@ def _build_parser() -> argparse.ArgumentParser:
         default=_DEFAULT_MODEL_ID,
         help="Model id used by direct_llm, stirrup_agent, and opencode_agent.",
     )
+    parser.add_argument(
+        "--stirrup-workspace-root",
+        type=Path,
+        default=None,
+        help=(
+            "Root directory for per-run Stirrup code-execution workspaces. "
+            "Workspaces are nested by agent/model/run_id."
+        ),
+    )
     parser.add_argument(
         "--gemini-model-id",
         default=_DEFAULT_GEMINI_MODEL_ID,
@@ -541,6 +565,14 @@ def _build_parser() -> argparse.ArgumentParser:
         action="store_true",
         help="Skip a scenario if its expected trajectory file already exists.",
     )
+    parser.add_argument(
+        "--preserve-workspaces",
+        action="store_true",
+        help=(
+            "Keep existing per-run workspaces instead of deleting them before a run. "
+            "For Stirrup, also preserve final code-execution files in --stirrup-workspace-root."
+        ),
+    )
     parser.add_argument(
         "--no-evaluate",
         action="store_true",
@@ -581,6 +613,13 @@ def main() -> None:
             )
         except ValueError as exc:
             parser.error(str(exc))
+    if args.stirrup_workspace_root is not None:
+        try:
+            validate_workspace_root_outside_repo(
+                args.stirrup_workspace_root, "--stirrup-workspace-root"
+            )
+        except ValueError as exc:
+            parser.error(str(exc))
     gemini_workspace_required = (
         args.gemini_allow_files or args.gemini_allow_bash or args.gemini_allow_edit
     )
@@ -621,6 +660,7 @@ def main() -> None:
             model_id=method.model_id,
             extra_args=method.extra_args,
             workspace_root=method_workspace_root,
+            reset_workspace=method.reset_workspace and not args.preserve_workspaces,
         )
 
         if not args.dry_run: