mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-08-21 12:12:22 +02:00
197 lines
8.2 KiB
Python
197 lines
8.2 KiB
Python
#!/usr/bin/env python3
|
|
"""RunToolScript: the model writes a short Python script that chains builtin tools, and only
|
|
the script's printed output enters context (PTC, lifted from hermes-agent's code_execution_tool,
|
|
MIT). A 10-page research sweep costs the tokens of its summary instead of ten raw page dumps.
|
|
|
|
Safety shape: the script subprocess gets a minimal env with NO auth token; every tool call is
|
|
brokered here against an explicit ALLOWLIST (never the deny-list's guess), so the gated surfaces
|
|
(MCP activation, delegation, HITL, canvas, schedules, SettingsWrite) are unreachable from scripts
|
|
by construction. Agents already hold always_allow Bash, so this adds reach for no new privilege."""
|
|
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
|
|
SCRIPT_TIMEOUT_S = 300.0
|
|
MAX_TOOL_CALLS = 50
|
|
MAX_STDOUT_BYTES = 50_000
|
|
|
|
# Read/write-safe, ungated tools only. Everything else is invisible to scripts on purpose.
|
|
SCRIPT_ALLOWED_TOOLS = ("WebSearch", "WebFetch", "MemoryRead", "MemoryWrite", "SettingsRead", "Skill")
|
|
|
|
TOOLS = [
|
|
{
|
|
"name": "RunToolScript",
|
|
"description": (
|
|
"Run a short Python script that chains multiple tool calls in ONE turn; only what "
|
|
"the script print()s comes back to you, so bulky intermediate results never enter "
|
|
"your context. Use this whenever a task needs 3+ tool calls whose raw outputs you "
|
|
"would only aggregate anyway (fetch N pages and extract one fact each, search then "
|
|
"fetch the top hits, sweep memory). The script gets one function: "
|
|
"call_tool(name, args_dict) -> str, valid names: "
|
|
+ ", ".join(SCRIPT_ALLOWED_TOOLS) + ". A failed tool raises PtcToolError (catchable). "
|
|
"Print ONLY your distilled findings. Budget: "
|
|
f"{MAX_TOOL_CALLS} tool calls, {SCRIPT_TIMEOUT_S:.0f}s, printed output capped at 50KB."
|
|
),
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"script": {
|
|
"type": "string",
|
|
"description": "Python source. Example: results = call_tool('WebSearch', {'query': 'x'}) then loop call_tool('WebFetch', {'url': u}) and print the aggregate.",
|
|
},
|
|
},
|
|
"required": ["script"],
|
|
},
|
|
},
|
|
]
|
|
|
|
|
|
def p_core():
|
|
"""The running combined sidecar (spawned as __main__), which owns tool routing. Tests
|
|
inject a stand-in via set_core_for_tests."""
|
|
if p_core_override is not None:
|
|
return p_core_override
|
|
main_mod = sys.modules.get("__main__")
|
|
if main_mod is not None and hasattr(main_mod, "P_ROUTE") and hasattr(main_mod, "p_call"):
|
|
return main_mod
|
|
return None
|
|
|
|
|
|
p_core_override = None
|
|
|
|
|
|
def set_core_for_tests(core) -> None:
|
|
global p_core_override
|
|
p_core_override = core
|
|
|
|
|
|
def p_result_text(result: dict) -> str:
|
|
parts = []
|
|
for c in result.get("content", []) or []:
|
|
if isinstance(c, dict) and c.get("type") == "text":
|
|
parts.append(str(c.get("text", "")))
|
|
return "\n".join(parts)
|
|
|
|
|
|
def p_dispatch(name: str, args: dict) -> dict:
|
|
"""One brokered inner call: allowlist first, then the sidecar's own routing."""
|
|
if name not in SCRIPT_ALLOWED_TOOLS:
|
|
return {"text": f"tool '{name}' is not callable from scripts; allowed: {', '.join(SCRIPT_ALLOWED_TOOLS)}", "is_error": True}
|
|
core = p_core()
|
|
if core is None:
|
|
return {"text": "tool routing unavailable", "is_error": True}
|
|
mod = core.P_ROUTE.get(name)
|
|
if mod is None:
|
|
return {"text": f"tool '{name}' is not loaded in this session", "is_error": True}
|
|
try:
|
|
result = core.p_call(mod, name, args)
|
|
except Exception as e:
|
|
return {"text": f"tool '{name}' raised: {e}", "is_error": True}
|
|
return {"text": p_result_text(result), "is_error": bool(result.get("isError"))}
|
|
|
|
|
|
def p_elide(text: str, cap: int = MAX_STDOUT_BYTES) -> str:
|
|
raw = text.encode("utf-8", errors="replace")
|
|
if len(raw) <= cap:
|
|
return text
|
|
head = int(cap * 0.4)
|
|
tail = cap - head
|
|
return (
|
|
raw[:head].decode("utf-8", errors="replace")
|
|
+ f"\n\n[... output elided: {len(raw)} bytes total, cap {cap} ...]\n\n"
|
|
+ raw[-tail:].decode("utf-8", errors="replace")
|
|
)
|
|
|
|
|
|
def p_runner_env() -> dict:
|
|
# Minimal on purpose: the child needs no secrets because the parent brokers every call.
|
|
keep = ("PATH", "HOME", "TMPDIR", "LANG", "LC_ALL", "SYSTEMROOT", "TEMP", "TMP")
|
|
env = {k: os.environ[k] for k in keep if k in os.environ}
|
|
env["PYTHONUNBUFFERED"] = "1"
|
|
env["PYTHONDONTWRITEBYTECODE"] = "1"
|
|
return env
|
|
|
|
|
|
def p_mcp_text(text: str, is_error: bool = False) -> dict:
|
|
out = {"content": [{"type": "text", "text": text}]}
|
|
if is_error:
|
|
out["isError"] = True
|
|
return out
|
|
|
|
|
|
def handle_tool_call(tool_name: str, arguments: dict) -> dict:
|
|
if tool_name != "RunToolScript":
|
|
return p_mcp_text(f"Unknown tool: {tool_name}", is_error=True)
|
|
script = str(arguments.get("script", "")).strip()
|
|
if not script:
|
|
return p_mcp_text("script is required", is_error=True)
|
|
runner = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ptc_script_runner.py")
|
|
try:
|
|
proc = subprocess.Popen(
|
|
[sys.executable, "-u", runner],
|
|
stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL,
|
|
env=p_runner_env(), text=True,
|
|
)
|
|
except Exception as e:
|
|
return p_mcp_text(f"could not start script runner: {e}", is_error=True)
|
|
deadline = time.monotonic() + SCRIPT_TIMEOUT_S
|
|
# The read loop blocks in readline, so the deadline needs teeth of its own: the timer kills the child, which turns the block into a clean EOF.
|
|
p_reaper = threading.Timer(SCRIPT_TIMEOUT_S, lambda: proc.poll() is None and proc.kill())
|
|
p_reaper.daemon = True
|
|
p_reaper.start()
|
|
calls_used = 0
|
|
try:
|
|
proc.stdin.write(json.dumps({"script": script}) + "\n")
|
|
proc.stdin.flush()
|
|
while True:
|
|
if time.monotonic() > deadline:
|
|
proc.kill()
|
|
return p_mcp_text(
|
|
f"script exceeded the {SCRIPT_TIMEOUT_S:.0f}s budget after {calls_used} tool calls; nothing was returned. Break the work into smaller scripts.",
|
|
is_error=True,
|
|
)
|
|
line = proc.stdout.readline()
|
|
if not line:
|
|
if time.monotonic() > deadline:
|
|
return p_mcp_text(
|
|
f"script exceeded the {SCRIPT_TIMEOUT_S:.0f}s budget after {calls_used} tool calls; nothing was returned. Break the work into smaller scripts.",
|
|
is_error=True,
|
|
)
|
|
return p_mcp_text(f"script runner exited unexpectedly after {calls_used} tool calls", is_error=True)
|
|
try:
|
|
msg = json.loads(line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
if msg.get("done"):
|
|
out = p_elide(str(msg.get("stdout", "")))
|
|
err = msg.get("error")
|
|
footer = f"\n\n[script ran {int(msg.get('calls', 0))} tool call(s)]"
|
|
if err:
|
|
return p_mcp_text(f"script raised {err}\n\npartial output:\n{out}{footer}", is_error=True)
|
|
if not out.strip():
|
|
return p_mcp_text(f"script printed nothing; print your findings next time.{footer}", is_error=True)
|
|
return p_mcp_text(out + footer)
|
|
call = msg.get("call")
|
|
if not isinstance(call, dict):
|
|
continue
|
|
calls_used += 1
|
|
if calls_used > MAX_TOOL_CALLS:
|
|
reply = {"seq": msg.get("seq"), "text": f"tool call budget ({MAX_TOOL_CALLS}) exhausted; print what you have", "is_error": True}
|
|
else:
|
|
reply = {"seq": msg.get("seq"), **p_dispatch(str(call.get("name", "")), dict(call.get("args") or {}))}
|
|
proc.stdin.write(json.dumps(reply) + "\n")
|
|
proc.stdin.flush()
|
|
except BrokenPipeError:
|
|
return p_mcp_text(f"script runner pipe broke after {calls_used} tool calls", is_error=True)
|
|
finally:
|
|
p_reaper.cancel()
|
|
try:
|
|
if proc.poll() is None:
|
|
proc.kill()
|
|
except Exception:
|
|
pass
|