diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 977d01f19..6d2948196 100644 --- a/strix/agents/factory.py +++ b/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, diff --git a/strix/tools/batch/__init__.py b/strix/tools/batch/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/strix/tools/batch/tools.py b/strix/tools/batch/tools.py new file mode 100644 index 000000000..e5e66a36b --- /dev/null +++ b/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) diff --git a/tests/tools/__init__.py b/tests/tools/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/tools/test_batch_tools.py b/tests/tools/test_batch_tools.py new file mode 100644 index 000000000..509745a7e --- /dev/null +++ b/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"]