Brunobkr/llama.cpp_AlgMor24_github

1import os2import shutil3import subprocess4 5import pytest6from utils import *7 8server: ServerProcess9 10# project root, used as the search directory for grep_search/file_glob_search11PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", ".."))12 13# marker for the grep_search test to find in this file14GREP_MARKER = "llama_cpp_test_tools_builtin_marker_grep_search"15 16# image the container runtime tests run their shell in17CONTAINER_IMAGE = "busybox"18 19 20@pytest.fixture(autouse=True)21def create_server():22    global server23    server = ServerPreset.router()24    server.server_tools = "all"25 26 27def call_tool(name: str, params: dict, headers: dict | None = None) -> dict:28    res = server.make_request("POST", "/tools", data={"tool": name, "params": params}, headers=headers)29    assert res.status_code == 200, res.body30    assert "error" not in res.body, res.body31    return res.body32 33 34def call_tool_expect_error(name: str, params: dict) -> str:35    res = server.make_request("POST", "/tools", data={"tool": name, "params": params})36    assert res.status_code == 200, res.body37    assert "error" in res.body, res.body38    return res.body["error"]39 40 41def test_tools_builtin_grep_search():42    global server43    server.start()44 45    res = call_tool("grep_search", {46        "path": PROJECT_ROOT,47        "pattern": GREP_MARKER,48        "include": "test_tools_builtin.py",  # bare pattern -> matches basename at any depth49    })50    text = res["plain_text_response"]51    assert "test_tools_builtin.py" in text52    assert GREP_MARKER in text53    assert "Total matches: 1" in text54 55 56def test_tools_builtin_read_file():57    global server58    server.start()59 60    this_file = os.path.join(PROJECT_ROOT, "tools", "server", "tests", "unit", "test_tools_builtin.py")61    res = call_tool("read_file", {"path": this_file})62    text = res["plain_text_response"]63    assert GREP_MARKER in text64    assert "def test_tools_builtin_read_file" in text65 66 67def test_tools_builtin_write_then_edit_file():68    global server69    server.start()70 71    log_path = os.path.join(PROJECT_ROOT, "test.log")72    try:73        write_res = call_tool("write_file", {"path": log_path, "content": "line1\nline2\nline3\n"})74        assert write_res["result"] == "file written successfully"75 76        read_before = call_tool("read_file", {"path": log_path})77        assert read_before["plain_text_response"] == "line1\nline2\nline3\n"78 79        edit_res = call_tool("edit_file", {80            "path": log_path,81            "edits": [82                {"old_text": "line2", "new_text": "line2-edited"},83                {"old_text": "line3\n", "new_text": "line3\nline4\n"},84            ],85        })86        assert edit_res["result"] == "file edited successfully"87        assert edit_res["edits_applied"] == 288 89        read_after = call_tool("read_file", {"path": log_path})90        assert read_after["plain_text_response"] == "line1\nline2-edited\nline3\nline4\n"91    finally:92        if os.path.exists(log_path):93            os.remove(log_path)94 95 96def test_tools_builtin_edit_file_rejects_non_unique_old_text():97    global server98    server.start()99 100    log_path = os.path.join(PROJECT_ROOT, "test.log")101    try:102        call_tool("write_file", {"path": log_path, "content": "dup\ndup\n"})103        err = call_tool_expect_error("edit_file", {104            "path": log_path,105            "edits": [{"old_text": "dup", "new_text": "changed"}],106        })107        assert "unique" in err108    finally:109        if os.path.exists(log_path):110            os.remove(log_path)111 112 113def test_tools_builtin_exec_shell_command_stream():114    global server115    server.start()116 117    events = list(server.make_stream_request("POST", "/tools", data={118        "tool": "exec_shell_command",119        "params": {"command": "echo hello"},120        "stream": True,121    }))122 123    assert len(events) >= 2124    assert events[-1]["done"] is True125    assert not events[-1].get("error")126    chunks = "".join(e["chunk"] for e in events[:-1])127    assert "hello" in chunks128    assert "[exit code: 0]" in chunks129 130 131def test_tools_builtin_cwd_header():132    global server133    server.start()134 135    cwd_dir = os.path.join(PROJECT_ROOT, "tools", "server", "tests", "unit")136    headers = {"x-tool-cwd": cwd_dir}137 138    res = call_tool("read_file", {"path": "test_tools_builtin.py"}, headers=headers)139    assert GREP_MARKER in res["plain_text_response"]140 141    # exec_shell_command should also run with that directory as its working directory:142    # writing to a relative filename must land inside cwd_dir143    marker_name = "llama_cpp_test_tools_builtin_cwd_marker.txt"144    marker_path = os.path.join(cwd_dir, marker_name)145    try:146        command = f"echo hello > {marker_name}"147        call_tool("exec_shell_command", {"command": command}, headers=headers)148        assert os.path.exists(marker_path)149    finally:150        if os.path.exists(marker_path):151            os.remove(marker_path)152 153 154def _container_engine_unavailable_reason(engine: str) -> str | None:155    """None if `engine` can run the image these tests use, otherwise the reason it can't."""156    engine_bin = shutil.which(engine)157    if engine_bin is None:158        return f"{engine} is not installed"159    try:160        # a daemon that answers `info` still cannot run a linux image when it serves windows161        # containers, so probe the image itself, which also pulls it before the tests162        subprocess.run([engine_bin, "run", "--rm", CONTAINER_IMAGE, "true"], capture_output=True, timeout=60, check=True)163    except Exception as e:164        return f"{engine} cannot run {CONTAINER_IMAGE}: {e}"165    return None166 167 168@pytest.fixture(params=["docker", "podman"])169def container_engine(request):170    engine = request.param171    reason = _container_engine_unavailable_reason(engine)172    if reason is not None:173        pytest.skip(reason)  # ty: ignore[too-many-positional-arguments, invalid-argument-type]174    return engine175 176 177@pytest.fixture178def container_id(container_engine: str):179    proc = subprocess.run(180        [container_engine, "run", "-d", "--rm", CONTAINER_IMAGE, "sleep", "300"],181        capture_output=True, text=True,182    )183    if proc.returncode != 0:184        pytest.skip(f"failed to start {container_engine} container: {proc.stderr.strip()}")  # ty: ignore[too-many-positional-arguments, invalid-argument-type]185 186    cid = proc.stdout.strip()187    try:188        yield cid189    finally:190        subprocess.run([container_engine, "rm", "-f", cid], capture_output=True)191 192 193def test_tools_builtin_runtime_header(container_engine: str, container_id: str):194    global server195    server.start()196 197    headers = {"x-tool-runtime": f"{container_engine}-container:{container_id}", "x-tool-cwd": "/tmp"}198 199    write_res = call_tool("write_file", {"path": "test.log", "content": "hello container\n"}, headers=headers)200    assert write_res["result"] == "file written successfully"201 202    read_res = call_tool("read_file", {"path": "test.log"}, headers=headers)203    assert read_res["plain_text_response"] == "hello container\n"204 205    exec_res = call_tool("exec_shell_command", {"command": "cat test.log"}, headers=headers)206    assert "hello container" in exec_res["plain_text_response"]207 208 209def test_tools_builtin_runtime_header_unknown_scheme():210    global server211    server.start()212 213    # an unknown runtime must fail, never silently fall back to running on the host214    res = server.make_request("POST", "/tools",215                              data={"tool": "exec_shell_command", "params": {"command": "echo hi"}},216                              headers={"x-tool-runtime": "fake:does-not-exist"})217    assert res.status_code == 500, res.body218    assert "unknown tool runtime" in str(res.body)219 220 221def test_tools_builtin_runtime_header_rejects_ssh_option_injection():222    global server223    server.start()224 225    # ssh reads options from its argv, so a target starting with '-' must be rejected226    res = server.make_request("POST", "/tools",227                              data={"tool": "exec_shell_command", "params": {"command": "echo hi"}},228                              headers={"x-tool-runtime": "ssh:-oProxyCommand=touch /tmp/pwned"})229    assert res.status_code == 500, res.body230    assert "invalid ssh target" in str(res.body)231 232 233@pytest.mark.parametrize("engine", ["docker", "podman"])234def test_tools_builtin_runtime_header_rejects_container_option_injection(engine: str):235    global server236    server.start()237 238    # the container id lands on the `<engine> exec` command line, so an id that looks239    # like an option must be rejected240    res = server.make_request("POST", "/tools",241                              data={"tool": "exec_shell_command", "params": {"command": "echo hi"}},242                              headers={"x-tool-runtime": f"{engine}-container:--privileged"})243    assert res.status_code == 500, res.body244    assert "invalid container id" in str(res.body)245 246 247def test_tools_builtin_docker_runtime_cleans_up_spawned_container():248    # docker-only: this reads the container hostname to get the spawned id, which only docker249    # sets to the short id. podman is covered by the attach path above250    reason = _container_engine_unavailable_reason("docker")251    if reason is not None:252        pytest.skip(reason)  # ty: ignore[too-many-positional-arguments, invalid-argument-type]253 254    global server255    server.server_tools_runtime = f"docker:{CONTAINER_IMAGE}"256    server.start()257 258    # exec_shell_command runs inside the container spawned for --tools-runtime; docker sets259    # the container's hostname to its own short id, so this also tells us which one to check260    res = call_tool("exec_shell_command", {"command": "hostname"})261    container_id = res["plain_text_response"].splitlines()[0].strip()262    assert len(container_id) >= 8, res263 264    running = subprocess.run(265        ["docker", "inspect", "-f", "{{.State.Running}}", container_id],266        capture_output=True, text=True,267    )268    assert running.returncode == 0 and running.stdout.strip() == "true", running.stderr269 270    server.stop()271 272    # a clean server shutdown must stop and remove the container it spawned (it runs with --rm),273    # not leave it behind as an abandoned child274    leftover = subprocess.run(["docker", "inspect", container_id], capture_output=True, text=True)275    assert leftover.returncode != 0, f"container {container_id} was not cleaned up after server exit"276 277 278def test_tools_builtin_edit_file_rejects_overlapping_edits():279    global server280    server.start()281 282    log_path = os.path.join(PROJECT_ROOT, "test.log")283    try:284        call_tool("write_file", {"path": log_path, "content": "line1\nline2\n"})285        err = call_tool_expect_error("edit_file", {286            "path": log_path,287            "edits": [288                {"old_text": "line1\nline2", "new_text": "a"},289                {"old_text": "line2", "new_text": "b"},290            ],291        })292        assert "overlap" in err293    finally:294        if os.path.exists(log_path):295            os.remove(log_path)296 297 298def test_tools_builtin_file_glob_search_type_dir(tmp_path):299    global server300    server.start()301 302    (tmp_path / "project-alpha" / "src").mkdir(parents=True)303    (tmp_path / "project-alpha" / "README.md").write_text("alpha")304    (tmp_path / "project-alpha" / "src" / "main.cpp").write_text("int main() {}")305    (tmp_path / "project-beta").mkdir()306    (tmp_path / "project-beta" / "notes.txt").write_text("beta")307 308    res = call_tool("file_glob_search", {"path": str(tmp_path), "type": "dir"})309    text = res["plain_text_response"]310    assert "project-alpha/" in text311    assert "project-beta/" in text312    assert "project-alpha/src/" in text313    assert "README.md" not in text314    types = {e["path"]: e["type"] for e in res["entries"]}315    assert types["project-alpha"] == "dir"316    assert types["project-alpha/src"] == "dir"317 318    res_all = call_tool("file_glob_search", {"path": str(tmp_path), "type": "all", "include": "*proj*"})319    paths = [e["path"] for e in res_all["entries"]]320    assert "project-alpha" in paths321    assert "project-beta" in paths322 323 324def test_tools_builtin_file_glob_search_max_depth_and_limit(tmp_path):325    global server326    server.start()327 328    (tmp_path / "a" / "b" / "c").mkdir(parents=True)329    (tmp_path / "top.txt").write_text("top")330    (tmp_path / "a" / "mid.txt").write_text("mid")331    (tmp_path / "a" / "b" / "deep.txt").write_text("deep")332 333    res = call_tool("file_glob_search", {"path": str(tmp_path), "max_depth": 1})334    assert "top.txt" in res["plain_text_response"]335    assert "mid.txt" not in res["plain_text_response"]336 337    res = call_tool("file_glob_search", {"path": str(tmp_path), "max_depth": 2})338    assert "mid.txt" in res["plain_text_response"]339    assert "deep.txt" not in res["plain_text_response"]340 341    res = call_tool("file_glob_search", {"path": str(tmp_path), "limit": 1})342    assert len(res["entries"]) == 1343    assert "Total matches: 3" in res["plain_text_response"]344 345 346def test_tools_builtin_file_glob_search_junk_dirs(tmp_path):347    global server348    server.start()349 350    (tmp_path / "build" / "nested").mkdir(parents=True)351    (tmp_path / "build" / "artifact.txt").write_text("built")352    (tmp_path / "src").mkdir()353    (tmp_path / "src" / "main.cpp").write_text("int main() {}")354 355    # a junk directory stays selectable as a working directory356    res = call_tool("file_glob_search", {"path": str(tmp_path), "type": "dir", "max_depth": 1})357    assert "build" in [e["path"] for e in res["entries"]]358 359    # but it is never walked, so nothing inside it shows up360    res = call_tool("file_glob_search", {"path": str(tmp_path), "type": "all"})361    paths = [e["path"] for e in res["entries"]]362    assert "src/main.cpp" in paths363    assert "build/artifact.txt" not in paths364    assert "build/nested" not in paths365 366 367def test_tools_builtin_file_glob_search_rejects_invalid_type(tmp_path):368    global server369    server.start()370 371    err = call_tool_expect_error("file_glob_search", {"path": str(tmp_path), "type": "bogus"})372    assert "invalid type" in err373 374 375def test_tools_builtin_cwd_header_overrides_model_param(tmp_path):376    global server377    server.start()378 379    workdir = tmp_path / "workdir"380    workdir.mkdir()381    (workdir / "marker.txt").write_text("marker")382 383    # a model-provided "cwd" in the params is overridden by the x-tool-cwd header384    res = call_tool("read_file", {"path": "marker.txt", "cwd": "/definitely/not/a/real/path"},385                    headers={"x-tool-cwd": str(workdir)})386    assert "marker" in res["plain_text_response"]387 388 389def test_tools_builtin_cwd_relative_paths(tmp_path):390    global server391    server.start()392 393    workdir = tmp_path / "workdir"394    workdir.mkdir()395    (workdir / "rel.txt").write_text("relative-content")396 397    headers = {"x-tool-cwd": str(workdir)}398 399    # relative paths in file tools resolve against the header cwd400    res = call_tool("read_file", {"path": "rel.txt"}, headers=headers)401    assert "relative-content" in res["plain_text_response"]402 403    res = call_tool("write_file", {"path": "sub/out.txt", "content": "written"}, headers=headers)404    assert (workdir / "sub" / "out.txt").read_text() == "written"405 406    res = call_tool("file_glob_search", {"path": ".", "include": "*.txt"}, headers=headers)407    assert "rel.txt" in res["plain_text_response"]408 409    # absolute paths are unaffected by the cwd410    other = tmp_path / "other"411    other.mkdir()412    (other / "abs.txt").write_text("absolute-content")413    res = call_tool("read_file", {"path": str(other / "abs.txt")}, headers=headers)414    assert "absolute-content" in res["plain_text_response"]415