8000 feat(extensions): Python code tools, run in a process of their own (admins only) by adeelahmad · Pull Request #103 · adeelahmad/lense · GitHub
[go: up one dir, main page]

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion DEVLOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ Todo:

- [x] Registry, versions, sharing, manifests in and out, check without saving, try a tool
- [x] Prompt and web tools, skills, hooks in the chat tool loop; approvals for change tools
- [ ] Code tools: Python in a sandboxed subprocess (admins create; time limit, no secrets, no network unless allowed)
- [x] Code tools: Python in a process of its own (admins write them; time, CPU, memory limits, no secrets, own folder, no network unless allowed; stops when its owner is no longer an admin). Refine later: a container or gVisor sandbox for non-admins, pip packages per tool
- [x] Canvas: a `tool` workflow scope whose arg nodes are the tool's parameters (ask_model, call_tool nodes)
- [x] Canvas editor for tool graphs in the web app (Extensions → Draw a tool)
- [x] Authoring from chat and voice: the assistant drafts and saves extensions behind an approval card
Expand Down
10 changes: 9 additions & 1 deletion docs/assistant.md
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,8 @@ shared with namespaces (admins can share with everyone), and switched on or off

- **Tool**: something the assistant can call, with typed parameters. Its body is a prompt (`{{param}}` in the text), a
web request (public addresses only), or a graph drawn on the canvas (`ask_model`, `call_tool` and the primitives,
from `arg` nodes to `return` nodes). A tool with `effect: change` asks for approval before it runs.
from `arg` nodes to `return` nodes), or Python code (admins only, below). A tool with `effect: change` asks for
approval before it runs.
- **Skill**: instructions and a line saying when to use them. The assistant sees the line, and reads the instructions
only when the skill applies.
- **Hook**: runs when a question arrives, before or after a tool, or after an answer, to add context, block a tool or
Expand All @@ -113,6 +114,13 @@ Make one however suits you:
- **Code**: a manifest, Markdown with YAML frontmatter (the body is a skill's instructions or a prompt tool's prompt),
or YAML. `POST /api/v1/extensions/check` checks it without saving, `GET /api/v1/extensions/{id}` gives it back.
- **Canvas**: a tool whose `run` is `{type: graph, graph}`.
- **Python** (admins only): a tool whose `run` is `{type: python, code, seconds, network}`, where `code` defines
`run(**args)` and returns something JSON can hold. It runs in a Python process of its own, from an empty folder it
alone reads and writes, with none of the server's environment (no keys or database address), the standard library
only, limits on CPU time (`seconds`, up to 60), memory and files, and no web unless `network: true` (a network
namespace of its own where the machine allows one, and a guard on sockets either way). What it prints comes back
next to its result. It stops running if its owner stops being an admin. The guard raises the bar, it isn't a wall,
which is why only admins write these.
- **Chat or voice**: ask the assistant ("make me a tool that translates text into French"). It drafts the manifest,
fixes what the check says, and asks for your approval before saving, changing or switching anything.

Expand Down
125 changes: 125 additions & 0 deletions fastapi_backend/app/domain/code_tools.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,125 @@
"""Running a code tool's Python apart from the server.

The code runs in its own Python process, started from an empty folder, with no environment variables from the server
(so no keys, tokens or database address), limits on CPU time, memory, open files and the size of files it writes,
and a wall-clock time limit after which it's killed. It reads its arguments as JSON and gives back what `run(**args)`
returns, as JSON.

The web is shut off unless the tool says it needs it: in a network namespace of its own when this machine allows one
(`unshare -rn`, no network but loopback), and, either way, by refusing sockets to anywhere but loopback and the start
of other programs inside the process (Python's audit hooks). The audit hooks raise the bar; they aren't a wall, which is
why only admins write code tools.
"""

from __future__ import annotations

import functools
import json
import shutil
import subprocess
import sys
import tempfile

