plato-sdk-v2 2.3.6__py3-none-any.whl → 2.3.8__py3-none-any.whl

plato/agents/runner.py

@@ -1,4 +1,11 @@
-"""Agent runner - run agents in Docker containers."""
+"""Agent runner - run agents in Docker containers.
+
+Agents emit their own OTel spans for trajectory events. This runner:
+1. Runs agents in Docker containers
+2. Streams stdout/stderr for logging
+3. Passes OTel environment variables for trace context propagation
+4. Uploads artifacts to S3 when complete
+"""
 
 from __future__ import annotations
 
@@ -8,12 +15,10 @@ import logging
 import os
 import platform
 import tempfile
-from pathlib import Path
 
 from opentelemetry import trace
 
 from plato.agents.artifacts import upload_artifacts
-from plato.agents.otel import get_tracer
 
 logger = logging.getLogger(__name__)
 
@@ -37,310 +42,142 @@ async def run_agent(
         workspace: Host directory to mount as /workspace
         logs_dir: Host directory for logs (temp dir if None)
         pull: Whether to pull the image first
+
+    Note: Agents handle their own OTel tracing. This runner only passes
+    the trace context (TRACEPARENT) so agent spans link to the parent step.
     """
     logs_dir = logs_dir or tempfile.mkdtemp(prefix="agent_logs_")
-    agent_name = image.split("/")[-1].split(":")[0]
 
     # Get session info from environment variables
     session_id = os.environ.get("SESSION_ID")
     otel_url = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT")
     upload_url = os.environ.get("UPLOAD_URL")
 
-    tracer = get_tracer("plato.agent")
-
-    with tracer.start_as_current_span(agent_name) as agent_span:
-        agent_span.set_attribute("span.type", "agent")
-        agent_span.set_attribute("source", "agent")
-        agent_span.set_attribute("image", image)
-        agent_span.set_attribute("content", f"Starting agent: {agent_name}")
-
-        # Pull image if requested
-        if pull:
-            with tracer.start_as_current_span("docker_pull") as pull_span:
-                pull_span.set_attribute("span.type", "docker_pull")
-                pull_span.set_attribute("image", image)
-                pull_proc = await asyncio.create_subprocess_exec(
-                    "docker",
-                    "pull",
-                    image,
-                    stdout=asyncio.subprocess.PIPE,
-                    stderr=asyncio.subprocess.STDOUT,
+    # Pull image if requested
+    if pull:
+        pull_proc = await asyncio.create_subprocess_exec(
+            "docker",
+            "pull",
+            image,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
+        )
+        await pull_proc.wait()
+
+    # Setup
+    os.makedirs(os.path.join(logs_dir, "agent"), exist_ok=True)
+    config_file = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False)
+    json.dump(config, config_file)
+    config_file.close()
+
+    try:
+        # Build docker command
+        docker_cmd = ["docker", "run", "--rm", "--privileged"]
+
+        # Determine if we need host networking
+        use_host_network = False
+        is_macos = platform.system() == "Darwin"
+
+        if not is_macos:
+            try:
+                proc = await asyncio.create_subprocess_exec(
+                    "iptables",
+                    "-L",
+                    "-n",
+                    stdout=asyncio.subprocess.DEVNULL,
+                    stderr=asyncio.subprocess.DEVNULL,
                 )
-                await pull_proc.wait()
-
-        # Setup
-        os.makedirs(os.path.join(logs_dir, "agent"), exist_ok=True)
-        config_file = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False)
-        json.dump(config, config_file)
-        config_file.close()
-
-        try:
-            # Build docker command
-            docker_cmd = ["docker", "run", "--rm"]
-
-            # Determine if we need host networking
-            use_host_network = False
-            is_macos = platform.system() == "Darwin"
-
-            if not is_macos:
-                try:
-                    proc = await asyncio.create_subprocess_exec(
-                        "iptables",
-                        "-L",
-                        "-n",
-                        stdout=asyncio.subprocess.DEVNULL,
-                        stderr=asyncio.subprocess.DEVNULL,
-                    )
-                    await proc.wait()
-                    has_iptables = proc.returncode == 0
-                except (FileNotFoundError, PermissionError):
-                    has_iptables = False
-
-                use_host_network = not has_iptables
-
-            if use_host_network:
-                docker_cmd.extend(["--network=host", "--add-host=localhost:127.0.0.1"])
-
+                await proc.wait()
+                has_iptables = proc.returncode == 0
+            except (FileNotFoundError, PermissionError):
+                has_iptables = False
+
+            use_host_network = not has_iptables
+
+        if use_host_network:
+            docker_cmd.extend(["--network=host", "--add-host=localhost:127.0.0.1"])
+
+        docker_cmd.extend(
+            [
+                "-v",
+                f"{workspace}:/workspace",
+                "-v",
+                f"{logs_dir}:/logs",
+                "-v",
+                f"{config_file.name}:/config.json:ro",
+                "-v",
+                "/var/run/docker.sock:/var/run/docker.sock",
+                "-w",
+                "/workspace",
+            ]
+        )
+
+        # Pass session info to agent
+        if otel_url:
+            traces_endpoint = f"{otel_url.rstrip('/')}/v1/traces"
+            docker_cmd.extend(["-e", f"OTEL_EXPORTER_OTLP_ENDPOINT={otel_url}"])
+            docker_cmd.extend(["-e", f"OTEL_EXPORTER_OTLP_TRACES_ENDPOINT={traces_endpoint}"])
+            docker_cmd.extend(["-e", "OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf"])
+        if session_id:
+            docker_cmd.extend(["-e", f"SESSION_ID={session_id}"])
+        if upload_url:
+            docker_cmd.extend(["-e", f"UPLOAD_URL={upload_url}"])
+
+        # Pass trace context to agent for parent linking
+        # Agent spans will be children of the current step span
+        current_span = trace.get_current_span()
+        span_context = current_span.get_span_context()
+        if span_context.is_valid:
+            trace_id = format(span_context.trace_id, "032x")
+            span_id = format(span_context.span_id, "016x")
+            # W3C Trace Context format for TRACEPARENT
+            traceparent = f"00-{trace_id}-{span_id}-01"
             docker_cmd.extend(
                 [
-                    "-v",
-                    f"{workspace}:/workspace",
-                    "-v",
-                    f"{logs_dir}:/logs",
-                    "-v",
-                    f"{config_file.name}:/config.json:ro",
-                    "-v",
-                    "/var/run/docker.sock:/var/run/docker.sock",
-                    "-w",
-                    "/workspace",
+                    "-e",
+                    f"TRACEPARENT={traceparent}",
+                    "-e",
+                    f"OTEL_TRACE_ID={trace_id}",
+                    "-e",
+                    f"OTEL_PARENT_SPAN_ID={span_id}",
                 ]
             )
 
-            # Pass session info to agent
-            if otel_url:
-                docker_cmd.extend(["-e", f"OTEL_EXPORTER_OTLP_ENDPOINT={otel_url}"])
-                # Use JSON protocol (not protobuf) for OTLP exports
-                docker_cmd.extend(["-e", "OTEL_EXPORTER_OTLP_PROTOCOL=http/json"])
-            if session_id:
-                docker_cmd.extend(["-e", f"SESSION_ID={session_id}"])
-            if upload_url:
-                docker_cmd.extend(["-e", f"UPLOAD_URL={upload_url}"])
-
-            # Pass trace context to agent for parent linking
-            current_span = trace.get_current_span()
-            span_context = current_span.get_span_context()
-            if span_context.is_valid:
-                trace_id = format(span_context.trace_id, "032x")
-                span_id = format(span_context.span_id, "016x")
-                docker_cmd.extend(
-                    [
-                        "-e",
-                        f"OTEL_TRACE_ID={trace_id}",
-                        "-e",
-                        f"OTEL_PARENT_SPAN_ID={span_id}",
-                    ]
-                )
-
-            for key, value in secrets.items():
-                docker_cmd.extend(["-e", f"{key.upper()}={value}"])
-
-            docker_cmd.append(image)
-
-            # Pass instruction via CLI arg
-            docker_cmd.extend(["--instruction", instruction])
-
-            # Run container and stream output
-            with tracer.start_as_current_span("agent_execution") as exec_span:
-                exec_span.set_attribute("span.type", "agent_execution")
-                exec_span.set_attribute("content", f"Running {agent_name}")
-
-                process = await asyncio.create_subprocess_exec(
-                    *docker_cmd,
-                    stdout=asyncio.subprocess.PIPE,
-                    stderr=asyncio.subprocess.STDOUT,
-                )
-
-                # Stream output line by line
-                output_lines: list[str] = []
-                turn_count = 0
-                assert process.stdout is not None
-                while True:
-                    line = await process.stdout.readline()
-                    if not line:
-                        break
-                    decoded_line = line.decode().rstrip()
-                    output_lines.append(decoded_line)
-
-                    # Try to parse JSON output from agent for structured trajectory spans
-                    try:
-                        data = json.loads(decoded_line)
-                        event_type = data.get("type", "")
-
-                        if event_type == "assistant":
-                            # Agent response - create a turn span
-                            turn_count += 1
-                            msg = data.get("message", {})
-                            content_items = msg.get("content", [])
-
-                            # Extract text and tool calls with full details
-                            text_parts = []
-                            tool_calls = []
-                            for item in content_items:
-                                if item.get("type") == "text":
-                                    text_parts.append(item.get("text", "")[:2000])
-                                elif item.get("type") == "tool_use":
-                                    tool_input = item.get("input", {})
-                                    # Truncate large inputs
-                                    input_str = json.dumps(tool_input) if tool_input else ""
-                                    if len(input_str) > 2000:
-                                        input_str = input_str[:2000] + "..."
-                                    tool_calls.append(
-                                        {
-                                            "tool": item.get("name"),
-                                            "id": item.get("id"),
-                                            "input": input_str,
-                                        }
-                                    )
-
-                            with tracer.start_as_current_span(f"turn_{turn_count}") as turn_span:
-                                turn_span.set_attribute("span.type", "agent_turn")
-                                turn_span.set_attribute("source", "agent")
-                                turn_span.set_attribute("turn_number", turn_count)
-                                turn_span.set_attribute("model", msg.get("model", "unknown"))
-
-                                if text_parts:
-                                    turn_span.set_attribute("content", "\n".join(text_parts)[:4000])
-                                if tool_calls:
-                                    turn_span.set_attribute("tool_calls", json.dumps(tool_calls))
-                                    # If no text content, show tool calls summary
-                                    if not text_parts:
-                                        turn_span.set_attribute(
-                                            "content", f"Tool calls: {', '.join(t['tool'] for t in tool_calls)}"
-                                        )
-
-                                # Usage info
-                                usage = msg.get("usage", {})
-                                if usage:
-                                    turn_span.set_attribute("input_tokens", usage.get("input_tokens", 0))
-                                    turn_span.set_attribute("output_tokens", usage.get("output_tokens", 0))
-
-                        elif event_type == "user":
-                            # Tool result
-                            tool_results = data.get("message", {}).get("content", [])
-                            for result in tool_results:
-                                if result.get("type") == "tool_result":
-                                    tool_id = result.get("tool_use_id", "")
-                                    content = result.get("content", "")
-                                    # Handle content that might be a list of content blocks
-                                    if isinstance(content, list):
-                                        text_parts = []
-                                        for item in content:
-                                            if isinstance(item, dict) and item.get("type") == "text":
-                                                text_parts.append(item.get("text", ""))
-                                            elif isinstance(item, str):
-                                                text_parts.append(item)
-                                        content = "\n".join(text_parts)
-                                    if isinstance(content, str):
-                                        content = content[:2000]  # Truncate large results
-                                    with tracer.start_as_current_span("tool_result") as tr_span:
-                                        tr_span.set_attribute("span.type", "tool_result")
-                                        tr_span.set_attribute("source", "agent")
-                                        tr_span.set_attribute("tool_use_id", tool_id)
-                                        tr_span.set_attribute("content", f"Tool result for {tool_id}")
-                                        tr_span.set_attribute("result", content if content else "")
-
-                        elif event_type == "result":
-                            # Final result
-                            result_text = data.get("result", "")[:1000]
-                            is_error = data.get("is_error", False)
-                            duration_ms = data.get("duration_ms", 0)
-                            total_cost = data.get("total_cost_usd", 0)
-
-                            with tracer.start_as_current_span("agent_result") as res_span:
-                                res_span.set_attribute("span.type", "agent_result")
-                                res_span.set_attribute("source", "agent")
-                                res_span.set_attribute("content", result_text if result_text else "Agent completed")
-                                res_span.set_attribute("is_error", is_error)
-                                res_span.set_attribute("duration_ms", duration_ms)
-                                res_span.set_attribute("total_cost_usd", total_cost)
-                                res_span.set_attribute("num_turns", data.get("num_turns", turn_count))
-
-                        elif event_type == "system" and data.get("subtype") == "init":
-                            # Agent initialization
-                            with tracer.start_as_current_span("agent_init") as init_span:
-                                init_span.set_attribute("span.type", "agent_init")
-                                init_span.set_attribute("source", "agent")
-                                init_span.set_attribute("model", data.get("model", "unknown"))
-                                init_span.set_attribute("tools", json.dumps(data.get("tools", [])))
-                                init_span.set_attribute("content", f"Agent initialized: {data.get('model', 'unknown')}")
-
-                        else:
-                            # Other output - just log it without creating a span
-                            logger.debug(f"[agent] {decoded_line}")
-                            continue
-
-                    except json.JSONDecodeError:
-                        # Not JSON - just log it
-                        logger.info(f"[agent] {decoded_line}")
-
-                await process.wait()
-
-                exit_code = process.returncode or 0
-                if exit_code != 0:
-                    error_context = "\n".join(output_lines[-50:]) if output_lines else "No output captured"
-
-                    exec_span.set_attribute("error", True)
-                    exec_span.set_attribute("exit_code", exit_code)
-                    exec_span.add_event(
-                        "agent_error",
-                        {
-                            "exit_code": exit_code,
-                            "output": error_context[:4000],
-                        },
-                    )
+        for key, value in secrets.items():
+            docker_cmd.extend(["-e", f"{key.upper()}={value}"])
 
-                    agent_span.set_attribute("error", True)
-                    agent_span.set_attribute("exit_code", exit_code)
+        docker_cmd.append(image)
 
-                    raise RuntimeError(f"Agent failed with exit code {exit_code}")
+        # Pass instruction via CLI arg
+        docker_cmd.extend(["--instruction", instruction])
 
-                exec_span.set_attribute("success", True)
+        # Run container - agents emit their own OTel spans
+        process = await asyncio.create_subprocess_exec(
+            *docker_cmd,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
+        )
 
-        finally:
-            os.unlink(config_file.name)
+        # Capture output for error reporting
+        output_lines: list[str] = []
+        assert process.stdout is not None
+        while True:
+            line = await process.stdout.readline()
+            if not line:
+                break
+            decoded_line = line.decode().rstrip()
+            output_lines.append(decoded_line)
 
-            # Load trajectory and log as event
-            trajectory_path = Path(logs_dir) / "agent" / "trajectory.json"
-            if trajectory_path.exists():
-                try:
-                    with open(trajectory_path) as f:
-                        trajectory = json.load(f)
-                    if isinstance(trajectory, dict) and "schema_version" in trajectory:
-                        # Add agent image to trajectory
-                        agent_data = trajectory.get("agent", {})
-                        extra = agent_data.get("extra") or {}
-                        extra["image"] = image
-                        agent_data["extra"] = extra
-                        trajectory["agent"] = agent_data
+        await process.wait()
 
-                        # Log trajectory as span event
-                        with tracer.start_as_current_span("trajectory") as traj_span:
-                            traj_span.set_attribute("span.type", "trajectory")
-                            traj_span.set_attribute("log_type", "atif")
-                            traj_span.set_attribute("source", "agent")
-                            # Store trajectory in span (truncated for OTel limits)
-                            traj_json = json.dumps(trajectory)
-                            if len(traj_json) > 10000:
-                                traj_span.set_attribute("trajectory_truncated", True)
-                                traj_span.set_attribute("trajectory_size", len(traj_json))
-                            else:
-                                traj_span.set_attribute("trajectory", traj_json)
-                except Exception as e:
-                    logger.warning(f"Failed to load trajectory: {e}")
+        exit_code = process.returncode or 0
+        if exit_code != 0:
+            error_context = "\n".join(output_lines[-50:]) if output_lines else "No output captured"
+            raise RuntimeError(f"Agent failed with exit code {exit_code}\n\nAgent output:\n{error_context}")
 
-            # Upload artifacts if we have upload URL configured
-            if upload_url:
-                await upload_artifacts(upload_url, logs_dir)
+    finally:
+        os.unlink(config_file.name)
 
-        agent_span.set_attribute("success", True)
-        agent_span.set_attribute("content", f"Agent {agent_name} completed successfully")
+        # Upload artifacts if we have upload URL configured
+        if upload_url:
+            await upload_artifacts(upload_url, logs_dir)

plato/v1/cli/sandbox.py

@@ -131,6 +131,9 @@ def sandbox_start(
     timeout: int = typer.Option(1800, "--timeout", help="VM lifetime in seconds (default: 30 minutes)"),
     no_reset: bool = typer.Option(False, "--no-reset", help="Skip initial reset after ready"),
     json_output: bool = typer.Option(False, "--json", "-j", help="Output as JSON"),
+    working_dir: Path = typer.Option(
+        None, "--working-dir", "-w", help="Working directory for .sandbox.yaml and .plato/"
+    ),
 ):
     """
     Start a sandbox environment.
@@ -377,7 +380,7 @@ def sandbox_start(
                     console.print("[cyan]  Generating SSH key pair...[/cyan]")
 
                 base_url = os.getenv("PLATO_BASE_URL", "https://plato.so")
-                ssh_info = setup_ssh_for_sandbox(base_url, job_id, username=ssh_username)
+                ssh_info = setup_ssh_for_sandbox(base_url, job_id, username=ssh_username, working_dir=working_dir)
                 ssh_host = ssh_info["ssh_host"]
                 ssh_config_path = ssh_info["config_path"]
                 ssh_private_key_path = ssh_info["private_key_path"]
@@ -489,7 +492,7 @@ def sandbox_start(
         # Add heartbeat PID
         if heartbeat_pid:
             state["heartbeat_pid"] = heartbeat_pid
-        save_sandbox_state(state)
+        save_sandbox_state(state, working_dir)
 
         # Close the plato client (heartbeat process keeps session alive)
         plato.close()

plato/v1/cli/ssh.py

@@ -9,21 +9,21 @@ from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import ed25519
 
 
-def get_plato_dir() -> Path:
+def get_plato_dir(working_dir: Path | str | None = None) -> Path:
     """Get the directory for plato config/SSH files.
 
-    Uses /workspace/.plato if /workspace exists (container environment),
-    otherwise uses ~/.plato (local development).
+    Args:
+        working_dir: If provided, returns working_dir/.plato (for container/agent use).
+                     If None, returns ~/.plato (local development).
     """
-    workspace = Path("/workspace")
-    if workspace.exists() and workspace.is_dir():
-        return workspace / ".plato"
+    if working_dir is not None:
+        return Path(working_dir) / ".plato"
     return Path.home() / ".plato"
 
 
-def get_next_sandbox_number() -> int:
+def get_next_sandbox_number(working_dir: Path | str | None = None) -> int:
     """Find next available sandbox number by checking existing config files."""
-    plato_dir = get_plato_dir()
+    plato_dir = get_plato_dir(working_dir)
     if not plato_dir.exists():
         return 1
 
@@ -41,13 +41,13 @@ def get_next_sandbox_number() -> int:
     return max_num + 1
 
 
-def generate_ssh_key_pair(sandbox_num: int) -> tuple[str, str]:
+def generate_ssh_key_pair(sandbox_num: int, working_dir: Path | str | None = None) -> tuple[str, str]:
     """
     Generate a new ed25519 SSH key pair for a specific sandbox.
 
     Returns (public_key_str, private_key_path).
     """
-    plato_dir = get_plato_dir()
+    plato_dir = get_plato_dir(working_dir)
     plato_dir.mkdir(mode=0o700, exist_ok=True)
 
     private_key_path = plato_dir / f"ssh_{sandbox_num}_key"
@@ -136,6 +136,7 @@ def create_ssh_config(
     username: str,
     private_key_path: str,
     sandbox_num: int,
+    working_dir: Path | str | None = None,
 ) -> str:
     """
     Create a temporary SSH config file for a specific sandbox.
@@ -172,7 +173,7 @@ def create_ssh_config(
     TCPKeepAlive yes
 """
 
-    plato_dir = get_plato_dir()
+    plato_dir = get_plato_dir(working_dir)
     plato_dir.mkdir(mode=0o700, exist_ok=True)
 
     config_path = plato_dir / f"ssh_{sandbox_num}.conf"
@@ -182,7 +183,12 @@ def create_ssh_config(
     return str(config_path)
 
 
-def setup_ssh_for_sandbox(base_url: str, job_public_id: str, username: str = "plato") -> dict:
+def setup_ssh_for_sandbox(
+    base_url: str,
+    job_public_id: str,
+    username: str = "plato",
+    working_dir: Path | str | None = None,
+) -> dict:
     """
     Set up SSH access for a sandbox - generates keys and creates config.
 
@@ -190,14 +196,14 @@ def setup_ssh_for_sandbox(base_url: str, job_public_id: str, username: str = "pl
 
     Returns dict with: ssh_host, config_path, public_key, private_key_path
     """
-    sandbox_num = get_next_sandbox_number()
+    sandbox_num = get_next_sandbox_number(working_dir)
     ssh_host = f"sandbox-{sandbox_num}"
 
     # Choose random port between 2200 and 2299
     local_port = random.randint(2200, 2299)
 
     # Generate SSH key pair
-    public_key, private_key_path = generate_ssh_key_pair(sandbox_num)
+    public_key, private_key_path = generate_ssh_key_pair(sandbox_num, working_dir)
 
     # Create SSH config file
     config_path = create_ssh_config(
@@ -208,6 +214,7 @@ def setup_ssh_for_sandbox(base_url: str, job_public_id: str, username: str = "pl
         username=username,
         private_key_path=private_key_path,
         sandbox_num=sandbox_num,
+        working_dir=working_dir,
     )
 
     return {

plato/v1/cli/utils.py

@@ -15,32 +15,52 @@ console = Console()
 SANDBOX_FILE = ".sandbox.yaml"
 
 
-def get_sandbox_state() -> dict | None:
-    """Read sandbox state from .sandbox.yaml in current directory."""
-    sandbox_file = Path.cwd() / SANDBOX_FILE
+def get_sandbox_state(working_dir: Path | str | None = None) -> dict | None:
+    """Read sandbox state from .sandbox.yaml.
+
+    Args:
+        working_dir: Directory containing .sandbox.yaml. If None, uses cwd.
+    """
+    base_dir = Path(working_dir) if working_dir else Path.cwd()
+    sandbox_file = base_dir / SANDBOX_FILE
     if not sandbox_file.exists():
         return None
     with open(sandbox_file) as f:
         return yaml.safe_load(f)
 
 
-def save_sandbox_state(state: dict) -> None:
-    """Save sandbox state to .sandbox.yaml in current directory."""
-    sandbox_file = Path.cwd() / SANDBOX_FILE
+def save_sandbox_state(state: dict, working_dir: Path | str | None = None) -> None:
+    """Save sandbox state to .sandbox.yaml.
+
+    Args:
+        state: State dict to save.
+        working_dir: Directory to save .sandbox.yaml in. If None, uses cwd.
+    """
+    base_dir = Path(working_dir) if working_dir else Path.cwd()
+    sandbox_file = base_dir / SANDBOX_FILE
     with open(sandbox_file, "w") as f:
         yaml.dump(state, f, default_flow_style=False)
 
 
-def remove_sandbox_state() -> None:
-    """Remove .sandbox.yaml from current directory."""
-    sandbox_file = Path.cwd() / SANDBOX_FILE
+def remove_sandbox_state(working_dir: Path | str | None = None) -> None:
+    """Remove .sandbox.yaml.
+
+    Args:
+        working_dir: Directory containing .sandbox.yaml. If None, uses cwd.
+    """
+    base_dir = Path(working_dir) if working_dir else Path.cwd()
+    sandbox_file = base_dir / SANDBOX_FILE
     if sandbox_file.exists():
         sandbox_file.unlink()
 
 
-def require_sandbox_state() -> dict:
-    """Get sandbox state or exit with error."""
-    state = get_sandbox_state()
+def require_sandbox_state(working_dir: Path | str | None = None) -> dict:
+    """Get sandbox state or exit with error.
+
+    Args:
+        working_dir: Directory containing .sandbox.yaml. If None, uses cwd.
+    """
+    state = get_sandbox_state(working_dir)
     if not state:
         console.print("[red]No sandbox found in current directory[/red]")
         console.print("\n[yellow]Start a sandbox with:[/yellow]")