okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

nekomimi.petrpc: pass host calls to routes as a method and arguments, and reply once59e13872026-10-10

"""File editing, search, and cell timing are exercised through the daemon's real Python tool."""

import json
import unittest

from harness import Albedo, Provider, Reply, python, text


class PythonToolsTests(unittest.TestCase):
    def setUp(self):
        self.arguments = None

        def script(request):
            if request["messages"][-1].get("role") == "user":
                if self.arguments is not None:
                    return Reply("python", tool_arguments=self.arguments)
                return python(self.code)
            return text("done")

        self.provider = Provider(script)
        self.addCleanup(self.provider.close)
        self.app = Albedo(self.provider)
        self.app.__enter__()
        self.addCleanup(self.app.__exit__, None, None, None)

    def run_cell(self, code, session):
        """The result of one prompt's cell, the latest in the session."""
        self.code = code
        self.app.prompt(session, "use the Python tools").close()
        self.app.idle(session)
        results = [
            json.loads(part["value"])
            for entry in self.app.history(session)["items"]
            if entry["kind"] == "tool_result" and entry["tool"]["name"] == "python"
            for part in entry["content"]
            if part["kind"] == "json" and part["field"] == "result"
        ]
        self.assertEqual(results[-1]["status"], "ok", results[-1])
        return results[-1]

    def execute(self, code):
        return self.run_cell(code, self.app.session())["output"]

    def test_host_replies_preserve_escaped_call_ids_and_argument_errors(self):
        output = self.execute(
            "kernel = sys.modules['__main__']\n"
            "key = 'host reply \" \\\\ λ\\n'\n"
            "pending = asyncio.get_running_loop().create_future()\n"
            "kernel.PENDING[key] = pending\n"
            "try:\n"
            "    kernel.send({'type': 'call', 'id': key, 'method': 'mail.submit', 'args': []})\n"
            "    print(json.dumps(await asyncio.wait_for(pending, 5)))\n"
            "finally:\n    kernel.PENDING.pop(key, None)"
        )
        self.assertEqual(
            json.loads(output),
            {
                "ok": False,
                "code": "invalid",
                "message": "mail.submit(to, body) takes two strings",
            },
        )

    def test_deleting_shadowed_tools_recovers_the_original_bindings(self):
        (self.app.workspace / "note.txt").write_text("recovered\n")
        session = self.app.session()
        output = self.run_cell(
            "original_tools = (files, run, jobs, cells, output, show_image)\n"
            "files, run, jobs, cells, output, show_image = range(6)\n"
            "assert (files, run, jobs, cells, output, show_image) == tuple(range(6))\n"
            "del files, run, jobs, cells, output, show_image\n"
            "assert (files, run, jobs, cells, output, show_image) == original_tools\n"
            "def recovered_read():\n    return files.read('note.txt')\n"
            "print(recovered_read())",
            session,
        )["output"]
        self.assertIn("recovered", output)
        output = self.run_cell(
            "assert (files, run, jobs, cells, output, show_image) == original_tools\n"
            "print(recovered_read())\n"
            "import builtins\nassert not hasattr(builtins, 'files')",
            session,
        )["output"]
        self.assertIn("recovered", output)

    def test_ambiguous_edit_requires_a_hint_and_preserves_other_matches(self):
        note = self.app.workspace / "note.txt"
        note.write_text("same\nother\nsame\n")
        output = self.execute(
            "try:\n    files.edit('note.txt', 'same', 'changed')\n"
            "except Exception as exc:\n    print('ambiguous:', exc)\n"
            "files.edit('note.txt', 'same', 'changed', line_hint=3)\n"
            "print(files.read('note.txt'))"
        )
        self.assertEqual(note.read_text(), "same\nother\nchanged\n")
        self.assertIn("ambiguous:", output)
        self.assertIn("same", output)
        self.assertIn("changed", output)

    def test_an_edit_miss_says_when_only_whitespace_differs(self):
        (self.app.workspace / "note.txt").write_text("if ok:\n    run()\n")
        output = self.execute(
            "for old in ('if ok:\\n\\trun()', 'absent'):\n"
            "    try:\n        files.edit('note.txt', old, 'x')\n"
            "    except ValueError as error:\n        print(str(error).splitlines()[0])"
        )
        lines = output.splitlines()
        self.assertIn("whitespace differs", lines[0])
        self.assertNotIn("whitespace differs", lines[1])

    def test_find_searches_workspace_and_returns_context(self):
        (self.app.workspace / "note.txt").write_text("before\nneedle\nafter\n")
        output = self.execute(
            "rows = await files.find('needle', '.', context=1)\n"
            "print(rows)\nprint(files.read('note.txt', start_line=2, end_line=2))"
        )
        self.assertIn("note.txt", output)
        self.assertIn("needle", output)
        self.assertIn("before", output)
        self.assertIn("after", output)

    def test_an_empty_find_says_what_was_searched_and_skipped(self):
        (self.app.workspace / ".ignore").write_text("ignored.txt\n")
        (self.app.workspace / "ignored.txt").write_text("needle\n")
        (self.app.workspace / ".dot.txt").write_text("needle\n")
        (self.app.workspace / "plain.txt").write_text("nothing\n")
        output = self.execute(
            "print(await files.find('needle', '.'))\n"
            "print(await files.find('needle', '.', hidden=True, ignored=True))"
        )
        self.assertIn("no matches for 'needle'", output)
        self.assertIn("1 files searched", output)
        self.assertIn("skipped hidden files, .gitignore/.ignore rules", output)
        self.assertIn("ignored.txt", output)
        self.assertIn("rerun with hidden=True and/or ignored=True", output)
        self.assertEqual(output.count("no matches"), 1, output)

    def test_find_in_a_missing_root_raises_instead_of_returning_nothing(self):
        output = self.execute(
            "try:\n    await files.find('needle', 'nowhere')\n"
            "except FileNotFoundError as exc:\n    print('missing:', exc)"
        )
        self.assertIn("missing: nowhere not found", output)

    def test_a_file_result_is_awaitable_and_offers_content_and_text(self):
        (self.app.workspace / "note.txt").write_text("hello\n")
        output = self.execute(
            "text = await files.read('note.txt')\n"
            "print('awaited:', text.strip())\n"
            "print('content:', (await files.read('note.txt')).content.strip())\n"
            "print('text:', files.read('note.txt').text.strip())"
        )
        self.assertIn("awaited: 1 | hello", output)
        self.assertIn("content: 1 | hello", output)
        self.assertIn("text: 1 | hello", output)

    def test_paths_searches_a_list_of_directories_with_a_list_of_globs(self):
        for directory, name in [("a", "one.md"), ("b", "two.txt"), ("c", "three.md")]:
            (self.app.workspace / directory).mkdir()
            (self.app.workspace / directory / name).write_text("x")
        output = self.execute(
            "rows = await files.paths(None, ['a', 'b'], glob=['*.md', '*.txt'])\n"
            "print(sorted(str(row) for row in rows))"
        )
        self.assertIn("one.md", output)
        self.assertIn("two.txt", output)
        self.assertNotIn("three.md", output)

    def test_paths_globs_a_path_from_the_workspace_root(self):
        (self.app.workspace / "src" / "deep").mkdir(parents=True)
        (self.app.workspace / "src" / "deep" / "stream_test.go").write_text("x")
        output = self.execute(
            "print([str(row) for row in await files.paths('src/*/*_test.go')])\n"
            "print([str(row) for row in await files.paths('./src/*/*_test.go')])"
        )
        self.assertEqual(output.count("['src/deep/stream_test.go']"), 2, output)

    def test_file_reads_decode_complete_lines_across_chunk_boundaries(self):
        fixtures = [
            (b"", []),
            (b"\n\n", ["", ""]),
            (b"a\r\r\n", ["a", ""]),
            (b"a\n", ["a"]),
            (b"a\nlast", ["a", "last"]),
            (
                "a\nb\rc\r\nd\ve\ff\x1cg\x1dh\x1ei\x85j\u2028k\u2029last".encode(),
                list("abcdefghijk") + ["last"],
            ),
            (b"a\xffb\xe2\x82", ["a\ufffdb\ufffd"]),
        ]
        for padding in (65_534, 65_535, 65_536):
            prefix = b"x" * padding
            for separator in (
                b"\r\n",
                "\x85".encode(),
                "\u2028".encode(),
                "\u2029".encode(),
            ):
                fixtures.append((prefix + separator + b"tail", ["x" * padding, "tail"]))
            fixtures.append((prefix + "é\n尾".encode(), ["x" * padding + "é", "尾"]))
            fixtures.append(
                (
                    prefix + b"\xe2\x82\nend\xe2\x82",
                    ["x" * padding + "\ufffd", "end\ufffd"],
                )
            )
        for index, (data, _lines) in enumerate(fixtures):
            (self.app.workspace / f"decode-{index}.txt").write_bytes(data)
        # Keep large expected lines inside the kernel, out of the tool output.
        small_expected = [lines for _data, lines in fixtures[:7]]
        output = self.execute(
            f"expected = {small_expected!r}\n"
            "for padding in (65534, 65535, 65536):\n"
            "    expected.extend([['x' * padding, 'tail']] * 4)\n"
            "    expected.append(['x' * padding + 'é', '尾'])\n"
            "    expected.append(['x' * padding + '\ufffd', 'end\ufffd'])\n"
            "for index, lines in enumerate(expected):\n"
            "    result = files.read(f'decode-{index}.txt', max_chars=200000)\n"
            "    numbered = '\\n'.join(f'{number:>6} | {line}' for number, line in enumerate(lines, 1))\n"
            "    if lines:\n"
            "        assert result == numbered, (index, len(result), len(numbered))\n"
            "    else:\n"
            "        assert ' | ' not in result\n"
            "        assert all(detail in result for detail in ('decode-0.txt', '0 lines', 'line 1'))\n"
            "print('decoded complete lines')"
        )
        self.assertEqual(output.strip(), "decoded complete lines")

    def test_file_reads_enforce_budgets_and_character_limit_precedence(self):
        (self.app.workspace / "budget.txt").write_text("a\nbb\nc\nd\n")
        (self.app.workspace / "numbered.txt").write_text("\n" * 999_999 + "x\nz")
        output = self.execute(
            "first = '     1 | a'\n"
            "second = '     2 | bb'\n"
            "both = first + '\\n' + second\n"
            "assert files.read('budget.txt', end_line=2, max_chars=23) == both\n"
            "assert files.read('budget.txt', end_line=1, max_chars=11) == first\n"
            "rejected = files.read('budget.txt', max_chars=10)\n"
            "assert ' | ' not in rejected\n"
            "assert all(detail in rejected for detail in ('line 1', '10 characters', 'max_chars=10', 'start_line=1', 'end_line=1', 'max_chars=11'))\n"
            "for budget, shown, unshown, resume in ((22, [first], 'lines 2-4', 2), (23, [first, second], 'lines 3-4', 3)):\n"
            "    result = files.read('budget.txt', max_chars=budget).splitlines()\n"
            "    assert result[:-1] == shown\n"
            "    assert all(detail in result[-1] for detail in (f'max_chars={budget}', unshown, f'start_line={resume}'))\n"
            "result = files.read('budget.txt', limit=2, max_chars=23).splitlines()\n"
            "assert result[:-1] == [first, second]\n"
            "assert all(detail in result[-1] for detail in ('limit=2', '2 more through line 4', 'start_line=3'))\n"
            "result = files.read('budget.txt', limit=3, max_chars=23).splitlines()\n"
            "assert result[:-1] == [first, second]\n"
            "assert all(detail in result[-1] for detail in ('max_chars=23', 'lines 3-3', 'start_line=3'))\n"
            "assert 'limit=' not in result[-1]\n"
            "assert files.read('numbered.txt', start_line=1000000, end_line=1000000, max_chars=12) == '1000000 | x'\n"
            "rejected = files.read('numbered.txt', start_line=1000000, end_line=1000000, max_chars=11)\n"
            "assert ' | ' not in rejected\n"
            "assert all(detail in rejected for detail in ('line 1000000', '11 characters', 'max_chars=11', 'start_line=1000000', 'end_line=1000000', 'max_chars=12'))\n"
            "print('budgets and numbering checked')"
        )
        self.assertEqual(output.strip(), "budgets and numbering checked")

    def test_file_reads_count_requested_windows_and_eof_exactly(self):
        (self.app.workspace / "window.txt").write_text("a\nb\nc\nd")
        pattern = "a\r\nb\rc\nd\ve\ff\x1cg\x1dh\x1ei\x85j\u2028k\u2029last\n".encode()
        (self.app.workspace / "bulk.txt").write_bytes(
            b"head\r\n" + pattern * 10_000 + b"tail\xe2\x82"
        )
        output = self.execute(
            "assert files.read('window.txt', start_line=2, end_line=3) == '     2 | b\\n     3 | c'\n"
            "assert files.read('window.txt', start_line=3, end_line=99) == '     3 | c\\n     4 | d'\n"
            "for options, row, remaining, last, resume in (({'start_line': 2}, '     2 | b', 2, 4, 3), ({'start_line': 2, 'end_line': 3}, '     2 | b', 1, 3, 3), ({'start_line': 3, 'end_line': 99}, '     3 | c', 1, 4, 4)):\n"
            "    result = files.read('window.txt', limit=1, **options).splitlines()\n"
            "    assert result[:-1] == [row]\n"
            "    assert all(detail in result[-1] for detail in ('limit=1', f'{remaining} more through line {last}', f'start_line={resume}'))\n"
            "assert files.read('window.txt', start_line=4, limit=1) == '     4 | d'\n"
            "empty = files.read('window.txt', start_line=5, end_line=99)\n"
            "assert all(detail in empty for detail in ('window.txt', '4 lines', 'line 5'))\n"
            "assert ' | ' not in empty\n"
            "result = files.read('bulk.txt', start_line=2, limit=1).splitlines()\n"
            "assert result[:-1] == ['     2 | a']\n"
            "assert all(detail in result[-1] for detail in ('limit=1', '120000 more through line 120002', 'start_line=3'))\n"
            "result = files.read('bulk.txt', start_line=2, end_line=100000, limit=1).splitlines()\n"
            "assert result[:-1] == ['     2 | a']\n"
            "assert all(detail in result[-1] for detail in ('limit=1', '99998 more through line 100000', 'start_line=3'))\n"
            "result = files.read('bulk.txt', start_line=120001, limit=1).splitlines()\n"
            "assert result[:-1] == ['120001 | last']\n"
            "assert all(detail in result[-1] for detail in ('limit=1', '1 more through line 120002', 'start_line=120002'))\n"
            "assert files.read('bulk.txt', start_line=120001, end_line=120001) == '120001 | last'\n"
            "assert files.read('bulk.txt', start_line=120002) == '120002 | tail\\ufffd'\n"
            "print('window counts checked')"
        )
        self.assertEqual(output.strip(), "window counts checked")

    def test_file_reads_count_giant_lines_without_shortening_them(self):
        (self.app.workspace / "giant.txt").write_text(
            "head\n" + "x" * 1_048_576 + "\ntail\n"
        )
        (self.app.workspace / "unicode.txt").write_text("é" * 110_000)
        output = self.execute(
            "assert files.read('giant.txt', start_line=3) == '     3 | tail'\n"
            "assert files.read('giant.txt', end_line=1) == '     1 | head'\n"
            "rejected = files.read('giant.txt', start_line=2, end_line=2)\n"
            "assert ' | ' not in rejected\n"
            "assert all(detail in rejected for detail in ('line 2', '1048585 characters', 'max_chars=16000', '200000'))\n"
            "assert 'max_chars=1048586' not in rejected\n"
            "result = files.read('giant.txt', max_chars=32).splitlines()\n"
            "assert result[:-1] == ['     1 | head']\n"
            "assert all(detail in result[-1] for detail in ('max_chars=32', 'lines 2-3', 'start_line=2'))\n"
            "result = files.read('giant.txt', limit=1).splitlines()\n"
            "assert result[:-1] == ['     1 | head']\n"
            "assert all(detail in result[-1] for detail in ('limit=1', '2 more through line 3', 'start_line=2'))\n"
            "for budget in (16000, 110009):\n"
            "    rejected = files.read('unicode.txt', max_chars=budget)\n"
            "    assert ' | ' not in rejected\n"
            "    assert all(detail in rejected for detail in ('line 1', '110009 characters', f'max_chars={budget}', 'start_line=1', 'end_line=1', 'max_chars=110010'))\n"
            "assert files.read('unicode.txt', max_chars=110010) == '     1 | ' + 'é' * 110000\n"
            "print('giant lines counted and returned intact')"
        )
        self.assertEqual(output.strip(), "giant lines counted and returned intact")

    def test_file_reads_validate_ranges_before_opening_and_report_path_errors(self):
        output = self.execute(
            "for options in ({'start_line': 0}, {'start_line': 2, 'end_line': 1}, {'limit': 0}, {'max_chars': 0}, {'max_chars': 200001}):\n"
            "    try:\n"
            "        files.read('absent.txt', **options)\n"
            "    except ValueError:\n"
            "        pass\n"
            "    else:\n"
            "        raise AssertionError(options)\n"
            "try:\n"
            "    files.read('absent.txt')\n"
            "except FileNotFoundError as error:\n"
            "    assert 'absent.txt' in str(error)\n"
            "else:\n"
            "    raise AssertionError('missing file accepted')\n"
            "try:\n"
            "    files.read('.')\n"
            "except IsADirectoryError as error:\n"
            "    assert \"files.ls('.')\" in str(error)\n"
            "else:\n"
            "    raise AssertionError('directory accepted')\n"
            "print('range and path errors preserved')"
        )
        self.assertEqual(output.strip(), "range and path errors preserved")

    def test_cells_get_the_api_a_model_reaches_for(self):
        output = self.execute(
            "job = run('echo', 'hello')\n"
            "await job\n"
            "assert (await job.tail()).strip() == 'hello'\n"
            "assert (await job.head(lines=1)).strip() == 'hello'\n"
            "assert output.read(job.id, 0, 3) == 'hel'\n"
            "try:\n"
            "    output.read(job.id, limit=10**7)\n"
            "except ValueError as error:\n"
            "    assert 'offset=' in str(error), error\n"
            "else:\n"
            "    raise AssertionError('over-cap limit shortened silently')\n"
            "from albedo import files as imported\n"
            "assert imported is files\n"
            "assert re.fullmatch('a+', 'aa') and Path('.').is_dir()\n"
            "assert hashlib.md5(b'').hexdigest() and json.dumps(os.sep)\n"
            "try:\n"
            "    run('true', timeout=120000)\n"
            "except ValueError as error:\n"
            "    assert 'seconds' in str(error) and 'milliseconds' in str(error), error\n"
            "try:\n"
            "    time.sleep(2)\n"
            "except PermissionError as error:\n"
            "    assert 'run()' in str(error) and 'asyncio.sleep' in str(error), error\n"
            "time.sleep(0.01)\n"
            "await asyncio.sleep(0.01)\n"
            "print('api shapes ok')"
        )
        self.assertEqual(output.strip(), "api shapes ok")

    def test_a_cell_reports_how_long_it_ran(self):
        session = self.app.session()
        slept = self.run_cell("import asyncio\nawait asyncio.sleep(0.3)", session)
        self.assertGreaterEqual(slept["duration"], 0.3)
        self.assertLess(slept["duration"], 5)
        # The journal keeps it, so cells.info reports it after the fact.
        info = self.run_cell(
            f"print((await cells.info({slept['cell_id']!r}))['duration'])", session
        )
        self.assertEqual(float(info["output"]), slept["duration"])

    def call(self, arguments, session):
        """The result of one python call made with exactly these arguments."""
        self.arguments = arguments
        self.app.prompt(session, "call the Python tool").close()
        self.app.idle(session)
        self.arguments = None
        return [
            json.loads(part["value"])
            for entry in self.app.history(session)["items"]
            if entry["kind"] == "tool_result" and entry["tool"]["name"] == "python"
            for part in entry["content"]
            if part["kind"] == "json" and part["field"] == "result"
        ][-1]

    def test_timeout_is_optional(self):
        result = self.call({"code": "print('no timeout')"}, self.app.session())
        self.assertEqual(result["status"], "ok", result)
        self.assertEqual(result["output"].strip(), "no timeout")

    def test_an_unusable_timeout_keeps_the_code_for_a_rerun(self):
        session = self.app.session()
        rejected = self.call(
            {"code": "open('made.txt', 'w').write('hi')", "timeout_ms": "soon"},
            session,
        )
        self.assertIn("timeout_ms", rejected["error"])
        self.assertFalse((self.app.workspace / "made.txt").exists())
        self.run_cell(f"await cells.run({rejected['cell_id']!r})", session)
        self.assertEqual((self.app.workspace / "made.txt").read_text(), "hi")