MAX_SECONDS = 60
MAX_OUTPUT = 200_000
MEMORY = 512 * 2**20

RUNNER = r"""
import json, os, resource, sys
cpu, mem, net = int(sys.argv[1]), int(sys.argv[2]), sys.argv[3] == "1"
HOME = os.path.realpath(os.getcwd())
LIBS = tuple({os.path.realpath(p) + os.sep for p in (sys.prefix, sys.base_prefix, sys.exec_prefix, sys.base_exec_prefix)})
resource.setrlimit(resource.RLIMIT_CPU, (cpu, cpu))
try:
resource.setrlimit(resource.RLIMIT_AS, (mem, mem))
except (ValueError, OSError):
pass
resource.setrlimit(resource.RLIMIT_FSIZE, (50 * 2**20, 50 * 2**20))
resource.setrlimit(resource.RLIMIT_NOFILE, (64, 64))
LOOP = ("127.0.0.1", "::1", "localhost")

def guard(event, args):
if event == "open" and not isinstance(args[0], int):
path = os.path.realpath(os.fsdecode(args[0]))
writing = any(c in str(args[1] or "") for c in "wax+") or (isinstance(args[2], int) and args[2] & (os.O_WRONLY | os.O_RDWR))
if not (path == HOME or path.startswith(HOME + os.sep) or (not writing and path.startswith(LIBS))):
raise PermissionError("a code tool reads and writes in its own folder only")
return
if event in ("ctypes.dlopen", "ctypes.cdata", "sys.setprofile", "sys.settrace", "sys.addaudithook"):
raise PermissionError("a code tool can't do that")
if event in ("subprocess.Popen", "os.system", "os.exec", "os.posix_spawn", "os.spawn", "os.fork", "os.forkpty", "pty.spawn"):
raise PermissionError("a code tool can't start other programs")
if not net and event in ("socket.connect", "socket.sendto", "socket.getaddrinfo"):
where = args[1] if event != "socket.getaddrinfo" else args[0]
host = where[0] if isinstance(where, tuple) else where
if host not in LOOP:
raise PermissionError("this code tool has no network: say it needs it (network: true)")

sys.addaudithook(guard)
src = sys.stdin.readline()
args = json.loads(sys.stdin.readline())
scope = {"__name__": "lens_tool"}
exec(compile(json.loads(src), "tool.py", "exec"), scope)
if not callable(scope.get("run")):
raise SystemExit("the code defines no run(**args) function")
out = scope["run"](**args)
sys.stdout.write("\n\x00LENS-RESULT\x00" + json.dumps(out, default=str))
"""


class CodeError(ValueError):
pass


@functools.lru_cache(maxsize=1)
def netns():
"""Whether this machine lets a process have a network namespace of its own (`unshare -rn`)."""
exe = shutil.which("unshare")
if not exe:
return None
try:
ok = subprocess.run([exe, "-rn", "true"], capture_output=True, timeout=5).returncode == 0
except (OSError, subprocess.SubprocessError):
ok = False
return exe if ok else None


def check(source):
"""The code compiles and defines run (SyntaxError and ValueError name the problem)."""
try:
tree = compile(source, "tool.py", "exec", flags=0x400, dont_inherit=True) # ast.PyCF_ONLY_AST
except SyntaxError as e:
raise CodeError(f"the code doesn't compile: {e.msg} (line {e.lineno})") from None
if not any(getattr(n, "name", None) == "run" for n in getattr(tree, "body", [])):
raise CodeError("the code defines a function run(**args) that the tool calls")


