feat(tools): batch_view_files / batch_terminal_execute — sandbox-side concurrency

strix/agents/factory.py

@@ -24,6 +24,7 @@
     view_agent_graph,
     wait_for_message,
 )
+from strix.tools.batch.tools import batch_terminal_execute, batch_view_files
 from strix.tools.finish.tool import finish_scan
 from strix.tools.load_skill.tool import load_skill
 from strix.tools.notes.tools import (
@@ -329,6 +330,8 @@ def _finish_tool_use_behavior(
 _BASE_TOOLS: tuple[Tool, ...] = (
     think,
     load_skill,
+    batch_view_files,
+    batch_terminal_execute,
     create_todo,
     list_todos,
     update_todo,

strix/tools/batch/tools.py

@@ -0,0 +1,166 @@
+"""``batch_view_files`` / ``batch_terminal_execute`` — sandbox-side concurrency.
+
+v1 sets ``parallel_tool_calls=False`` (the SDK issues one tool call per turn), so
+an agent that needs to read 20 files or run 10 independent commands pays a full
+model round-trip each. These tools let it fan out in a SINGLE call: the reads /
+execs run concurrently inside the sandbox (bounded), and the combined results
+come back in one tool response — the wall-clock win on the recon-heavy phase of
+a scan.
+
+Concurrency is bounded (a semaphore) so a large batch can't exhaust the sandbox.
+Each item is isolated: one failing read/exec returns its own error entry, never
+aborts the batch. Order is preserved (results[i] corresponds to inputs[i]).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+from pathlib import PurePosixPath
+from typing import Any
+
+from agents import RunContextWrapper, function_tool
+
+
+logger = logging.getLogger(__name__)
+
+_MAX_CONCURRENCY = 8          # bounded fan-out — don't swamp the sandbox
+_MAX_ITEMS = 50               # cap batch size (a 200-file batch is a smell)
+
+# Per-item timeouts are bounded so a full, valid batch cannot exceed the
+# enclosing @function_tool deadline (below) and get cancelled before returning
+# its per-item results. Worst case is ceil(_MAX_ITEMS / _MAX_CONCURRENCY) waves
+# of per-item timeouts: ceil(50/8) = 7. So keep 7 * per_item < tool deadline:
+#   view: 7 * 20s = 140s < 180s ;  exec: 7 * 40s = 280s < 300s.
+_VIEW_ITEM_TIMEOUT = 20
+_EXEC_ITEM_TIMEOUT = 40
+_VIEW_TOOL_TIMEOUT = 180
+_EXEC_TOOL_TIMEOUT = 300
+
+
+def _session(ctx: RunContextWrapper) -> Any:
+    inner = ctx.context if isinstance(ctx.context, dict) else {}
+    return inner.get("sandbox_session")
+
+
+def _safe_workspace_path(path: str) -> str | None:
+    """Resolve a caller path to an absolute /workspace path, or None if it would
+    escape the workspace root.
+
+    The path is passed to ``cat`` verbatim, so a ``..`` hop (e.g.
+    ``workspace/../../etc/passwd``) would read container files outside the
+    target. Normalise and require the result to stay under ``/workspace``.
+    """
+    rel = path.lstrip("/")
+    rel = rel[len("workspace/") :] if rel.startswith("workspace/") else rel
+    # PurePosixPath: sandbox paths are always posix regardless of host OS.
+    normalized = PurePosixPath("/workspace") / rel
+    resolved = PurePosixPath(os.path.normpath(str(normalized)))
+    if resolved != PurePosixPath("/workspace") and PurePosixPath(
+        "/workspace"
+    ) not in resolved.parents:
+        return None
+    return str(resolved)
+
+
+async def _gather_bounded(coros: list[Any]) -> list[Any]:
+    sem = asyncio.Semaphore(_MAX_CONCURRENCY)
+
+    async def _run(coro: Any) -> Any:
+        async with sem:
+            return await coro
+
+    return await asyncio.gather(*(_run(c) for c in coros), return_exceptions=True)
+
+
+async def _view_files_impl(session: Any, paths: list[str]) -> dict[str, Any]:
+    """Core of batch_view_files — session-injected so it's directly testable."""
+    if session is None:
+        return {"success": False, "error": "no sandbox session in context"}
+    if not paths:
+        return {"success": True, "results": []}
+    if len(paths) > _MAX_ITEMS:
+        return {"success": False,
+                "error": f"batch too large ({len(paths)} > {_MAX_ITEMS}); split it"}
+
+    async def _read(path: str) -> dict[str, Any]:
+        safe = _safe_workspace_path(path)
+        if safe is None:
+            return {"path": path, "error": "path escapes /workspace (rejected)"}
+        result = await session.exec("cat", safe, timeout=_VIEW_ITEM_TIMEOUT)
+        if getattr(result, "exit_code", 1) != 0:
+            return {"path": path, "error": (getattr(result, "stderr", "") or
+                                            "read failed").strip()[:200]}
+        content = result.stdout or ""
+        truncated = len(content) > 100_000
+        return {"path": path, "content": content[:100_000], "truncated": truncated}
+
+    settled = await _gather_bounded([_read(p) for p in paths])
+    results = [
+        r if isinstance(r, dict)
+        else {"path": paths[i], "error": f"{type(r).__name__}: {r}"}
+        for i, r in enumerate(settled)
+    ]
+    return {"success": True, "results": results}
+
+
+@function_tool(timeout=_VIEW_TOOL_TIMEOUT, strict_mode=False)
+async def batch_view_files(
+    ctx: RunContextWrapper,
+    paths: list[str],
+) -> dict[str, Any]:
+    """Read multiple files from the target in one call (concurrent, bounded).
+
+    Prefer this over N single reads when mapping a codebase — one tool call, one
+    model round-trip. Each result carries the path, content (or an error), and a
+    truncation flag; order matches ``paths``.
+
+    Args:
+        paths: repo-relative file paths to read (max 50).
+    """
+    return await _view_files_impl(_session(ctx), paths)
+
+
+async def _exec_impl(session: Any, commands: list[str]) -> dict[str, Any]:
+    """Core of batch_terminal_execute — session-injected so it's directly testable."""
+    if session is None:
+        return {"success": False, "error": "no sandbox session in context"}
+    if not commands:
+        return {"success": True, "results": []}
+    if len(commands) > _MAX_ITEMS:
+        return {"success": False,
+                "error": f"batch too large ({len(commands)} > {_MAX_ITEMS}); split it"}
+
+    async def _run(cmd: str) -> dict[str, Any]:
+        result = await session.exec("sh", "-c", cmd, timeout=_EXEC_ITEM_TIMEOUT)
+        return {
+            "command": cmd,
+            "exit_code": getattr(result, "exit_code", None),
+            "stdout": (result.stdout or "")[:50_000],
+            "stderr": (getattr(result, "stderr", "") or "")[:10_000],
+        }
+
+    settled = await _gather_bounded([_run(c) for c in commands])
+    results = [
+        r if isinstance(r, dict)
+        else {"command": commands[i], "error": f"{type(r).__name__}: {r}"}
+        for i, r in enumerate(settled)
+    ]
+    return {"success": True, "results": results}
+
+
+@function_tool(timeout=_EXEC_TOOL_TIMEOUT, strict_mode=False)
+async def batch_terminal_execute(
+    ctx: RunContextWrapper,
+    commands: list[str],
+) -> dict[str, Any]:
+    """Run multiple INDEPENDENT shell commands in the sandbox in one call
+    (concurrent, bounded). Use only for commands with no ordering dependency
+    between them (parallel greps, per-service probes, etc.) — order of
+    completion is not guaranteed, but results[i] maps to commands[i].
+
+    Args:
+        commands: shell command strings to run (max 50).
+    """
+    return await _exec_impl(_session(ctx), commands)

tests/tools/test_batch_tools.py

@@ -0,0 +1,149 @@
+"""batch_view_files / batch_terminal_execute — sandbox-side concurrency tools.
+
+Exercised with a fake session (no sandbox) covering: order preservation,
+per-item error isolation, the batch-size cap, missing-session guard, and that
+the reads/execs actually fan out concurrently (bounded).
+"""
+from __future__ import annotations
+
+import asyncio
+import math
+
+from strix.tools.batch import tools as bt
+
+
+class _Result:
+    def __init__(self, stdout="", exit_code=0, stderr=""):
+        self.stdout = stdout
+        self.exit_code = exit_code
+        self.stderr = stderr
+
+
+class _FakeSession:
+    def __init__(self, handler):
+        self._handler = handler
+        self.max_in_flight = 0
+        self._in_flight = 0
+
+    async def exec(self, *args, timeout=60):  # noqa: ARG002 — matches session.exec signature
+        self._in_flight += 1
+        self.max_in_flight = max(self.max_in_flight, self._in_flight)
+        await asyncio.sleep(0.01)  # force overlap so concurrency is observable
+        try:
+            return self._handler(args)
+        finally:
+            self._in_flight -= 1
+
+
+def _view(session, paths):
+    return asyncio.run(bt._view_files_impl(session, paths))
+
+
+def _exec(session, commands):
+    return asyncio.run(bt._exec_impl(session, commands))
+
+
+# ---- batch_view_files -----------------------------------------------------
+
+def test_view_files_reads_all_in_order():
+    sess = _FakeSession(lambda args: _Result(stdout=f"content of {args[-1]}"))
+    res = _view(sess, ["a.py", "b.py", "c.py"])
+    assert res["success"] is True
+    assert [r["path"] for r in res["results"]] == ["a.py", "b.py", "c.py"]
+    assert "content of /workspace/a.py" in res["results"][0]["content"]
+
+
+def test_view_files_error_isolated():
+    def handler(args):
+        if "bad" in args[-1]:
+            return _Result(exit_code=1, stderr="No such file")
+        return _Result(stdout="ok")
+    res = _view(_FakeSession(handler), ["good.py", "bad.py"])
+    assert "content" in res["results"][0]           # good read succeeded
+    assert "error" in res["results"][1]              # bad read isolated
+    assert res["success"] is True                    # batch didn't abort
+
+
+def test_view_files_runs_concurrently():
+    sess = _FakeSession(lambda _a: _Result(stdout="x"))
+    _view(sess, [f"f{i}.py" for i in range(6)])
+    assert sess.max_in_flight > 1                    # actually overlapped
+
+
+def test_view_files_cap_enforced():
+    res = _view(_FakeSession(lambda _a: _Result()), [f"f{i}" for i in range(51)])
+    assert res["success"] is False and "too large" in res["error"]
+
+
+def test_view_files_no_session():
+    res = _view(None, ["a"])
+    assert res["success"] is False and "no sandbox session" in res["error"]
+
+
+def test_view_files_rejects_path_traversal():
+    # A '..' hop must not escape /workspace — the path is cat'd verbatim, so
+    # workspace/../../etc/passwd would read container files otherwise.
+    catted: list[str] = []
+
+    def handler(args):
+        catted.append(args[-1])
+        return _Result(stdout="x")
+
+    res = _view(_FakeSession(handler), ["../../etc/passwd", "workspace/../secret", "ok.py"])
+    # traversal entries are rejected with an error, never handed to cat
+    assert res["results"][0]["error"].startswith("path escapes")
+    assert res["results"][1]["error"].startswith("path escapes")
+    assert "content" in res["results"][2]  # the safe one still reads
+    assert all("/workspace/../" not in c and "/etc/passwd" not in c for c in catted)
+
+
+def test_view_files_normalizes_workspace_prefix():
+    # Both bare and workspace-prefixed forms resolve to the same /workspace path.
+    seen: list[str] = []
+
+    def handler(args):
+        seen.append(args[-1])
+        return _Result(stdout="x")
+
+    _view(_FakeSession(handler), ["src/a.py", "workspace/src/a.py", "/src/a.py"])
+    assert seen == ["/workspace/src/a.py"] * 3
+
+
+def test_batch_deadline_fits_tool_timeout():
+    # Worst-case wave count * per-item timeout must stay under the enclosing
+    # @function_tool deadline, or a valid batch is cancelled before returning.
+    waves = math.ceil(bt._MAX_ITEMS / bt._MAX_CONCURRENCY)
+    assert waves * bt._VIEW_ITEM_TIMEOUT < bt._VIEW_TOOL_TIMEOUT
+    assert waves * bt._EXEC_ITEM_TIMEOUT < bt._EXEC_TOOL_TIMEOUT
+
+
+def test_view_files_empty_is_ok():
+    res = _view(_FakeSession(lambda _a: _Result()), [])
+    assert res["success"] is True and res["results"] == []
+
+
+# ---- batch_terminal_execute -----------------------------------------------
+
+def test_exec_runs_all_in_order():
+    sess = _FakeSession(lambda args: _Result(stdout=f"ran: {args[-1]}", exit_code=0))
+    res = _exec(sess, ["echo a", "echo b"])
+    assert res["success"] is True
+    assert [r["command"] for r in res["results"]] == ["echo a", "echo b"]
+    assert res["results"][0]["exit_code"] == 0
+    assert "ran: echo a" in res["results"][0]["stdout"]
+
+
+def test_exec_error_isolated():
+    def handler(args):
+        cmd = args[-1]
+        if "fail" in cmd:
+            return _Result(exit_code=1, stderr="boom")
+        return _Result(stdout="fine", exit_code=0)
+    res = _exec(_FakeSession(handler), ["true", "fail-cmd"])
+    assert res["results"][0]["exit_code"] == 0
+    assert res["results"][1]["exit_code"] == 1 and res["results"][1]["stderr"] == "boom"
+
+
+def test_exec_cap_enforced():
+    res = _exec(_FakeSession(lambda _a: _Result()), [f"echo {i}" for i in range(51)])
+    assert res["success"] is False and "too large" in res["error"]