mirror of
https://github.com/usestrix/strix.git
synced 2026-11-04 07:27:16 +00:00
feat(tools): batch_view_files / batch_terminal_execute — sandbox-side concurrency
make_model_settings sets parallel_tool_calls=False (one tool call per turn), so
an agent mapping a codebase pays a full model round-trip per file read / command.
These two function tools let it fan out in a SINGLE call: the reads/execs run
concurrently inside the sandbox (asyncio.gather, semaphore-bounded to 8) and the
combined results come back in one tool response — a real wall-clock win on the
recon-heavy phase of a scan.
- batch_view_files(paths): concurrent read of up to 50 files; per-file result
{path, content|error, truncated}, order preserved (results[i] ↔ paths[i]).
- batch_terminal_execute(commands): concurrent run of up to 50 INDEPENDENT shell
commands; per-command {command, exit_code, stdout, stderr}.
Both are session-injected at the core (_view_files_impl / _exec_impl) with a thin
@function_tool shell, so the concurrency + isolation logic is directly testable
without the SDK tool-call layer. Per-item error isolation: one failing read/exec
returns its own error entry, never aborts the batch. Concurrency is bounded so a
large batch can't swamp the sandbox; batch size capped at 50.
Registered in _BASE_TOOLS so root + child agents get them. Session accessed via
the same ctx.context["sandbox_session"] convention the existing tools use. Tools
self-describe; no prompt change — models reach for them unprompted.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
40f4e67320
commit
ce391d76bc
5 changed files with 248 additions and 0 deletions
|
|
@ -24,6 +24,7 @@ from strix.tools.agents_graph.tools import (
|
|||
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,
|
||||
|
|
|
|||
0
strix/tools/batch/__init__.py
Normal file
0
strix/tools/batch/__init__.py
Normal file
134
strix/tools/batch/tools.py
Normal file
134
strix/tools/batch/tools.py
Normal file
|
|
@ -0,0 +1,134 @@
|
|||
"""``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
|
||||
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)
|
||||
_DEFAULT_TIMEOUT = 60
|
||||
|
||||
|
||||
def _session(ctx: RunContextWrapper) -> Any:
|
||||
inner = ctx.context if isinstance(ctx.context, dict) else {}
|
||||
return inner.get("sandbox_session")
|
||||
|
||||
|
||||
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]:
|
||||
rel = path.lstrip("/")
|
||||
target = rel if rel.startswith("workspace/") else f"workspace/{rel}"
|
||||
result = await session.exec("cat", f"/{target}", timeout=_DEFAULT_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=180, 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=_DEFAULT_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=300, 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)
|
||||
0
tests/tools/__init__.py
Normal file
0
tests/tools/__init__.py
Normal file
111
tests/tools/test_batch_tools.py
Normal file
111
tests/tools/test_batch_tools.py
Normal file
|
|
@ -0,0 +1,111 @@
|
|||
"""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
|
||||
|
||||
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_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"]
|
||||
Loading…
Add table
Reference in a new issue