def run(source, args, seconds=20, network=False):
"""Run `run(**args)` from `source` in a process of its own; what it returned. CodeError says what went wrong
(with the end of what it printed)."""
seconds = max(1, min(int(seconds or 20), MAX_SECONDS))
with tempfile.TemporaryDirectory(prefix="lens-tool-") as tmp:
cmd = [sys.executable, "-I", "-S", "-B", "-c", RUNNER, str(seconds), str(MEMORY), "1" if network else "0"]
if not network and netns():
cmd = [netns(), "-rn", *cmd]
env = {"PATH": "/usr/bin:/bin", "HOME": tmp, "TMPDIR": tmp, "LANG": "C.UTF-8", "PYTHONDONTWRITEBYTECODE": "1"}
stdin = json.dumps(source) + "\n" + json.dumps(args) + "\n"
try:
p = subprocess.run(
cmd, input=stdin, capture_output=True, text=True, cwd=tmp, env=env, timeout=seconds + 5, start_new_session=True
)
except subprocess.TimeoutExpired:
raise CodeError(f"the code ran longer than {seconds} seconds and was stopped") from None
out, err = p.stdout[-MAX_OUTPUT:], p.stderr[-2000:]
head, sep, result = out.rpartition("\n\x00LENS-RESULT\x00")
if p.returncode != 0 or not sep:
why = (err.strip().splitlines() or [f"it exited with {p.returncode}"])[-1]
if p.returncode in (-9, -24, 137):
why = "it used more CPU time or memory than a code tool may"
raise CodeError(f"the code failed: {why}")
try:
value = json.loads(result)
except ValueError:
raise CodeError("the code's result isn't JSON") from None
printed = head.strip()
return {"result": value, **({"printed": printed[-4000:]} if printed else {})}
51 changes: 44 additions & 7 deletions fastapi_backend/app/domain/extensions.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,9 @@
- A tool has parameters, an effect and a body. A `read` tool answers at once; a `change` tool waits for the person's
approval, like the built-in tools that change data (a setup conversation, which acts, runs it at once). Bodies:
`prompt` (the model, with the parameters written into a template) and `http` (a web request, to public addresses
only, unless an admin allowed a private network for web pages) and `graph` (drawn on the canvas: tool_nodes.py).
only, unless an admin allowed a private network for web pages), `graph` (drawn on the canvas: tool_nodes.py) and
`python` (code with a `run(**args)` function, in a process of its own: code_tools.py). Only admins write Python
tools, and one runs only while an admin owns its extension.
- A skill is instructions with a line saying when to use them. The assistant sees each skill's name and that line, and
reads the instructions (`use_skill`) only when the skill applies.
- A hook runs on an event: `message` (a question arrives), `before_tool`, `after_tool`, `answer` (an answer was
Expand All @@ -28,6 +30,8 @@
import urllib.parse
import urllib.request

import yaml

from . import auth, store

R = store.R
Expand All @@ -37,7 +41,7 @@
PARAM_RX = re.compile(r"^[a-z][a-z0-9_]{0,40}$")
PARAM_KINDS = {"text": "string", "number": "number", "integer": "integer", "bool": "boolean", "json": "object", "list": "array"}
EFFECTS = ("read", "change")
RUNS = ("prompt", "http", "graph")
RUNS = ("prompt", "http", "graph", "python")
METHODS = ("GET", "POST", "PUT", "PATCH", "DELETE")
EVENTS = ("message", "before_tool", "after_tool", "answer")
ACTIONS = {
Expand Down Expand Up @@ -117,6 +121,24 @@ def _check_run(run, params, me, db=None):
if db is None:
raise ValueError("a canvas tool is checked when it's saved")
return {"type": t, "graph": tool_nodes.check(db, run.get("graph") or {}, params)}
< BD0 /td> if t == "python": # code, run apart from the server (code_tools.py)
from . import code_tools

if set(run) - {"type", "code", "seconds", "network"}:
raise ValueError("a Python tool is {type, code, seconds, network}")
if not me["admin"]:
raise ValueError("only admins write Python tools; draw it on the canvas or use a prompt or web tool instead")
code = _text(run.get("code"), "the tool's code", MAX_TEXT * 5, required=True)
try:
code_tools.check(code)
except code_tools.CodeError as e:
raise ValueError(str(e)) from None
seconds = run.get("seconds", 20)
if isinstance(seconds, bool) or not isinstance(seconds, int) or not 1 <= seconds <= code_tools.MAX_SECONDS:
raise ValueError(f"seconds is a whole number from 1 to {code_tools.MAX_SECONDS}")
if not isinstance(run.get("network", False), bool):
raise ValueError("network is true or false")
return {"type": t, "code": code, "seconds": seconds, "network": bool(run.get("network"))}
if t == "prompt":
if set(run) - {"type", "system", "prompt"}:
raise ValueError("a prompt tool is {type, system, prompt}")
Expand Down Expand Up @@ -284,8 +306,6 @@ def check_manifest(m, me, db=None):
def parse_manifest(text):
"""A manifest written as code: YAML or JSON, or Markdown with YAML frontmatter, whose body is a skill's
instructions or a prompt tool's prompt. ValueError when it can't be read."""
import yaml

text = (text or "").lstrip("")
body = None
fm = re.match(r"^---\s*\n(.*?)\n---\s*(?:\n(.*))?$", text, re.S)
Expand All @@ -312,19 +332,27 @@ def parse_manifest(text):
return head


class _Dumper(yaml.SafeDumper):
"""YAML with text over several lines (a Python tool's code) written as a block, as people write it."""


_Dumper.add_representer(
str,
lambda d, s: d.represent_scalar("tag:yaml.org,2002:str", s, style="|" if "\n" in s else None),
)


def to_manifest(ext):
"""An extension as a manifest people can read, change and import again: Markdown with frontmatter for skills and
prompt tools, YAML for the rest."""
import yaml

spec = json.loads(json.dumps(ext["spec"]))
head = store.clean({"name": ext["name"], "kind": ext["kind"], "title": ext.get("title"), "description": ext.get("description")})
body = None
if ext["kind"] == "skill":
body = spec.pop("instructions")
elif ext["kind"] == "tool" and spec["run"]["type"] == "prompt":
body = spec["run"].pop("prompt")
text = yaml.safe_dump({**head, **spec}, sort_keys=False, allow_unicode=True, width=120)
text = yaml.dump({**head, **spec}, Dumper=_Dumper, sort_keys=False, allow_unicode=True, width=120)
return f"---\n{text}---\n\n{body}\n" if body is not None else text


Expand Down Expand Up @@ -569,6 +597,8 @@ def __init__(self, db, me):
for it in items:
if it["kind"] == "hook" and not trusted:
continue
if it["kind"] == "tool" and it["spec"]["run"]["type"] == "python" and not _owner_admin(db, g):
continue # code runs only while an admin owns it
item = {
"ext": g["id"],
"version": g["version"],
Expand Down Expand Up @@ -689,6 +719,13 @@ def run_tool(db, cfg, spec, args, model=None, toolbox=None):
from . import tool_nodes

return {"result": tool_nodes.run(db, cfg, run["graph"], args, toolbox)}
if run["type"] == "python":
from . import code_tools

try:
return code_tools.run(run["code"], args, seconds=run.get("seconds"), network=run.get("network"))
except code_tools.CodeError as e:
raise ValueError(str(e)) from None
return http_call(cfg, run, args)


Expand Down
78 changes: 78 additions & 0 deletions fastapi_backend/tests/api/test_extensions.py
Original file line number Diff line number Diff line change
Expand Up @@ -361,3 +361,81 @@ def test_made_in_chat_behind_an_approval(app, db, cfg, folder, new_client, llm):
c.post(f"/api/v1/approvals/{a['id']}", headers=h, json={"decision": "approve"})
got = c.get(f"/api/v1/extensions/{out['extension']}", headers=h).json()
assert (got["version"], got["enabled"]) == (2, False) and "Swedish" in got["spec"]["params"][1]["options"]


PY_TOOL = """name: word_count
kind: tool
description: Count the words in a piece of text.
params:
- {name: text, kind: text, required: true}
run:
type: python
seconds: 10
code: |
def run(text):
print("counting")
return {"words": len(text.split())}
"""


def test_python_tools_run_apart_and_only_admins_write_them(app, db, cfg, folder, new_client):
s = Assist(app, db, cfg, folder, new_client)
c, h = s.cl["editor"]
ca, ha = s.cl["admin"]
# only admins write code tools
r = c.post("/api/v1/extensions", headers=h, json={"text": PY_TOOL})
assert r.status_code == 400 and "only admins" in r.json()["detail"]
bad = ca.post("/api/v1/extensions/check", headers=ha, json={"text": PY_TOOL.replace("def run(text)", "def go(text)")})
assert bad.status_code == 400 and "run(**args)" in bad.json()["detail"]
bad = ca.post("/api/v1/extensions/check", headers=ha, json={"text": PY_TOOL.replace("print(", "print((")})
assert bad.status_code == 400 and "doesn't compile" in bad.json()["detail"]

r = ca.post("/api/v1/extensions", headers=ha, json={"text": PY_TOOL})
assert r.status_code == 200, r.text
eid = r.json()["id"]
got = ca.get(f"/api/v1/extensions/{eid}", headers=ha).json()
assert "code: |" in got["manifest"] # the code goes back out as a block, as it was written
again = extensions.parse_manifest(got["manifest"])
assert again["run"]["code"] == got["spec"]["run"]["code"]
out = ca.post(f"/api/v1/extensions/{eid}/test", headers=ha, json={"args": {"text": "one two three"}}).json()["output"]
assert out == {"result": {"words": 3}, "printed": "counting"}

# shared with the editor's namespace, it runs for them; they can't change its code
ca.patch(f"/api/v1/extensions/{eid}", headers=ha, json={"visibility": "namespace", "namespaces": ["pods"]})
me = {"id": db.one("SELECT record::id(id) AS id FROM account WHERE email = 'ed@x.io'")["id"], "email": "ed@x.io"}
assert "word_count" in ai_tools.Toolbox(db, cfg, me, {s.pods}, {s.pods}, {}, None).ext.tools
# once its owner is no longer an admin, it stops running
root = db.one("SELECT record::id(id) AS id FROM account WHERE email = 'root@x.io'")["id"]
db.q("UPDATE $r SET admin = false", r=extensions.R("account", root))
assert "word_count" not in ai_tools.Toolbox(db, cfg, me, {s.pods}, {s.pods}, {}, None).ext.tools


def run_code(code, args=None, **kw):
from app.domain import code_tools

return code_tools.run(code, args or {}, **kw)


def test_python_code_is_kept_apart_from_the_server(monkeypatch):
from app.domain import code_tools

monkeypatch.setenv("LENS_SECRET_FOR_TEST", "s3cret")
assert "LENS_SECRET_FOR_TEST" not in run_code("import os\ndef run():\n return list(os.environ)")["result"]
# its own folder only
assert run_code("def run():\n open('notes.txt', 'w').write('x')\n return open('notes.txt').read()")["result"] == "x"
for code, why in [
("def run():\n return open('/etc/passwd').read()", "own folder"),
("def run():\n open('/tmp/lens-escape', 'w')", "own folder"),
("import subprocess\ndef run():\n return subprocess.run(['id']).returncode", "other programs"),
("import os\ndef run():\n return os.system('id')", "other programs"),
("import socket\ndef run():\n socket.create_connection(('example.com', 80), timeout=2)", "no network"),
("def run():\n while True: pass", "stopped"),
("def run():\n raise RuntimeError('boom')", "boom"),
("def run():\n return 1", None),
]:
if why is None:
assert run_code(code)["result"] == 1
continue
with pytest.raises(code_tools.CodeError) as e:
run_code(code, seconds=2)
assert why in str(e.value) or (why == "stopped" and "CPU time" in str(e.value)), (code, str(e.value))
Loading
Loading
0