From 73e917005fdec552689f09f5bd35df4d8378308a Mon Sep 17 00:00:00 2001 From: khushalsonawat Date: Wed, 11 Nov 2026 01:49:43 +0530 Subject: [PATCH 1/3] fix(harness): decode a Python file before searching it when its codec can spell a name A coding line can name a codec that spells ASCII letters in other bytes (raw_unicode_escape, utf-7, the iso2022 codecs, hz, mac-arabic and mac-farsi), so a file could import playwright or traceai_openai while its raw bytes never held the name. The browser and instrumentation scans searched those raw bytes, and the import walk decoded such files before ending their lines, or read unicode-escape as UTF-8. Every codec not proven safe to search as bytes is now decoded first, as Python reads source (lines ended, a final line end, NUL cut), under the scan's one 4 MiB decode total. The 81 safe codecs are proven by a sweep test. A decoded-first file passes the gate when a needle is in its text, or in its raw bytes when the text isn't ASCII, so no file base lets in is stopped. The walk reads the same way, with its own total, and follows the union of both readings for a unicode-escape file whose text holds a CR, charged once. Refs: TH-9490 Signed-off-by: khushalsonawat --- .github/workflows/sdk-smoke.yml | 1 + src/fi/alk/harness/credentials.py | 69 +- src/fi/alk/harness/python_source.py | 189 ++- tests/harness/test_agent_type_refusal.py | 6 +- tests/harness/test_python_source_codecs.py | 1351 ++++++++++++++++++++ 5 files changed, 1579 insertions(+), 37 deletions(-) create mode 100644 tests/harness/test_python_source_codecs.py diff --git a/.github/workflows/sdk-smoke.yml b/.github/workflows/sdk-smoke.yml index 728d41b49..ec3c66a73 100644 --- a/.github/workflows/sdk-smoke.yml +++ b/.github/workflows/sdk-smoke.yml @@ -95,6 +95,7 @@ jobs: tests/harness/test_hosted_entrypoint_provision_retry.py tests/harness/test_instrumentation_evidence.py tests/harness/test_agent_type_refusal.py + tests/harness/test_python_source_codecs.py tests/harness/test_executor_completed_counts.py tests/harness/test_reading_entry.py tests/harness/test_agent_reading.py diff --git a/src/fi/alk/harness/credentials.py b/src/fi/alk/harness/credentials.py index 93bad5bf9..4284e5174 100644 --- a/src/fi/alk/harness/credentials.py +++ b/src/fi/alk/harness/credentials.py @@ -26,7 +26,7 @@ from fi.simulate.runtime.spec import SecretRef -from .python_source import ParseBudget, parsed, source_encoding, utf8_size +from .python_source import EscapeBudget, ParseBudget, decoded_first, parsed, utf8_size from .secrets import process_secret_environment_names from .source_folder import SubfolderNotFound, agent_folder, agent_scope, in_agent_scope, is_within @@ -1434,6 +1434,10 @@ def script_module(path: Path, name: str) -> Path | None: return hit return None + # One total for the whole walk: bundling has no lookup limit and no parse total, so this is + # all that bounds what decoding first costs there, a unicode-escape file's second reading + # included (at most this much more parsing per walk). + escapes = EscapeBudget() seen: set[Path] = set() pending = deque(sorted(_code_in(folder))) own = set(pending) # Read once already, so importing one doesn't queue it again. @@ -1460,25 +1464,62 @@ def script_module(path: Path, name: str) -> Path | None: if len(raw) > _MAX_FILE_BYTES: continue # Grew since it was looked at. bytes_read += len(raw) - # As read_text reads it, line ends included; a .py in the encoding its coding line names, - # as Python reads it. - encoding = (source_encoding(raw) if path.suffix == ".py" else None) or "utf-8-sig" - try: - # Strictly: replacing runs Python's own error handler for each byte in most codecs, - # 60 ms a megabyte, where UTF-8's is in C. Python itself refuses such a file. - text = raw.decode(encoding) - except (UnicodeError, LookupError): - # Read as UTF-8, as is a codec source_encoding refuses: a crafted file Python reads in - # one of those can then give fewer imports here than it does to Python. - text = raw.decode("utf-8-sig", errors="replace") - content = text.replace("\r\n", "\n").replace("\r", "\n") + # A .py in a codec that can spell a name in other bytes is read as Python reads it while + # the walk's total lasts: decoded after its line ends, so hz's `~` before a CRLF joins. + taken, decoded, encoding = ( + decoded_first(raw, escapes) if path.suffix == ".py" else (False, None, None) + ) + read_before = None + if taken: + if decoded is None: + # Python refuses the file, so it imports nothing. + continue + content = decoded.replace("\r\n", "\n").replace("\r", "\n") + if encoding is None and "\r" in decoded: + # unicode-escape is also read as UTF-8, where `\r` is two characters: decoded, it's + # a CR the parser ends a comment or a string at where Python reads on, so what that + # reading imports is followed too. + read_before = raw.decode("utf-8-sig", errors="replace") + read_before = read_before.replace("\r\n", "\n").replace("\r", "\n") + else: + # As read_text reads it, line ends included; a .py in the encoding its coding line + # names, as Python reads it. + encoding = encoding or "utf-8-sig" + try: + # Strictly: replacing runs Python's own error handler for each byte in most + # codecs, 60 ms a megabyte, where UTF-8's is in C. Python itself refuses such a + # file. + text = raw.decode(encoding) + except (UnicodeError, LookupError): + # Read as UTF-8, as is a codec refused by name (and unicode-escape past the + # total): a crafted file Python reads in one of those can then give fewer imports + # here than it does to Python. + text = raw.decode("utf-8-sig", errors="replace") + content = text.replace("\r\n", "\n").replace("\r", "\n") hits: list[Path] = [] if path.suffix == ".py": # A file with no import in it can't lead anywhere, so it doesn't use up the parse total # before the files that do (Pine's TH-9490 review); a file the total cuts says so. if "import" not in content: continue - names = list(dict.fromkeys(_python_imports(content, parse_limit, parse_budget))) + readings = [content] if read_before is None else [content, read_before] + larger = max(readings, key=utf8_size) + if len(readings) == 1 or utf8_size(larger) > parse_limit: + # Each reading is charged as it's parsed; one past the limit is neither. + found = [ + name + for reading in readings + for name in _python_imports(reading, parse_limit, parse_budget) + ] + elif parse_budget is None or parse_budget.spend(larger): + # One file read two ways is charged once, as the larger reading, so it spends no + # more of the total than it did read one way. + found = [ + name for reading in readings for name in _python_imports(reading, parse_limit) + ] + else: + found = [] + names = list(dict.fromkeys(found)) cut |= parse_budget is not None and parse_budget.left < 0 # One file can hold more names than the walk may try; the rest aren't tried. cut |= len(names) > import_limit - imports_tried diff --git a/src/fi/alk/harness/python_source.py b/src/fi/alk/harness/python_source.py index 047159e3f..06de87d4b 100644 --- a/src/fi/alk/harness/python_source.py +++ b/src/fi/alk/harness/python_source.py @@ -32,10 +32,11 @@ # token is charged its line. This much is at most about 0.03 s of that work on 3.12.3. Over 5,000 # real files of 4 to 64 KB, 64 KB of code comes to about 0.5 M, and 1.9 M at the most. LARGEST_TOKENIZED_WORK = 8 * 1024 * 1024 -# What one scan decodes in unicode-escape, in bytes. Python runs such a file, and its escapes can -# spell any name (`import \x70laywright`), so it's decoded as Python does; but its decode looks up -# each \N{...} name, and 3,000 files of 128 KB held the save 6-7 s with no total. This much -# decodes in about 0.05 s; past it, the file is searched as bytes. +# What one scan decodes before searching, in bytes. Python runs a file in a codec that can spell a +# name in other bytes (`import \x70laywright` in unicode-escape, `import +AHA-laywright` in +# utf-7), so such a file is decoded as Python does; but unicode-escape's decode looks up each +# \N{...} name, and 3,000 files of 128 KB held the save 6-7 s with no total. This much decodes in +# about 0.05 s; past it, the file is searched as bytes. LARGEST_ESCAPED_TOTAL = 4 * 1024 * 1024 @@ -45,6 +46,99 @@ _LONGEST_MARK_RUN = 32 # Text codecs a coding line may name that aren't read: see _coding. _UNREAD_CODECS = frozenset({"idna", "punycode", "undefined", "unicode-escape"}) +# The codecs a coding line may name in which no byte form spells a name character, NUL, CR or LF +# but its own single byte, nor spells nothing: a needle missing from the raw bytes is missing from +# the text Python reads too. The sweep in test_python_source_codecs proves each one, so a +# codec missing from here is decoded before it's searched. +_RAW_SAFE_CODECS = frozenset( + { + "ascii", + "big5", + "big5hkscs", + "charmap", + "cp1006", + "cp1125", + "cp1250", + "cp1251", + "cp1252", + "cp1253", + "cp1254", + "cp1255", + "cp1256", + "cp1257", + "cp1258", + "cp437", + "cp720", + "cp737", + "cp775", + "cp850", + "cp852", + "cp855", + "cp856", + "cp857", + "cp858", + "cp860", + "cp861", + "cp862", + "cp863", + "cp864", + "cp865", + "cp866", + "cp869", + "cp874", + "cp932", + "cp949", + "cp950", + "euc_jis_2004", + "euc_jisx0213", + "euc_jp", + "euc_kr", + "gb18030", + "gb2312", + "gbk", + "hp-roman8", + "iso8859-1", + "iso8859-10", + "iso8859-11", + "iso8859-13", + "iso8859-14", + "iso8859-15", + "iso8859-16", + "iso8859-2", + "iso8859-3", + "iso8859-4", + "iso8859-5", + "iso8859-6", + "iso8859-7", + "iso8859-8", + "iso8859-9", + "johab", + "koi8-r", + "koi8-t", + "koi8-u", + "kz1048", + "mac-croatian", + "mac-cyrillic", + "mac-greek", + "mac-iceland", + "mac-latin2", + "mac-roman", + "mac-romanian", + "mac-turkish", + "palmos", + "ptcp154", + "shift_jis", + "shift_jis_2004", + "shift_jisx0213", + "tis-620", + "utf-8", + "utf-8-sig", + } +) +# Refused by name in _coding: never decoded, whatever the total. +_NEVER_DECODED_FIRST = frozenset({"idna", "punycode", "undefined"}) +# Python reads a lone CR as a line end before it decodes, so hz's `~` and a line end join. +_CR_TO_LF = bytes.maketrans(b"\r", b"\n") _NON_ASCII_RUN = re.compile(r"[^\x00-\x7f]{2,}") @@ -65,7 +159,8 @@ def spend(self, text: str) -> bool: class EscapeBudget: - """What is left for one scan to decode in unicode-escape, in bytes.""" + """What is left for one scan to decode before searching, in bytes: every file read that way (a + text codec not proven safe to search as bytes, unicode-escape among them) is charged here.""" def __init__(self, total: int | None = None) -> None: self.left = LARGEST_ESCAPED_TOTAL if total is None else total @@ -114,11 +209,12 @@ def read_module( in the encoding its coding line or byte-order mark names (UTF-8 without one), or, given ``needles``, naming none of them. - Given ``escapes``, a file in unicode-escape is decoded as Python does while the scan's total - lasts, and the needles are looked for in what it decodes to. Given ``undecoded``, a file that - names a needle but isn't read because its coding line names a codec refused by name (or - unicode-escape past the total) is put there as bytes: Python still runs it (plain ASCII under - idna as it would without the line), so a scan looking for names can still find them. + Given ``escapes``, a file in a codec not proven safe to search as bytes (unicode-escape, + utf-7, hz and others) is decoded as Python does while the scan's total lasts, and a needle + counts in what it decodes to, or in its bytes when that isn't ASCII. Given ``undecoded``, a file that names a needle + but isn't read because its coding line names a codec refused by name (or unicode-escape past + the total) is put there as bytes: Python still runs it (plain ASCII under idna as it would + without the line), so a scan looking for names can still find them. """ try: # A link can name a device or a pipe, which reads without end. @@ -131,18 +227,17 @@ def read_module( if len(raw) > limit: return None coding = _coding(raw) - if ( - coding is not None - and coding[1].name == "unicode-escape" - and escapes is not None - and escapes.spend(raw) - ): - try: - # Python ends lines first, so a backslash before a CRLF joins the lines it ends. - text = raw.replace(b"\r\n", b"\n").replace(b"\r", b"\n").decode("unicode-escape") - except UnicodeError: + taken, text = _decoded_first(raw, coding, escapes) + if taken: + if text is None: return None - if needles and not any(needle.decode("ascii") in text for needle in needles): + # The bytes too, as for any other file: the parser folds a fullwidth name the text spells + # (`ｐlaywright`) to the needle, and a needle in the bytes is what lets such a file in. On + # ASCII text nothing folds, so whatever a scan can find is in the text itself. + if needles and not any( + needle.decode("ascii") in text or not text.isascii() and needle in raw + for needle in needles + ): return None return text # A literal match: a name spelled in other characters that Python folds to the same one, @@ -177,6 +272,58 @@ def source_encoding(raw: bytes) -> str | None: return encoding +def decoded_first( + raw: bytes, escapes: EscapeBudget | None +) -> tuple[bool, str | None, str | None]: + """``raw`` read as ``read_module`` reads a file before searching it, for a reader that + searches no needles: whether it was decoded first, its text (None when Python refuses it), + and the encoding ``source_encoding`` gives, to read a file that wasn't. The coding line is + read once for both.""" + coding = _coding(raw) + taken, text = _decoded_first(raw, coding, escapes) + encoding = None if coding is None or coding[1].name in _UNREAD_CODECS else coding[0] + return taken, text, encoding + + +def _decoded_first( + raw: bytes, coding: tuple[str, codecs.CodecInfo] | None, escapes: EscapeBudget | None +) -> tuple[bool, str | None]: + """Whether ``raw``, in the codec ``coding`` names, is decoded before it's searched, and its + text as Python reads it: None when Python refuses it. + + A codec not proven safe to search as bytes can spell a name in other bytes, so its file is + decoded while ``escapes`` lasts. Past it, or with no total, it isn't taken, and the caller + reads it as any other file. + """ + if coding is None or escapes is None: + return False, None + encoding, info = coding + if ( + info.name in _RAW_SAFE_CODECS + or info.name in _NEVER_DECODED_FIRST + # Not a text encoding, such as rot13 or base64: bytes.decode refuses it before reading. + or not getattr(info, "_is_text_encoding", True) + ): + return False, None + if not escapes.spend(raw): + return False, None + if b"\r" in raw: + # Python ends lines before it decodes, so a backslash or hz's `~` before a CRLF joins + # the lines it ends. + raw = raw.replace(b"\r\n", b"\n").translate(_CR_TO_LF) + if not raw.endswith(b"\n"): + # Python ends the last line too: a file ending in an open escape still runs. + raw += b"\n" + try: + text = raw.decode(encoding) + except (UnicodeError, LookupError): + # LookupError only from a codec registered at run time: Python's text codecs don't. + return True, None + # Python reads nothing past a NUL its decode makes at a line's start, and refuses the file + # for one elsewhere: so a name after it is never run. + return True, text.partition("\0")[0] + + def _first_lines(raw: bytes) -> bytes: """The first two lines of ``raw``, each ended as Python ends a line (LF, CRLF or a lone CR), joined by LF: found where they end, so the rest of the file is never rewritten.""" diff --git a/tests/harness/test_agent_type_refusal.py b/tests/harness/test_agent_type_refusal.py index 97d3cf884..dcb24c5a8 100644 --- a/tests/harness/test_agent_type_refusal.py +++ b/tests/harness/test_agent_type_refusal.py @@ -701,9 +701,11 @@ def test_a_file_in_a_refused_codec_still_names_its_tools( assert (decoded, labels, parsed) == ([], [], []) -def test_an_escaped_import_in_a_unicode_escape_file_is_found(tmp_path): +@pytest.mark.parametrize("ending", [b"\n", b""], ids=["final-lf", "no-final-lf"]) +def test_an_escaped_import_in_a_unicode_escape_file_is_found(tmp_path, ending): # Python decodes the escapes before it parses, so the import is real and so is the finding. - (tmp_path / "helper.py").write_bytes(b"# coding: unicode-escape\nimport \\x70laywright\n") + # It ends the last line itself, so a file without a final line end runs the same. + (tmp_path / "helper.py").write_bytes(b"# coding: unicode-escape\nimport \\x70laywright" + ending) assert browser_tools_found(tmp_path) == ["playwright"] diff --git a/tests/harness/test_python_source_codecs.py b/tests/harness/test_python_source_codecs.py new file mode 100644 index 000000000..7fc63bd46 --- /dev/null +++ b/tests/harness/test_python_source_codecs.py @@ -0,0 +1,1351 @@ +"""A Python file whose codec can spell a name in other bytes is read as Python reads it. + +Each row that says Python runs a name checks it with ``compile``, which is what an import runs. +""" + +from __future__ import annotations + +import ast +import codecs +import collections +import encodings +import encodings.aliases +import pkgutil +import string +import sys +import types +from pathlib import Path + +import pytest + +from fi.alk.harness import agent_type_check, python_source +from fi.alk.harness.agent_type_check import ( + browser_tools_found, + browser_tools_in_agent, + code_the_agent_reaches, +) +from fi.alk.harness.credentials import imported_by_agent +from fi.alk.harness.instrumentation_evidence import detect_instrumentation +from fi.alk.harness.python_source import ( + _RAW_SAFE_CODECS, + EscapeBudget, + _coding, + _decoded_first, + read_module, +) + +# Python reads a decoded CR as a line end, and a line it starts as blank, only from 3.12. +_FROM_312 = pytest.mark.skipif( + sys.version_info < (3, 12), reason="Python reads a decoded CR this way from 3.12" +) +_ISO2022 = ( + "iso2022_jp", + "iso2022_jp_1", + "iso2022_jp_2", + "iso2022_jp_2004", + "iso2022_jp_3", + "iso2022_jp_ext", + "iso2022_kr", +) +# The codecs Python accepts on a coding line that spell a name character from other bytes. +_SPELLING_CODECS = frozenset( + { + "raw-unicode-escape", + "unicode-escape", + "utf-7", + "hz", + "mac-arabic", + "mac-farsi", + *_ISO2022, + } +) +_UTF7 = b"# coding: utf-7\n" +_ESCAPED = b"# coding: unicode-escape\n" +_RAW_ESCAPED = b"# coding: raw-unicode-escape\n" + + +def _in(codec: str, body: bytes) -> bytes: + return f"# coding: {codec}\n".encode() + body + + +def _names_python_runs(raw: bytes) -> set[str]: + """The names the code compiled from ``raw`` uses, its functions' included.""" + names: set[str] = set() + pending = [compile(raw, "", "exec", dont_inherit=True)] + while pending: + code = pending.pop() + names.update(code.co_names) + pending.extend(c for c in code.co_consts if isinstance(c, types.CodeType)) + return names + + +def _imported(tree: ast.AST) -> set[str]: + """The top-level modules ``tree`` imports.""" + found: set[str] = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + found.update(alias.name.split(".")[0] for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.module: + found.add(node.module.split(".")[0]) + return found + + +def _write(folder: Path, raw: bytes, name: str = "agent.py") -> Path: + folder.mkdir(parents=True, exist_ok=True) + path = folder / name + path.write_bytes(raw) + return path + + +def _shared(root: Path, agent: bytes) -> Path: + """A repository whose agent, in agents/a, may import common.web, which imports the names + both scans look for.""" + root = root.resolve() + _write(root / "agents" / "a", agent) + _write(root / "common", b"", "__init__.py") + _write(root / "common", b"import playwright\nimport traceai_openai\n", "web.py") + return root + + +def _scan(folder: Path, raw: bytes, name: str) -> object: + """What the reader that looks for ``name`` makes of a file holding ``raw``: whether the walk + reaches common/web.py, the instrumentation verdict, or the browser tools.""" + if name == "common.web": + repo = _shared(folder, raw) + return repo / "common" / "web.py" in imported_by_agent(repo, "agents/a") + _write(folder, raw) + if name.startswith("traceai_"): + return detect_instrumentation(folder) + return browser_tools_found(folder) + + +def _found(name: str) -> object: + if name == "common.web": + return True + return "traceai" if name.startswith("traceai_") else [name.replace("_", "-")] + + +def _missed(name: str) -> object: + if name == "common.web": + return False + return "none" if name.startswith("traceai_") else [] + + +@pytest.mark.parametrize( + "raw, name", + [ + pytest.param( + _RAW_ESCAPED + b"import \\u0070laywright\n", + "playwright", + id="raw-unicode-escape", + ), + pytest.param( + _ESCAPED + b"import \\x70laywright\n", "playwright", id="unicode-escape" + ), + pytest.param(_UTF7 + b"import +AHA-laywright\n", "playwright", id="utf-7"), + *( + pytest.param( + _in(codec, b"import play\x1b(Bwright\n"), "playwright", id=codec + ) + for codec in _ISO2022 + ), + pytest.param(_in("hz", b"import play~{~}wright\n"), "playwright", id="hz"), + pytest.param(_in("hz", b"import play~\nwright\n"), "playwright", id="hz-lf"), + pytest.param( + b"# coding: hz\r\nimport play~\r\nwright\r\n", "playwright", id="hz-crlf" + ), + pytest.param(b"# coding: hz\rimport play~\rwright\r", "playwright", id="hz-cr"), + # 0xDF is `_` here, and `playwright` before it would be found in the bytes. + pytest.param( + _in("mac-arabic", b"import undetected\xdfchromedriver\n"), + "undetected_chromedriver", + id="mac-arabic", + ), + pytest.param( + _in("mac-farsi", b"import undetected\xdfchromedriver\n"), + "undetected_chromedriver", + id="mac-farsi", + ), + ], +) +def test_a_browser_import_spelled_in_a_codec_that_can_spell_one_is_found( + tmp_path, raw, name +): + assert name in _names_python_runs(raw) + _write(tmp_path, raw) + assert browser_tools_found(tmp_path) == [name.replace("_", "-")] + + +@pytest.mark.parametrize( + "raw", + [ + pytest.param( + _RAW_ESCAPED + b"import \\u0074raceai_openai\n", id="raw-unicode-escape" + ), + pytest.param(_UTF7 + b"import +AHQ-raceai_openai\n", id="utf-7"), + pytest.param(_in("mac-arabic", b"import traceai\xdfopenai\n"), id="mac-arabic"), + pytest.param( + _in("iso2022_jp", b"import trace\x1b(Bai_openai\n"), id="iso2022_jp" + ), + pytest.param(_in("hz", b"import trace~{~}ai_openai\n"), id="hz"), + ], +) +def test_traceai_spelled_in_another_codec_is_found(tmp_path, raw): + assert "traceai_openai" in _names_python_runs(raw) + _write(tmp_path, raw) + assert detect_instrumentation(tmp_path) == "traceai" + + +@pytest.mark.parametrize( + "agent", + [ + # hz joins a `~` and the line end after it, once Python has made the CRLF an LF. + pytest.param(b"# coding: hz\r\nimport commo~\r\nn.web\r\n", id="hz-crlf"), + pytest.param(_ESCAPED + b"import \\x63ommon.web\n", id="unicode-escape"), + # Decoded whole, the NUL refuses the parse; Python reads up to it. + pytest.param(_UTF7 + b"import +AGM-ommon.web\n+AAA-\n", id="utf-7"), + ], +) +def test_the_walk_follows_an_import_spelled_in_another_codec(tmp_path, agent): + assert "common.web" in _names_python_runs(agent) + repo = _shared(tmp_path, agent) + assert repo / "common" / "web.py" in imported_by_agent(repo, "agents/a") + assert browser_tools_in_agent(repo, "agents/a") == ["playwright"] + reached = code_the_agent_reaches(repo, "agents/a") + assert detect_instrumentation(repo / "agents" / "a", reached) == "traceai" + + +@pytest.mark.parametrize( + "raw", + [ + pytest.param( + _in("iso2022_jp", b"import play\x1b(Bwright\n#\x1b"), id="iso2022_jp" + ), + pytest.param(_in("hz", b"import play~{~}wright\n#~"), id="hz"), + pytest.param(_ESCAPED + b"import \\x70laywright\n#\\", id="unicode-escape"), + ], +) +def test_a_file_ending_in_an_open_sequence_is_read(tmp_path, raw): + assert "playwright" in _names_python_runs(raw) + _write(tmp_path, raw) + assert browser_tools_found(tmp_path) == ["playwright"] + + +@pytest.mark.parametrize( + "raw, found", + [ + pytest.param( + _UTF7 + b"import +AHA-laywright\n+AAA-\nimport os\n", + True, + id="utf-7-before", + ), + pytest.param( + _UTF7 + b"import os\n+AAA-\nimport +AHA-laywright\n", + False, + id="utf-7-after", + ), + pytest.param( + _RAW_ESCAPED + b"import \\u0070laywright\n\\u0000\nimport os\n", + True, + id="raw-unicode-escape-before", + ), + pytest.param( + _RAW_ESCAPED + b"import os\n\\u0000\nimport playwright\n", + False, + id="raw-unicode-escape-after", + ), + pytest.param( + _ESCAPED + b"import \\x70laywright\n\\x00\nimport os\n", + True, + id="unicode-escape-before", + ), + pytest.param( + _ESCAPED + b"import os\n\\x00\nimport playwright\n", + False, + id="unicode-escape-after", + ), + ], +) +def test_a_decoded_nul_ends_the_text(tmp_path, raw, found): + assert ("playwright" in _names_python_runs(raw)) is found + _write(tmp_path, raw) + assert browser_tools_found(tmp_path) == (["playwright"] if found else []) + + +@pytest.mark.parametrize( + "raw, name", + [ + # A line that starts with a decoded CR, before a token that can't start a statement. + pytest.param( + _UTF7 + b"import playwright\n+AA0-1\n", "playwright", id="cr-then-1" + ), + pytest.param( + _UTF7 + b"import playwright\n+AA0-@dec\ndef f(): pass\n", + "playwright", + id="cr-then-decorator", + ), + pytest.param( + _ESCAPED + b"import playwright\n\\r(1)\n", + "playwright", + id="ue-cr-then-call", + ), + pytest.param( + _UTF7 + b"import traceai_openai\n+AA0-1\n", + "traceai_openai", + id="cr-then-1-traceai", + ), + pytest.param( + _UTF7 + b"import common.web\n+AA0-1\n", "common.web", id="cr-then-1-walk" + ), + # A decoded CR in a comment, then code. + pytest.param( + _UTF7 + b"x = f # c+AA0-(1); import playwright\n", + "playwright", + id="comment-cr", + ), + pytest.param( + _ESCAPED + b"x = f # c\\r(1); import playwright\n", + "playwright", + id="ue-comment-cr", + ), + pytest.param( + _UTF7 + b"x = f # c+AA0-(1); import traceai_openai\n", + "traceai_openai", + id="comment-cr-traceai", + ), + # A line of blanks and a `\`, then a line that starts with a decoded CR. + pytest.param( + _UTF7 + b"import playwright\n \\\n +AA0-1\n", + "playwright", + id="continued-blank-then-cr", + ), + pytest.param( + _RAW_ESCAPED + b" \\\n\\u000d-c\nimport playwright\n", + "playwright", + id="rue-continued-blank-then-cr", + ), + # A `\` before a decoded CR and an LF: one line end, as Python 3.12 reads it. + pytest.param( + _UTF7 + b"x = 'a\\+AA0-\nb'\nimport playwright\n", + "playwright", + id="string-continued", + ), + pytest.param( + _UTF7 + b"x = 'a\\+AA0-\nb'\nimport traceai_openai\n", + "traceai_openai", + id="string-continued-traceai", + ), + pytest.param( + _UTF7 + b"import os, \\+AA0-\nplaywright\n", + "playwright", + id="line-continued", + ), + pytest.param( + _ESCAPED + b"import os, \\\\\\r\nplaywright\n", + "playwright", + id="ue-line-continued", + ), + pytest.param( + _ESCAPED + b"import \\\\\\r\n\\x70laywright\n", + "playwright", + id="ue-line-continued-escaped-name", + ), + pytest.param( + _RAW_ESCAPED + b"import os, \\u005c\\u000d\n\\u0070laywright\n", + "playwright", + id="rue-line-continued", + ), + pytest.param( + _UTF7 + b"import \\+AA0-\n+AHA-laywright\n", + "playwright", + id="line-continued-spelled-name", + ), + ], +) +@_FROM_312 +def test_a_decoded_cr_reads_as_todays_code_reads_it(tmp_path, raw, name): + assert name in _names_python_runs(raw) + assert _scan(tmp_path, raw, name) == _found(name) + + +@pytest.mark.parametrize( + "raw, name", + [ + # Python keeps the CR in the string; the parser ends the line there. + pytest.param( + _ESCAPED + b'x = "\\r"\nimport playwright\n', + "playwright", + id="cr-in-string", + ), + # Python ends the comment at the CR and reads on; the parser keeps it in the comment. + pytest.param( + _UTF7 + b"import os # hi+AA0-, playwright\n", + "playwright", + id="comment-cr", + marks=_FROM_312, + ), + pytest.param( + _ESCAPED + b"x = 1 # hi\\r; import playwright\n", + "playwright", + id="ue-comment-cr", + marks=_FROM_312, + ), + # Python reads a line a decoded CR starts as blank, inside the block; the parser as + # a line at column 0, which ends it. + pytest.param( + _UTF7 + b"def f():\n a = 1\n+AA0-(1)\n import +AHA-laywright\n", + "playwright", + id="cr-in-a-block", + marks=_FROM_312, + ), + pytest.param( + _UTF7 + + b"try:\n import +AHA-laywright\n+AA0-[0]\nexcept ImportError:\n pass\n", + "playwright", + id="cr-in-a-try", + marks=_FROM_312, + ), + pytest.param( + _UTF7 + + b"try:\n import +AHQ-raceai_openai\n+AA0-[0]\nexcept ImportError:\n pass\n", + "traceai_openai", + id="cr-in-a-try-traceai", + marks=_FROM_312, + ), + pytest.param( + _UTF7 + b"+AA0-\\+AA0-\n import playwright\n", + "playwright", + id="cr-then-continued", + marks=_FROM_312, + ), + # More failed cookie places than _coding reads before the cookie: read as a refused + # codec's file, searched as bytes. + pytest.param( + b"#" + b"coding:#" * 65 + b" coding: utf-7\nimport +AHA-laywright\n", + "playwright", + id="cookie-after-65-failed-places", + ), + ], +) +def test_the_known_misses_stay_unread(tmp_path, raw, name): + assert name in _names_python_runs(raw) + assert _scan(tmp_path, raw, name) == _missed(name) + + +@pytest.mark.parametrize( + "hidden, scan, found, missed", + [ + pytest.param( + _UTF7 + b"import +AHA-laywright\n", + browser_tools_found, + ["playwright"], + [], + id="browser", + ), + pytest.param( + _UTF7 + b"import +AHQ-raceai_openai\n", + detect_instrumentation, + "traceai", + "none", + id="instrumentation", + ), + ], +) +def test_decode_first_files_share_the_scans_total( + tmp_path, monkeypatch, hidden, scan, found, missed +): + # Names nothing either scan looks for, but is decoded first all the same. + first = _UTF7 + b"import os\n# " + b"-" * 40 + b"\n" + # No larger, so a total of its own would let it through. + assert len(hidden) <= len(first) + for total, expected in ((len(first), missed), (len(first) + len(hidden), found)): + monkeypatch.setattr(python_source, "LARGEST_ESCAPED_TOTAL", total) + folder = tmp_path / str(total) + # Both scans read in name order: a.py first. + _write(folder, first, "a.py") + _write(folder, hidden, "b.py") + assert scan(folder) == expected + + +def test_past_the_total_a_file_reads_as_today(tmp_path, monkeypatch): + monkeypatch.setattr(python_source, "LARGEST_ESCAPED_TOTAL", 0) + _write(tmp_path / "plain", _UTF7 + b"import playwright\n") + _write(tmp_path / "spelled", _UTF7 + b"import +AHA-laywright\n") + assert browser_tools_found(tmp_path / "plain") == ["playwright"] + assert browser_tools_found(tmp_path / "spelled") == [] + + +@pytest.mark.parametrize( + "codec, word", + [ + ("latin-1", "café"), + ("cp1252", "€"), + ("shift_jis", "日本"), + ("gb18030", "中文"), + ("euc_kr", "한국"), + ], +) +def test_a_raw_safe_codec_reads_as_today(tmp_path, codec, word): + # CRLF and no final LF: decoding first would end both lines with an LF. + raw = f"# coding: {codec}\r\nimport playwright\r\nNAME = '{word}'".encode(codec) + budget = EscapeBudget() + assert _decoded_first(raw, _coding(raw), budget) == (False, None) + path = _write(tmp_path, raw) + assert read_module(path, needles=(b"playwright",), escapes=budget) == raw.decode( + codec + ) + assert budget.left == python_source.LARGEST_ESCAPED_TOTAL + + +def test_without_a_total_every_codec_reads_as_today(tmp_path): + unended = _UTF7 + b"import +AHA-laywright" + assert _decoded_first(unended, _coding(unended), None) == (False, None) + assert read_module(_write(tmp_path / "utf-7", unended)) == unended.decode("utf-7") + escaped = _ESCAPED + b"import \\x70laywright\n" + assert read_module(_write(tmp_path / "unicode-escape", escaped)) is None + + +@pytest.mark.parametrize("codec", ["rot13", "base64"]) +def test_a_non_text_codec_is_never_charged(tmp_path, codec): + raw = _in(codec, b"import playwright\n") + assert _coding(raw) is not None + budget = EscapeBudget() + assert read_module(_write(tmp_path, raw), escapes=budget) is None + assert budget.left == python_source.LARGEST_ESCAPED_TOTAL + + +def test_a_text_codec_whose_decode_raises_lookup_error_gives_none(): + def decode(data, errors="strict"): + raise LookupError("no table") + + stub = codecs.CodecInfo(codecs.utf_8_encode, decode, name="lookup-failing") + + def search(name): + return stub if name == "lookup_failing" else None + + codecs.register(search) + try: + coding = ("lookup-failing", codecs.lookup("lookup-failing")) + assert _decoded_first(b"x = 1\n", coding, EscapeBudget()) == (True, None) + finally: + codecs.unregister(search) + + +def test_the_walk_keeps_the_decoded_reading_beside_an_escaped_cr(tmp_path): + # Only the decoded reading spells the import; the escaped CR makes the walk read both. + agent = _ESCAPED + b"X = 1 # \\r\nimport \\x63ommon.web\n" + assert "common.web" in _names_python_runs(agent) + repo = _shared(tmp_path, agent) + assert repo / "common" / "web.py" in imported_by_agent(repo, "agents/a") + + +def test_a_file_read_two_ways_is_charged_once(tmp_path): + pad = b"# " + b"x" * 76 + b"\n" + repo = _shared( + tmp_path, _ESCAPED + b'SEP = "\\r"\nimport common.web\n' + pad * 1200 + ) + _write(repo / "agents" / "a", b"import common.tools\n" + pad * 1200, "b.py") + _write(repo / "common", b"", "tools.py") + stopped: list[bool] = [] + budget = python_source.ParseBudget() + found = imported_by_agent( + repo, + "agents/a", + parse_limit=python_source.LARGEST_PARSED_MODULE, + parse_budget=budget, + stopped=stopped, + ) + assert repo / "common" / "tools.py" in found and stopped == [False] + # Charged once, as the larger reading (the UTF-8 one, where `\\r` is two characters). + parsed = [repo / "agents" / "a" / n for n in ("agent.py", "b.py")] + parsed.append(repo / "common" / "web.py") + assert budget.left == python_source.LARGEST_PARSED_TOTAL - sum( + p.stat().st_size for p in parsed + ) + + +def test_the_union_is_for_unicode_escape_only(tmp_path): + # Read as UTF-8, `+ACM-` is arithmetic and the import is code; decoded, it's a `#`. + agent = _UTF7 + b'"""+AA0-"""\nx = 1 +ACM- 2; import common.web\n' + assert "common" not in _names_python_runs(agent) + repo = _shared(tmp_path, agent) + assert repo / "common" / "web.py" not in imported_by_agent(repo, "agents/a") + + +def test_a_failed_decode_is_still_charged(): + bad = _UTF7 + b"import os\n# +\xff\n" + budget = EscapeBudget() + assert _decoded_first(bad, _coding(bad), budget) == (True, None) + assert budget.left == python_source.LARGEST_ESCAPED_TOTAL - len(bad) + + +@pytest.mark.parametrize( + "agent", + [ + pytest.param(_ESCAPED + b"import common.web # see \\r, x\n", id="comment"), + pytest.param(_ESCAPED + b'SEP = "\\r"\nimport common.web\n', id="string"), + pytest.param(_ESCAPED + b"CR = b'\\r'\nimport common.web\n", id="bytes"), + pytest.param( + _ESCAPED + b'"""\\r"""\nimport \\x63ommon.web\n', id="and-an-escaped-import" + ), + pytest.param( + _ESCAPED + b'SEP = "\\r" # caf\xe9\nimport common.web\n', + id="and-a-latin-1-byte", + ), + ], +) +def test_the_walk_still_follows_a_unicode_escape_file_with_an_escaped_cr( + tmp_path, agent +): + # Decoded, the CR breaks the parse; read as UTF-8 as well, the file is followed. + assert "common.web" in _names_python_runs(agent) + repo = _shared(tmp_path, agent) + assert repo / "common" / "web.py" in imported_by_agent(repo, "agents/a") + assert browser_tools_in_agent(repo, "agents/a") == ["playwright"] + reached = code_the_agent_reaches(repo, "agents/a") + assert detect_instrumentation(repo / "agents" / "a", reached) == "traceai" + + +def test_a_cookie_with_many_coding_words_after_it_is_still_decoded_first(tmp_path): + raw = b"# coding: utf-7" + b" coding" * 65 + b"\nimport +AHA-laywright\n" + assert "playwright" in _names_python_runs(raw) + _write(tmp_path, raw) + assert browser_tools_found(tmp_path) == ["playwright"] + + +def test_the_walk_skips_a_python_file_it_cant_decode(tmp_path): + # The bad byte in a comment, so read as UTF-8 the import would still parse. + agent = _UTF7 + b"import common.web\n# +\xff\n" + with pytest.raises(SyntaxError): + compile(agent, "", "exec", dont_inherit=True) + repo = _shared(tmp_path, agent) + assert repo / "common" / "web.py" not in imported_by_agent(repo, "agents/a") + + +def test_the_walk_reads_scripts_as_before(tmp_path): + folder = tmp_path.resolve() / "agents" / "a" + _write(folder, b"", "util.js") + _write(folder, b"const u = require('./util')\n", "index.js") + assert folder / "util.js" in imported_by_agent(tmp_path, "agents/a") + # A coding line means nothing to a script: utf-7 would spell `./util` here. + _write(folder, _UTF7 + b"const u = require('./u+AHQ-il')\n", "index.js") + assert folder / "util.js" not in imported_by_agent(tmp_path, "agents/a") + + +def test_the_walks_total_is_spent_across_one_walk(tmp_path, monkeypatch): + first = _ESCAPED + b"import os\n# " + b"-" * 20 + b"\n" + # Past the total it's read as UTF-8, where the escape spells nothing. + second = _ESCAPED + b"import \\x63ommon.web\n" + assert len(second) <= len(first) + repo = _shared(tmp_path, b"") + folder = repo / "agents" / "a" + # The walk reads its folder's files in name order. + _write(folder, first, "a.py") + _write(folder, second, "b.py") + for total, followed in ((len(first), False), (len(first) + len(second), True)): + monkeypatch.setattr(python_source, "LARGEST_ESCAPED_TOTAL", total) + assert ( + repo / "common" / "web.py" in imported_by_agent(repo, "agents/a") + ) is followed + + +def _punycode_decodes(monkeypatch) -> list[int]: + import encodings.punycode + + decoded: list[int] = [] + real = encodings.punycode.punycode_decode + monkeypatch.setattr( + encodings.punycode, + "punycode_decode", + lambda *args: decoded.append(len(args[0])) or real(*args), + ) + return decoded + + +def test_the_walk_never_decodes_a_refused_codec(tmp_path, monkeypatch): + import encodings.idna + + decoded = _punycode_decodes(monkeypatch) + labels: list[bytes] = [] + real = encodings.idna.ToUnicode + monkeypatch.setattr( + encodings.idna, "ToUnicode", lambda label: labels.append(label) or real(label) + ) + repo = _shared(tmp_path, b"# coding: punycode\nimport common.web\n-" + b"9" * 64) + _write( + repo / "agents" / "a", + b"# coding: idna\nimport common.web\nHOST = 'xn--bcher-kva.example'\n", + "host.py", + ) + imported_by_agent(repo, "agents/a") + assert (decoded, labels) == ([], []) + + +def test_the_walk_reads_each_coding_line_once(tmp_path, monkeypatch): + asked: list[bytes] = [] + real = python_source._coding + monkeypatch.setattr( + python_source, "_coding", lambda raw: asked.append(raw) or real(raw) + ) + repo = _shared(tmp_path, b"import common.web\n") + _write(repo / "common", b"# a package\n", "__init__.py") + _write(repo / "agents" / "a", b"const u = require('./util')\n", "index.js") + _write(repo / "agents" / "a", b"", "util.js") + imported_by_agent(repo, "agents/a") + python = [ + repo / "agents" / "a" / "agent.py", + repo / "common" / "__init__.py", + repo / "common" / "web.py", + ] + assert sorted(asked) == sorted(path.read_bytes() for path in python) + + +def test_ci_runs_this_file(): + workflows = ( + Path(__file__).resolve().parents[2] / ".github" / "workflows" / "sdk-smoke.yml" + ) + assert "tests/harness/test_python_source_codecs.py" in workflows.read_text( + encoding="utf-8" + ) + + +@pytest.mark.parametrize( + "raw", + [ + # A fullwidth `ｐ`, which the parser folds to `p`; the needle is in a comment's bytes. + pytest.param(_UTF7 + b"import +/1A-laywright\n# +playwrightA-\n", id="utf-7"), + pytest.param( + b"# coding: hz\nimport ~{#p#l#a#y#w#r#i#g#h#t~}\n# ~{playwright~}\n", + id="hz", + ), + ], +) +def test_a_raw_needle_still_lets_the_file_in(tmp_path, raw): + assert "playwright" in _names_python_runs(raw) + _write(tmp_path, raw) + assert browser_tools_found(tmp_path) == ["playwright"] + + +# Every byte form that could join, end or start a name when it decodes. +_NAME_OR_LINE = frozenset(string.ascii_letters + string.digits + "_." + "\0\r\n") +_INCOMPLETE = frozenset({"incomplete multibyte sequence", "unexpected end of data"}) +_DEEPEST_FORM = 4 +# Complete forms by length, for the codecs whose long forms are enumerated by structure. +_UTF8_FORMS = { + # Unicode's scalar values by length: U+0000-U+007F, U+0080-U+07FF, U+0800-U+FFFF less the + # 2,048 surrogates, and U+10000-U+10FFFF. utf-8-sig's byte-order mark is a three-byte form. + 1: 0x80, + 2: 0x800 - 0x80, + 3: 0x10000 - 0x800 - 0x800, + 4: 0x110000 - 0x10000, +} +_FORMS = { + "utf-8": _UTF8_FORMS, + "utf-8-sig": _UTF8_FORMS, + # gb18030 spells each scalar value once: ASCII in one byte, 126 leads (81-FE) by 190 trails + # (40-7E, 80-FE) in two, and in four the rest of the plane below U+10000 (less the + # surrogates) and every value above it. + "gb18030": { + 1: 0x80, + 2: 126 * 190, + 4: (0x10000 - 0x800 - 0x80 - 126 * 190) + (0x110000 - 0x10000), + }, + # CPython's KS X 1001 table (_codecs_kr.c): 2,350 Hangul, 4,888 Hanja and 987 others in two + # bytes (the filler, A4 D4, starts the eight-byte form instead; the euro and registered signs + # came in 1998), and in eight the 19 * 21 * 28 modern syllables the make-up forms compose. + "euc_kr": {1: 0x80, 2: 2350 + 4888 + 987, 8: 19 * 21 * 28}, +} + + +def _accepted_codecs() -> set[str]: + """Every codec Python accepts on a coding line, by its codec info's name.""" + names = set(encodings.aliases.aliases) | set(encodings.aliases.aliases.values()) + names |= {module.name for module in pkgutil.iter_modules(encodings.__path__)} + accepted: set[str] = set() + for name in names - {"aliases"}: + try: + codec = codecs.lookup(name).name + compile( + f"# coding: {name}\nx = 1\n".encode(), + "", + "exec", + dont_inherit=True, + ) + except (LookupError, SyntaxError): + continue + accepted.add(codec) + return accepted + + +def _read(codec: str, form: bytes) -> tuple[str, str | None]: + """Whether ``form`` is complete, incomplete or refused in ``codec``, and its text.""" + try: + return "complete", form.decode(codec) + except UnicodeDecodeError as error: + if error.reason in _INCOMPLETE and error.end == len(form): + return "incomplete", None + return "refused", None + + +class _Sweep: + """Every complete byte form of one codec, found depth-first from each first byte, and those + that could spell a name from other bytes.""" + + def __init__(self, codec: str) -> None: + self.codec = codec + self.forms: collections.Counter[int] = collections.Counter() + self.unsafe: list[str] = [] + + def walk(self, prefix: bytes = b"") -> None: + for byte in range(256): + form = prefix + bytes((byte,)) + if self._by_structure(form): + continue + state, text = _read(self.codec, form) + if state == "complete": + self._check(form, text) + elif state == "incomplete": + if len(form) == _DEEPEST_FORM: + self.unsafe.append(f"{form.hex()}: still incomplete") + else: + self.walk(form) + + def _by_structure(self, form: bytes) -> bool: + # These decoders ask for a whole long form before they read its bytes, or would cost a + # tree of every byte at each place, so their long forms are built from their ranges. + if ( + self.codec in ("utf-8", "utf-8-sig") + and len(form) == 1 + and 0xE0 <= form[0] <= 0xF4 + ): + self._utf8(form, 3 if form[0] < 0xF0 else 4) + elif self.codec == "gb18030" and len(form) == 2 and 0x30 <= form[1] <= 0x39: + for third in range(0x81, 0xFF): + for fourth in range(0x30, 0x3A): + self._leaf(form + bytes((third, fourth))) + elif self.codec == "euc_kr" and form == b"\xa4\xd4": + for initial in range(0xA1, 0xFF): + for medial in range(0xA1, 0xFF): + for final in range(0xA1, 0xFF): + self._leaf( + bytes( + (0xA4, 0xD4, 0xA4, initial, 0xA4, medial, 0xA4, final) + ) + ) + else: + return False + return True + + def _utf8(self, prefix: bytes, length: int) -> None: + # Only a continuation byte, 80-BF, can follow a lead. + for byte in range(0x80, 0xC0): + form = prefix + bytes((byte,)) + if len(form) == length: + self._leaf(form) + elif _read(self.codec, form)[0] == "incomplete": + self._utf8(form, length) + + def _leaf(self, form: bytes) -> None: + state, text = _read(self.codec, form) + if state == "complete": + self._check(form, text) + elif state == "incomplete": + self.unsafe.append(f"{form.hex()}: incomplete at its full length") + + def _check(self, form: bytes, text: str | None) -> None: + self.forms[len(form)] += 1 + if len(form) == 1 and text == chr(form[0]): + return + if not text: + # A byte-order mark is read only at the file's start, before any name. + if not (self.codec == "utf-8-sig" and form == codecs.BOM_UTF8): + self.unsafe.append(f"{form.hex()}: decodes to nothing") + elif not _NAME_OR_LINE.isdisjoint(text): + self.unsafe.append(f"{form.hex()}: decodes to {text!r}") + + +def _outside_the_structure() -> list[str]: + """Each form with a byte the structural ranges leave out that isn't refused: a decoder that + took one would have forms the sweep never tried.""" + taken: list[str] = [] + for lead in range(0x81, 0xFF): + for digit in range(0x30, 0x3A): + for third in (*range(0x00, 0x81), 0xFF): + for fourth in (0x30, 0x39): + form = bytes((lead, digit, third, fourth)) + if _read("gb18030", form)[0] != "refused": + taken.append(f"gb18030 {form.hex()}") + for fourth in (*range(0x00, 0x30), *range(0x3A, 0x100)): + form = bytes((lead, digit, 0x81, fourth)) + if _read("gb18030", form)[0] != "refused": + taken.append(f"gb18030 {form.hex()}") + # Each place of a make-up form after the filler: A4, or a jamo byte of A1-FE. + syllable = bytes((0xA4, 0xD4, 0xA4, 0xA1, 0xA4, 0xBF, 0xA4, 0xA1)) + assert _read("euc_kr", syllable)[0] == "complete" + for place in range(2, 8): + allowed = {0xA4} if place % 2 == 0 else set(range(0xA1, 0xFF)) + for byte in sorted(set(range(256)) - allowed): + form = syllable[:place] + bytes((byte,)) + syllable[place + 1 :] + if _read("euc_kr", form)[0] != "refused": + taken.append(f"euc_kr {form.hex()}") + return taken + + +def test_every_raw_safe_codec_is_proven(): + accepted = _accepted_codecs() + assert _RAW_SAFE_CODECS <= accepted + assert len(_RAW_SAFE_CODECS) == 81 + assert _RAW_SAFE_CODECS.isdisjoint(_SPELLING_CODECS) + unsafe: dict[str, list[str]] = {} + forms: dict[str, dict[int, int]] = {} + for codec in sorted(_RAW_SAFE_CODECS): + sweep = _Sweep(codec) + sweep.walk() + if sweep.unsafe: + unsafe[codec] = sweep.unsafe[:5] + if codec in _FORMS: + forms[codec] = dict(sweep.forms) + assert unsafe == {} + assert forms == _FORMS + assert _outside_the_structure() == [] + # The escape codecs refuse a lone `\` rather than call it incomplete, so the sweep can't see + # them; the disjointness above keeps them out. Every other spelling codec it must flag. + for codec in sorted(_SPELLING_CODECS - {"unicode-escape", "raw-unicode-escape"}): + sweep = _Sweep(codec) + sweep.walk() + assert sweep.unsafe, codec + + +def test_the_accepted_codecs_are_exactly_the_three_groups(): + # A codec a new Python adds fails here until it's put in one group or the other. + assert _accepted_codecs() == _RAW_SAFE_CODECS | _SPELLING_CODECS | {"idna"} + + +# Each holds a needle in bytes its codec decodes to other characters, beside an import of +# playwright Python runs: the bytes let the file in, and the text read is Python's. Sorted by +# their bytes; two need Python 3.12's reading of a decoded CR. +_DECOYS = [ + ( + b"# coding: hz\n" + b"# ~{playwright~}\n" + b"import os, \\\n" + b"~{#p#l#a#y#w#r#i#g#h#t~}\n" + b"import ~{playwright~}\n" + ), + ( + b"# coding: hz\n" + b"import os, \\\r\n" + b"~{#p#l#a#y#w#r#i#g#h#t~}\r\r\n" + b"import os\r\r\n" + b"x = 1 # ~{playwright~}\r\n" + ), + ( + b"# coding: hz\n" + b"try:\r import ~{#p#l#a#y#w#r#i#g#h#t~}\rexcept E:\r pass\rx = 1\rtry:\r import ~{playwright~}\rexcept E:\r pass" + ), + ( + b"# coding: hz\n" + b"x = f # c~{playwright~}(1); import ~{playwright~}\n" + b"import os\n" + b"def f():\n" + b" import ~{#p#l#a#y#w#r#i#g#h#t~}\n" + ), + ( + b"# coding: hz\r\n" + b"def f():\r\n" + b" import ~{#p#l#a#y#w#r#i#g#h#t~}\r\r\n" + b"x = f # c~{playwright~}(1); import ~{#p#l#a#y#w#r#i#g#h#t~}\r\n" + ), + ( + b"# coding: hz\r\n" + b"import os, \\\r\n" + b"~{#p#l#a#y#w#r#i#g#h#t~}\r\r\n" + b"x = f # c~{#p#l#a#y#w#r#i#g#h#t~}(1); import ~{playwright~}\r\n" + ), + ( + b"# coding: hz\r\n" + b"import ~{#p#l#a#y#w#r#i#g#h#t~}\r\r\n" + b"x = 1 # ~{playwright~}\r\r\n" + b"import ~{playwright~}\r\r\n" + b"import os, \\\r\n" + b"~{#p#l#a#y#w#r#i#g#h#t~}\r\n" + ), + ( + b"# coding: hz\r\n" + b"try:\n" + b" import ~{playwright~}\n" + b"except E:\n" + b" pass\n" + b"def g():\n" + b" ~{playwright~}\n" + b" import ~{#p#l#a#y#w#r#i#g#h#t~}\n" + ), + ( + b"# coding: hz\rdef f():\r\n" + b" import ~{#p#l#a#y#w#r#i#g#h#t~}\r\r\n" + b"~{playwright~}" + ), + b"# coding: hz\rdef f():\r import ~{#p#l#a#y#w#r#i#g#h#t~}\rimport ~{playwright~} # ~{playwright~}\rdef f():\r import ~{playwright~}\r", + ( + b"# coding: hz\rdef g():\n" + b" ~{playwright~}\n" + b" import ~{#p#l#a#y#w#r#i#g#h#t~}\n" + ), + ( + b"# coding: hz\rimport os # ~{#p#l#a#y#w#r#i#g#h#t~}, ~{playwright~}\n" + b"# ~{#p#l#a#y#w#r#i#g#h#t~}\n" + b"try:\n" + b" import ~{#p#l#a#y#w#r#i#g#h#t~}\n" + b"except E:\n" + b" pass\n" + b"x = 1 # ~{#p#l#a#y#w#r#i#g#h#t~}" + ), + b"# coding: hz\rimport ~{#p#l#a#y#w#r#i#g#h#t~} # ~{playwright~}", + ( + b"# coding: hz\rimport ~{playwright~} # ~\n" + b"\n" + b"import ~{#p#l#a#y#w#r#i#g#h#t~} # ~{playwright~}\n" + b"~\n" + b"\n" + ), + ( + b"# coding: hz\rtry:\r\n" + b" import ~{#p#l#a#y#w#r#i#g#h#t~}\r\n" + b"except E:\r\n" + b" pass\r\r\n" + b" \\\r\n" + b"~{playwright~}1\r\n" + ), + b"# coding: hz\rx = '~{playwright~}'\rimport ~{#p#l#a#y#w#r#i#g#h#t~}", + ( + b"# coding: hz\rx = 1\n" + b"def f():\n" + b" import ~{playwright~}\n" + b"import ~{#p#l#a#y#w#r#i#g#h#t~}\n" + ), + b"# coding: hz\rx = 1\rx = 1\rimport os\rdef f():\r import ~{playwright~}\rtry:\r import ~{#p#l#a#y#w#r#i#g#h#t~}\rexcept E:\r pass\r", + ( + b"# coding: hz\rx = 1 # ~~\r\r\n" + b"import os\r\r\n" + b"~{playwright~}\r\r\n" + b"import ~{#p#l#a#y#w#r#i#g#h#t~} # ~{#p#l#a#y#w#r#i#g#h#t~}\r\n" + ), + ( + b"# coding: iso2022_kr\r\n" + b"x = 1\n" + b"\x1b$)C\n" + b"try:\n" + b" import \x0e#p#l#a#y#w#r#i#g#h#t\x0f\n" + b"except E:\n" + b" pass\n" + b"import \x0eplaywright\x0f" + ), + ( + b"# coding: utf-7\n" + b"def g():\r\n" + b" +AAo-\r\n" + b" import +/1A-laywright\r\r\n" + b"try:\r\n" + b" import +playwrightA-\r\n" + b"except E:\r\n" + b" pass\r\n" + ), + (b"# coding: utf-7\ndef g():\r +playwrightA-\r import +/1A-laywright\r"), + ( + b"# coding: utf-7\n" + b"if a:\n" + b" pass\n" + b"import +/1A-laywright\n" + b"x = '+-'\n" + b"import +playwrightA- # +/1A-laywright\n" + ), + (b"# coding: utf-7\nimport os, \\\r+/1A-laywright\r+-+playwrightA-"), + ( + b"# coding: utf-7\n" + b"x = 1 # +-\rtry:\r import +/1A-laywright\rexcept E:\r pass\rx = '+-'\rtry:\r import +playwrightA-\rexcept E:\r pass" + ), + (b"# coding: utf-7\r\nimport +/1A-laywright\nx = 1 # +playwrightA-"), + (b"# coding: utf-7\r\nimport +/1A-laywright # +playwrightA-"), + ( + b"# coding: utf-7\r\n" + b"import os # +-, +playwrightA-\r\r\n" + b"import os, \\\r\n" + b"+/1A-laywright\r\n" + ), + (b"# coding: utf-7\r\nimport os, \\\r+/1A-laywright\rx = '+playwrightA-'\r"), + ( + b"# coding: utf-7\r\n" + b"x = 1\rif a:\r pass\rimport +/1A-laywright # +/1A-laywright\rx = '+playwrightA-'\r" + ), + b"# coding: utf-7\r# +playwrightA-\rtry:\r import +/1A-laywright\rexcept E:\r pass+AHA", + (b"# coding: utf-7\r# -\nx = 1\nimport +/1A-laywright # +playwrightA-\n"), + b"# coding: utf-7\r+playwrightA-x = 1\rtry:\r import +/1A-laywright\rexcept E:\r pass\rx = 1 # +playwrightA-\r", + ( + b"# coding: utf-7\rif a:\n" + b" pass\n" + b"import +/1A-laywright\n" + b"'''+AHA'''\n" + b"import os, \\\n" + b"+playwrightA-+playwrightA-" + ), + b"# coding: utf-7\rimport +/1A-laywright\rtry:\r import +playwrightA-\rexcept E:\r pass\r", + ( + b"# coding: utf-7\rtry:\r\n" + b" import +/1A-laywright\r\n" + b"except E:\r\n" + b" pass\r\r\n" + b"x = 1 # +playwrightA-" + ), + ( + b"# coding: utf-7\rx = f # c+AA0ACg-(1); import +playwrightA-\r\r\n" + b"try:\r\n" + b" import +/1A-laywright\r\n" + b"except E:\r\n" + b" pass\r\r\n" + b"x = 1\r\n" + ), + ( + b"# coding=hz\n" + b"def f():\r\n" + b" import ~{#p#l#a#y#w#r#i#g#h#t~}\r\r\n" + b"import os\r\r\n" + b"import os, \\\r\n" + b"~{playwright~}\r\r\n" + b"x = 1" + ), + ( + b"# coding=hz\n" + b"def g():\r ~{playwright~}\r import ~{#p#l#a#y#w#r#i#g#h#t~}\r" + ), + ( + b"# coding=hz\n" + b"import os, \\\r\n" + b"~{#p#l#a#y#w#r#i#g#h#t~}\r\r\n" + b"'''~~'''\r\r\n" + b"def f():\r\n" + b" import ~{playwright~}\r\r\n" + b"x = 1" + ), + ( + b"# coding=hz\n" + b"import os, \\\r~{playwright~}\rimport ~{playwright~}\rimport ~{#p#l#a#y#w#r#i#g#h#t~}\r" + ), + ( + b"# coding=hz\n" + b"import ~{#p#l#a#y#w#r#i#g#h#t~}\rtry:\r import ~{playwright~}\rexcept E:\r pass\r" + ), + ( + b"# coding=hz\n" + b"import ~{playwright~}\n" + b"import os, \\\n" + b"~{#p#l#a#y#w#r#i#g#h#t~}\n" + b"if a:\n" + b" pass\n" + b"if a:\n" + b" pass" + ), + ( + b"# coding=hz\n" + b"import ~{playwright~}\r\r\n" + b"import os, \\\r\n" + b"~{#p#l#a#y#w#r#i#g#h#t~}\r\n" + ), + ( + b"# coding=hz\n" + b"x = 1\n" + b"~{playwright~}\n" + b"try:\n" + b" import ~{#p#l#a#y#w#r#i#g#h#t~}\n" + b"except E:\n" + b" pass\n" + ), + pytest.param( + ( + b"# coding=utf-7\n" + b"+playwrightA-\r'''+-'''\rimport os, \\\r+/1A-laywright+AA0ACg-" + ), + marks=_FROM_312, + ), + pytest.param( + ( + b"# coding=utf-7\n" + b"def f():\n" + b" import +playwrightA-\n" + b"+AA0-\n" + b"def g():\n" + b" +AA0ACg-\n" + b" import +/1A-laywright\n" + ), + marks=_FROM_312, + ), + ( + b"# coding=utf-7\n" + b"if a:\n" + b" x = 1\n" + b"+AAo-\n" + b" import +/1A-laywright\n" + b"def f():\n" + b" import +playwrightA-\n" + ), + ( + b"# coding=utf-7\n" + b"import +playwrightA- # +playwrightA-\r\r\n" + b"import os\r\r\n" + b"try:\r\n" + b" import +/1A-laywright\r\n" + b"except E:\r\n" + b" pass\r\r\n" + b"'''+AHA'''\r\n" + ), + ( + b"# coding=utf-7\n" + b"import os, \\\n" + b"+/1A-laywright\n" + b"+/1A-laywrightx = 1\n" + b"# +AAo-\n" + b"def f():\n" + b" import +/1A-laywright+playwrightA-" + ), + ( + b"# coding=utf-7\n" + b"x = 1 # +-\r\r\n" + b"+/1A-laywright\r\r\n" + b"'''+playwrightA-'''\r\r\n" + b"try:\r\n" + b" import +/1A-laywright\r\n" + b"except E:\r\n" + b" pass" + ), + ( + b"#!/usr/bin/env python\n" + b"# -*- coding: hz -*-\n" + b"# ~{playwright~}\r\r\n" + b"try:\r\n" + b" import ~{#p#l#a#y#w#r#i#g#h#t~}\r\n" + b"except E:\r\n" + b" pass\r\r\n" + b"~{#p#l#a#y#w#r#i#g#h#t~}" + ), + ( + b"#!/usr/bin/env python\n" + b"# -*- coding: hz -*-\n" + b"# ~{playwright~}\rimport ~{#p#l#a#y#w#r#i#g#h#t~}\r~{#p#l#a#y#w#r#i#g#h#t~}" + ), + ( + b"#!/usr/bin/env python\n" + b"# -*- coding: hz -*-\n" + b"if a:\r pass\rimport ~{#p#l#a#y#w#r#i#g#h#t~}\rimport ~{#p#l#a#y#w#r#i#g#h#t~}~{playwright~}" + ), + ( + b"#!/usr/bin/env python\n" + b"# -*- coding: hz -*-\n" + b"import os, \\\n" + b"~{#p#l#a#y#w#r#i#g#h#t~}\n" + b"import os, \\\n" + b"~{playwright~}\n" + ), + ( + b"#!/usr/bin/env python\n" + b"# -*- coding: hz -*-\n" + b"import ~{#p#l#a#y#w#r#i#g#h#t~}\n" + b"'''~{playwright~}'''\n" + b"x = 1\n" + ), + ( + b"#!/usr/bin/env python\n" + b"# -*- coding: hz -*-\n" + b"import ~{#p#l#a#y#w#r#i#g#h#t~}\n" + b"def f():\n" + b" import ~{playwright~}\n" + ), + ( + b"#!/usr/bin/env python\n" + b"# -*- coding: hz -*-\n" + b"import ~{#p#l#a#y#w#r#i#g#h#t~}\n" + b"x = 1~\n" + b"\n" + b"import os, \\\n" + b"~{#p#l#a#y#w#r#i#g#h#t~}\n" + b"x = '~\n" + b"'\n" + b"import os, \\\n" + b"~{playwright~}\n" + ), + ( + b"#!/usr/bin/env python\n" + b"# -*- coding: hz -*-\n" + b"import ~{playwright~}\r\r\n" + b"def f():\r\n" + b" import ~{playwright~}\r\r\n" + b"def f():\r\n" + b" import ~{#p#l#a#y#w#r#i#g#h#t~}" + ), + ( + b"#!/usr/bin/env python\n" + b"# -*- coding: hz -*-\n" + b"try:\n" + b" import ~{playwright~}\n" + b"except E:\n" + b" pass\n" + b"# ~~\n" + b"import os, \\\n" + b"~{#p#l#a#y#w#r#i#g#h#t~}" + ), + ( + b"#!/usr/bin/env python\n" + b"# -*- coding: hz -*-\n" + b"x = f # c~{playwright~}(1); import ~{#p#l#a#y#w#r#i#g#h#t~}\r\r\n" + b"def f():\r\n" + b" import ~{#p#l#a#y#w#r#i#g#h#t~}\r\r\n" + b"x = 1" + ), + ( + b"#!/usr/bin/env python\n" + b"# -*- coding: iso2022_kr -*-\n" + b"x = 1\x1b$)C\rtry:\r import \x0eplaywright\x0f\rexcept E:\r pass\rtry:\r import \x0eplaywright\x0f\rexcept E:\r pass\rimport \x0e#p#l#a#y#w#r#i#g#h#t\x0f # \x1b\r" + ), + ( + b"#!/usr/bin/env python\n" + b"# -*- coding: utf-7 -*-\n" + b"+/1A-laywright\r\r\n" + b"def g():\r\n" + b" +playwrightA-\r\n" + b" import +/1A-laywright" + ), + ( + b"#!/usr/bin/env python\n" + b"# -*- coding: utf-7 -*-\n" + b"def f():\r\n" + b" import +playwrightA-\r\r\n" + b"import +/1A-laywright # +/1A-laywright\r\n" + ), + ( + b"#!/usr/bin/env python\n" + b"# -*- coding: utf-7 -*-\n" + b"if a:\r\n" + b" pass\r\r\n" + b"def g():\r\n" + b" +playwrightA-\r\n" + b" import +/1A-laywright\r\r\n" + b"x = 1+AAo-" + ), + ( + b"#!/usr/bin/env python\n" + b"# -*- coding: utf-7 -*-\n" + b"import +playwrightA- # +/1A-laywright\n" + b"import +/1A-laywright # +-" + ), +] + + +@pytest.mark.parametrize( + "raw", _DECOYS, ids=[f"decoy-{n:02}" for n in range(len(_DECOYS))] +) +def test_no_decoy_case_regresses(tmp_path, raw): + assert "playwright" in _imported(ast.parse(raw)) + path = _write(tmp_path, raw) + text = read_module( + path, needles=agent_type_check._PYTHON_NAMES, escapes=EscapeBudget() + ) + assert text is not None + assert "playwright" in _imported(ast.parse(text)) From f103efd3f28a66256c7e5b5000e21f14e08d83b7 Mon Sep 17 00:00:00 2001 From: khushalsonawat Date: Wed, 11 Nov 2026 05:05:04 +0530 Subject: [PATCH 2/3] fix(harness): the save's import walk says when it leaves code unread, and reads more of it (TH-9543) - The walk sets stopped when it leaves code it reached unfollowed, beyond its limits: - a .py holding import whose decoded reading isn't parsed (too big, past the total, or no tree); - a code file over the read limit, or a link to one; - an import name dropped for its length or depth, scripts included; - a link target past 128 parts; - a file in a codec decoded first, left undecoded once the walk's decode total ran out. - The save walk gets its own 512 KB parse total (the other scans keep 256 KB), and each reading of a file read two ways is charged as it's parsed. - One flag replaces the walk's two. Signed-off-by: khushalsonawat --- src/fi/alk/harness/agent_type_check.py | 6 +- src/fi/alk/harness/credentials.py | 130 +++++--- tests/harness/test_agent_type_refusal.py | 368 ++++++++++++++++++++- tests/harness/test_python_source_codecs.py | 24 +- tests/test_harness_credentials.py | 13 + 5 files changed, 467 insertions(+), 74 deletions(-) diff --git a/src/fi/alk/harness/agent_type_check.py b/src/fi/alk/harness/agent_type_check.py index 1679bb511..b7ed80fe6 100644 --- a/src/fi/alk/harness/agent_type_check.py +++ b/src/fi/alk/harness/agent_type_check.py @@ -157,6 +157,10 @@ _WALKED_SCRIPTS = 2 * 1024 * 1024 # The import names it tries then. _WALKED_IMPORTS = 5_000 +# And the Python it parses then, each reading of a file charged: twice the other scans' total, so an +# agent's own code of a few hundred KB isn't cut, while crafted files that only this walk parses +# hold the save about 1.3 s on 3.12 and 3.1 s on 3.11. +_WALKED_PARSED = 512 * 1024 # And the lookups on disk it makes then, a file found costing one for each part of its path: a name # can resolve a package for each of its parts, and 1,000 crafted names took the save 12 s. _WALKED_LOOKUPS = 100_000 @@ -642,7 +646,7 @@ def _walked(root: Path, subfolder: str) -> frozenset[Path]: root, subfolder, parse_limit=LARGEST_PARSED_MODULE, - parse_budget=ParseBudget(), + parse_budget=ParseBudget(_WALKED_PARSED), script_budget=ParseBudget(_WALKED_SCRIPTS), import_limit=_WALKED_IMPORTS, lookup_limit=_WALKED_LOOKUPS, diff --git a/src/fi/alk/harness/credentials.py b/src/fi/alk/harness/credentials.py index 4284e5174..2585be9d1 100644 --- a/src/fi/alk/harness/credentials.py +++ b/src/fi/alk/harness/credentials.py @@ -672,23 +672,38 @@ def record( def _file_inside( - candidate: Path, root: Path, deepest: int, left: list[float] | None = None + candidate: Path, + root: Path, + deepest: int, + left: list[float] | None = None, + unread: list[bool] | None = None, ) -> Path | None: """``candidate`` resolved, if it's a plain file in ``root`` at most ``deepest`` parts long. ``left`` is the lookups on disk still allowed: one for the file, then one for each part of - its path, which resolving it looks up in turn. + its path, which resolving it looks up in turn. Given ``unread``, its one item is set when a + candidate is refused for its depth, or as a code file too big to read or a link to one: code + may be there that the walk doesn't follow. """ if len(candidate.parts) > deepest: + _mark(unread) return None if left is not None: left[0] -= 1 - if not _plain_file(candidate): + if not _plain_file(candidate, unread): return None if left is not None: left[0] -= len(candidate.parts) resolved = candidate.resolve() - return resolved if len(resolved.parts) <= deepest and is_within(resolved, root) else None + if len(resolved.parts) > deepest: + _mark(unread) + return None + return resolved if is_within(resolved, root) else None + + +def _mark(unread: list[bool] | None) -> None: + if unread is not None: + unread[0] = True # A uv workspace's members read for their names, at most. @@ -1071,7 +1086,11 @@ def _top_names(base: Path, left: list[float], limit: int) -> tuple[set[str], boo def _real_folder( - folder: Path, root: Path, known: dict[tuple[str, ...], Path | None], left: list[float] + folder: Path, + root: Path, + known: dict[tuple[str, ...], Path | None], + left: list[float], + unread: list[bool] | None = None, ) -> Path | None: """``folder`` with each link in it followed, if it's a folder inside ``root``; else None. @@ -1082,11 +1101,13 @@ def _real_folder( which has no link in it, so it costs nothing; above ``root`` nothing is looked up, and a path that stays out gives None. So do a missing folder, a loop and a chain of links past ``_LINKS_FOLLOWED``; what that limit cut isn't kept, since the same link reached by a shorter - chain may be followed. + chain may be followed. A link whose target is past ``_LONGEST_TARGET`` parts isn't followed + either, though the kernel would follow it: given ``unread``, its one item says so. """ # Each name here is reached first, so what the chain limit cuts below it is cut for any way to # it, and can be kept; a link's target is walked by _step, which keeps only what it didn't cut. - capped = [False] + # The second item: a link target refused for its length. + capped = [False, False] parts, top = folder.parts, len(root.parts) # A path longer than a link target may be (a uv ``path`` source is written by the customer) # isn't walked: each of its prefixes is a key, so its length would cost its square. @@ -1104,7 +1125,7 @@ def _real_folder( # Kept for the whole path too: each candidate in this folder (a name has 13, in two # folders) would otherwise search back down to the missing one again, uncharged. known[parts] = None - return None + break name = parts[end - 1] if name == "..": real = real.parent @@ -1114,6 +1135,8 @@ def _real_folder( else: real = _step(real / name, root, known, left, 0, capped) known[parts[:end]] = real + if capped[1]: + _mark(unread) return real if real is not None and is_within(real, root) else None @@ -1173,6 +1196,7 @@ def _followed( except (OSError, ValueError): return None if len(target.parts) > _LONGEST_TARGET: + capped[1] = True return None # Each part is a step on the way, charged as a lookup is. left[0] -= len(target.parts) @@ -1206,24 +1230,34 @@ def _code_in(folder: Path) -> Iterator[Path]: yield Path(directory, name) -def _plain_file(path: Path) -> bool: - """A regular file small enough to read: a committed link can point at anything.""" +def _plain_file(path: Path, unread: list[bool] | None = None) -> bool: + """A regular file small enough to read: a committed link can point at anything, so it isn't + followed. Given ``unread``, its one item is set when a code file is refused for its size or + for being a link, since Python and Node run what one holds or points at.""" try: details = path.lstat() except (OSError, ValueError): return False - return stat.S_ISREG(details.st_mode) and details.st_size <= _MAX_FILE_BYTES + if stat.S_ISREG(details.st_mode) and details.st_size <= _MAX_FILE_BYTES: + return True + if (stat.S_ISREG(details.st_mode) or stat.S_ISLNK(details.st_mode)) and ( + path.suffix == ".py" or path.suffix in _SCRIPT_SUFFIXES + ): + _mark(unread) + return False def _python_imports( content: str, limit: int = _MAX_FILE_BYTES, budget: ParseBudget | None = None -) -> list[tuple[int, str]]: +) -> list[tuple[int, str]] | None: + """The names ``content`` imports, each with its level; None when it isn't parsed: over + ``limit``, past ``budget``, or not Python the parser reads.""" # In bytes, as the file's size is: the parser's cost follows them, not the characters. if utf8_size(content) > limit or (budget is not None and not budget.spend(content)): - return [] + return None tree = parsed(content) if tree is None: - return [] + return None names: list[tuple[int, str]] = [] for node in ast.walk(tree): if isinstance(node, ast.Import): @@ -1316,7 +1350,9 @@ def imported_by_agent( lookups on disk (no limit when not given): a file found costs one for each part of its path, and one name can resolve a package for each of its parts. Nothing more than ``_DEEPEST_FOLLOWED`` folders below the repository is followed. Given ``stopped``, its one - item says whether a limit stopped the walk with files or names left to follow. + item says whether the walk left code it reached unfollowed: a limit stopped it with files or + names left, a name was dropped for its length or depth, or a code file was too big to read, + a link, or not parsed in the reading Python runs. """ root = Path(root).resolve() try: @@ -1333,7 +1369,7 @@ def imported_by_agent( # dozens of parts deep: 1,000 names took the save 12 s. What is looked up again costs nothing. lookups_left = [math.inf if lookup_limit is None else float(lookup_limit)] # The uv workspace members the agent's project depends on are installed with it, so an import - # names one by its package (Dileep's third review): each one's folder and its src/ are looked + # names one by its package: each one's folder and its src/ are looked # in too. project, members, members_cut = _workspace_dependencies(root, folder, lookups_left, {}) # The agent's own project, when the folder is below it, is installed too, src/ layout or not. @@ -1354,7 +1390,9 @@ def imported_by_agent( # follow every link in a path uncharged, and a chain of them took 2,300 imports 11.6 s. So the # save reaches what bundling, which resolves as the kernel does, reaches. real_folders: dict[tuple[str, ...], Path | None] | None = None if lookup_limit is None else {} - spent = [False] + # Whether the walk left code it reached unfollowed: a limit, a name dropped for its length or + # depth, or a file it couldn't read or parse. + unread = [members_cut] def found_at(candidates: Iterable[Path]) -> Path | None: for candidate in candidates: @@ -1362,7 +1400,7 @@ def found_at(candidates: Iterable[Path]) -> Path | None: hit = looked_up[candidate] elif lookups_left[0] <= 0: # A name left untried: the walk didn't reach all it could. - spent[0] = True + unread[0] = True return None else: place = candidate @@ -1372,13 +1410,15 @@ def found_at(candidates: Iterable[Path]) -> Path | None: and len(candidate.parts) <= deepest ): # Through its folders' links, each followed once and charged. - folder_found = _real_folder(candidate.parent, root, real_folders, lookups_left) + folder_found = _real_folder( + candidate.parent, root, real_folders, lookups_left, unread + ) place = None if folder_found is None else folder_found / candidate.name if place is None: lookups_left[0] -= 1 hit = None else: - hit = _file_inside(place, root, deepest, lookups_left) + hit = _file_inside(place, root, deepest, lookups_left, unread) if hit is not None or candidate.name == "__init__.py": looked_up[candidate] = hit if hit is not None: @@ -1389,6 +1429,7 @@ def python_modules(path: Path, level: int, name: str) -> list[Path]: module = name.split(".") # A path past the system's limit can't be opened, so it isn't built. if len(module) > _DEEPEST_FOLLOWED or len(name) > _LONGEST_NAME: + unread[0] = True return [] starts = ( (path.parents[level - 1],) @@ -1398,10 +1439,11 @@ def python_modules(path: Path, level: int, name: str) -> list[Path]: for start in starts: if lookups_left[0] <= 0: # Spent: each place left would be refused uncharged. - spent[0] = True + unread[0] = True return [] if len(start.parts) + len(module) > deepest: # Past the depth cut wherever it's looked for: not built at all. + unread[0] = True continue target = start.joinpath(*module) if hit := found_at((target.with_suffix(".py"), target / "__init__.py")): @@ -1417,6 +1459,7 @@ def python_modules(path: Path, level: int, name: str) -> list[Path]: def script_module(path: Path, name: str) -> Path | None: if len(name) > _LONGEST_NAME: + unread[0] = True return None for start in (path.parent,) if name.startswith(".") else (folder, root): # Joined as written, as Node joins it before it reads the disk: a ".." after a folder @@ -1442,7 +1485,6 @@ def script_module(path: Path, name: str) -> Path | None: pending = deque(sorted(_code_in(folder))) own = set(pending) # Read once already, so importing one doesn't queue it again. read = bytes_read = imports_tried = 0 - cut = members_cut import_limit = _MAX_WALK_IMPORTS if import_limit is None else import_limit while ( pending @@ -1453,7 +1495,10 @@ def script_module(path: Path, name: str) -> Path | None: ): # The folder's own files first, so the budgets can't spend themselves on a chain. path = pending.popleft() - if len(path.parts) > deepest or not _plain_file(path): + if len(path.parts) > deepest: + unread[0] = True + continue + if not _plain_file(path, unread): continue read += 1 try: @@ -1462,6 +1507,7 @@ def script_module(path: Path, name: str) -> Path | None: except OSError: continue if len(raw) > _MAX_FILE_BYTES: + unread[0] = True continue # Grew since it was looked at. bytes_read += len(raw) # A .py in a codec that can spell a name in other bytes is read as Python reads it while @@ -1499,30 +1545,20 @@ def script_module(path: Path, name: str) -> Path | None: hits: list[Path] = [] if path.suffix == ".py": # A file with no import in it can't lead anywhere, so it doesn't use up the parse total - # before the files that do (Pine's TH-9490 review); a file the total cuts says so. + # before the files that do; a file the total cuts says so. if "import" not in content: continue readings = [content] if read_before is None else [content, read_before] - larger = max(readings, key=utf8_size) - if len(readings) == 1 or utf8_size(larger) > parse_limit: - # Each reading is charged as it's parsed; one past the limit is neither. - found = [ - name - for reading in readings - for name in _python_imports(reading, parse_limit, parse_budget) - ] - elif parse_budget is None or parse_budget.spend(larger): - # One file read two ways is charged once, as the larger reading, so it spends no - # more of the total than it did read one way. - found = [ - name for reading in readings for name in _python_imports(reading, parse_limit) - ] - else: - found = [] - names = list(dict.fromkeys(found)) - cut |= parse_budget is not None and parse_budget.left < 0 + # Each reading is charged as it's parsed: a file parsed twice and charged once let a + # crafted pair cost the walk twice its total. + results = [_python_imports(reading, parse_limit, parse_budget) for reading in readings] + # The decoded reading is the one Python runs: refused, the walk left code unfollowed, + # whatever the other reading gives. A refusal of the other alone says nothing. + unread[0] |= results[0] is None + unread[0] |= parse_budget is not None and parse_budget.left < 0 + names = list(dict.fromkeys(name for found in results if found for name in found)) # One file can hold more names than the walk may try; the rest aren't tried. - cut |= len(names) > import_limit - imports_tried + unread[0] |= len(names) > import_limit - imports_tried names = names[: import_limit - imports_tried] imports_tried += len(names) for level, name in names: @@ -1531,11 +1567,11 @@ def script_module(path: Path, name: str) -> Path | None: except (IndexError, ValueError, OSError): continue # An import that can't be a path in this repository. elif script_budget is not None and not script_budget.spend(content): - cut = True + unread[0] = True continue else: names = list(dict.fromkeys(_script_imports(content))) - cut |= len(names) > import_limit - imports_tried + unread[0] |= len(names) > import_limit - imports_tried names = names[: import_limit - imports_tried] imports_tried += len(names) for name in names: @@ -1548,9 +1584,11 @@ def script_module(path: Path, name: str) -> Path | None: seen.add(hit) if hit not in own: pending.append(hit) + # A file in a codec decoded first, read without decoding once the total ran out. + unread[0] |= escapes.left < 0 if stopped is not None: # The loop ends with files left only when a limit ended it. - stopped[:] = [cut or spent[0] or bool(pending)] + stopped[:] = [unread[0] or bool(pending)] return frozenset(seen) diff --git a/tests/harness/test_agent_type_refusal.py b/tests/harness/test_agent_type_refusal.py index dcb24c5a8..39d7495fb 100644 --- a/tests/harness/test_agent_type_refusal.py +++ b/tests/harness/test_agent_type_refusal.py @@ -3127,7 +3127,7 @@ def test_the_walk_while_the_conversation_waits_follows_only_small_files(tmp_path def test_the_walk_while_the_conversation_waits_parses_only_up_to_its_total( tmp_path, monkeypatch, total, expected ): - from fi.alk.harness import python_source + from fi.alk.harness import agent_type_check repo = _tree( tmp_path / "repo", @@ -3137,7 +3137,8 @@ def test_the_walk_while_the_conversation_waits_parses_only_up_to_its_total( "common/web.py": "import playwright\n", }, ) - monkeypatch.setattr(python_source, "LARGEST_PARSED_TOTAL", total) + # The walk's own total: the browser scan's stays as it is. + monkeypatch.setattr(agent_type_check, "_WALKED_PARSED", total) assert browser_tools_in_agent(repo, "agents/a") == expected @@ -3548,9 +3549,9 @@ def test_the_walk_looks_up_each_package_once(tmp_path, monkeypatch): looked: list[Path] = [] look_up = credentials._file_inside - def counted(candidate, *args): + def counted(candidate, *args, **named): looked.append(candidate) - return look_up(candidate, *args) + return look_up(candidate, *args, **named) monkeypatch.setattr(credentials, "_file_inside", counted) @@ -3566,9 +3567,9 @@ def _counted_lookups(monkeypatch) -> list[Path]: looked: list[Path] = [] look_up = credentials._file_inside - def counted(candidate, *args): + def counted(candidate, *args, **named): looked.append(candidate) - return look_up(candidate, *args) + return look_up(candidate, *args, **named) monkeypatch.setattr(credentials, "_file_inside", counted) return looked @@ -4025,9 +4026,9 @@ def test_the_walk_charges_a_file_found_for_each_part_of_its_path(tmp_path, monke charged: list[tuple[Path, float]] = [] real = credentials._file_inside - def spy(candidate, root, deepest, left=None): + def spy(candidate, root, deepest, left=None, unread=None): before = left[0] - hit = real(candidate, root, deepest, left) + hit = real(candidate, root, deepest, left, unread) charged.append((candidate, before - left[0])) return hit @@ -4318,7 +4319,7 @@ def test_a_member_is_looked_in_only_for_its_own_package_names(tmp_path, monkeypa repo = _tree(tmp_path / "repo", files) tried: list[Path] = [] real = credentials._file_inside - monkeypatch.setattr(credentials, "_file_inside", lambda place, *rest: tried.append(place) or real(place, *rest)) + monkeypatch.setattr(credentials, "_file_inside", lambda place, *rest, **named: tried.append(place) or real(place, *rest, **named)) found = credentials.imported_by_agent(repo, "agents/a", lookup_limit=100_000) assert (repo / "libs" / "m7" / "m7" / "core.py").resolve() in found @@ -4775,7 +4776,7 @@ def test_a_candidate_past_the_depth_cut_isnt_walked_folder_by_folder(tmp_path, m monkeypatch.setattr( credentials, "_real_folder", - lambda folder, *rest: walked.append(len(folder.parts)) or real(folder, *rest), + lambda folder, *rest, **named: walked.append(len(folder.parts)) or real(folder, *rest, **named), ) credentials.imported_by_agent(repo, "agents/a", lookup_limit=100_000) @@ -4845,7 +4846,7 @@ def test_each_folder_on_the_way_to_a_candidate_is_walked_once(tmp_path, monkeypa (repo / "link").symlink_to("p") stepped: list[Path] = [] real = credentials._step - monkeypatch.setattr(credentials, "_step", lambda name, *rest: stepped.append(name) or real(name, *rest)) + monkeypatch.setattr(credentials, "_step", lambda name, *rest, **named: stepped.append(name) or real(name, *rest, **named)) credentials.imported_by_agent(repo, "agents/a", lookup_limit=100_000) # Each folder from the root down is stepped through once, whatever the number of names in it. @@ -5547,15 +5548,13 @@ def test_what_follows_a_cookie_isnt_counted_against_it(): def test_plain_files_ahead_of_the_agents_imports_dont_use_up_the_parse_total(tmp_path): from fi.alk.harness import agent_type_check - # Pine's TH-9490 review: three plain 126 KB files that sort first spent the 256 KB total, and - # the walk lost common/web.py with no word. + # Plain 126 KB files that sort first spent the parse total, and the walk lost common/web.py + # with no word: five of them are more than the save walk's 512 KB. pad = "x = 1\n" * 21_000 repo = _tree( tmp_path / "repo", { - "agents/a/a0.py": pad, - "agents/a/a1.py": pad, - "agents/a/a2.py": pad, + **{f"agents/a/a{n}.py": pad for n in range(5)}, "agents/a/main.py": "from common import web\n", "common/__init__.py": "", "common/web.py": "", @@ -5744,3 +5743,340 @@ def test_a_path_source_to_a_folder_with_no_project_isnt_followed(tmp_path): _, followed, _ = credentials._workspace_dependencies(repo, repo / "agents" / "a", [1e9], {}) assert followed == [] + + +def _save_walk(repo: Path, subfolder: str = "agents/a"): + """The walk with the save's limits, as _walked runs it: the files reached, and stopped.""" + from fi.alk.harness import agent_type_check + from fi.alk.harness.credentials import imported_by_agent + from fi.alk.harness.python_source import LARGEST_PARSED_MODULE, ParseBudget + + stopped: list[bool] = [] + found = imported_by_agent( + repo, + subfolder, + parse_limit=LARGEST_PARSED_MODULE, + parse_budget=ParseBudget(agent_type_check._WALKED_PARSED), + script_budget=ParseBudget(agent_type_check._WALKED_SCRIPTS), + import_limit=agent_type_check._WALKED_IMPORTS, + lookup_limit=agent_type_check._WALKED_LOOKUPS, + stopped=stopped, + ) + return found, stopped[0] + + +def _behind(tmp_path: Path, helper: bytes) -> Path: + """agents/a/main.py imports common.helpers (``helper``), which may import common.web.""" + repo = _tree( + tmp_path / "repo", + {"agents/a/main.py": "import common.helpers\n", "common/__init__.py": "", "common/web.py": "import playwright\n"}, + ).resolve() + (repo / "common" / "helpers.py").write_bytes(helper) + return repo + + +@pytest.mark.parametrize( + "helper", + [ + pytest.param( + b"from common import web\ndef first[T = int](x: T) -> T:\n return x\n", + id="3.13-syntax", + marks=pytest.mark.skipif(sys.version_info >= (3, 13), reason="parses on 3.13"), + ), + pytest.param(b"from common import web\nNAME = 'a" + "\u0301".encode() * 40 + b"'\n", id="marks"), + pytest.param( + b"# coding: latin-1\nfrom common import web\n# " + b"\xe9" * 70_000 + b"\n", id="oversize-as-utf8" + ), + ], +) +def test_a_file_the_walk_cant_parse_says_it_stopped(tmp_path, helper): + repo = _behind(tmp_path, helper) + + found, stopped = _save_walk(repo) + assert stopped and repo / "common" / "web.py" not in found + + +def test_a_file_past_the_parse_total_says_the_walk_stopped(tmp_path, monkeypatch): + from fi.alk.harness import agent_type_check + + repo = _behind(tmp_path, b"from common import web\n" + b"x = 1\n" * 50) + monkeypatch.setattr(agent_type_check, "_WALKED_PARSED", 40) + + found, stopped = _save_walk(repo) + assert stopped and repo / "common" / "web.py" not in found + + +_ESCAPED_LINE = b"# coding: unicode-escape\n" + + +@pytest.mark.parametrize( + "helper, stopped_too, reached", + [ + # Decoded, the CR ends the comment and Python runs the import; read as UTF-8 the comment + # runs on and the file parses, with only os. + (_ESCAPED_LINE + b"import os # c\\r; import common.web\n", True, False), + # Decoded, it's two imports; read as UTF-8, a stray backslash: only the decoded parses. + (_ESCAPED_LINE + b"import os\\rimport common.web\n", False, True), + # Both readings parse. + (_ESCAPED_LINE + b"import common.web # see \\r\n", False, True), + ], + ids=["decoded-refused", "utf8-refused", "both-parse"], +) +def test_a_file_read_two_ways_is_judged_by_its_decoded_reading(tmp_path, helper, stopped_too, reached): + repo = _behind(tmp_path, helper) + + found, stopped = _save_walk(repo) + assert stopped is stopped_too + assert (repo / "common" / "web.py" in found) is reached + + +def test_files_read_two_ways_base_parsed_are_still_parsed(tmp_path): + # Base charged each such file once, as its larger reading, against 256 KB; each reading is + # charged now, against 512 KB, so whatever base's total let through still fits. + pad = b"# " + b"x" * 76 + b"\n" + repo = _tree( + tmp_path / "repo", + {"agents/a/main.py": "import common.one\nimport common.two\n", "common/__init__.py": "", "common/web.py": "", "common/page.py": ""}, + ).resolve() + for name, target in (("one", "web"), ("two", "page")): + (repo / "common" / f"{name}.py").write_bytes( + _ESCAPED_LINE + f"import common.{target} # c\\r\n".encode() + pad * 1_500 + ) + larger = sum(len((repo / "common" / f"{n}.py").read_bytes()) for n in ("one", "two")) + assert larger < 256 * 1024 + + found, stopped = _save_walk(repo) + assert not stopped and {repo / "common" / "web.py", repo / "common" / "page.py"} <= found + + +@pytest.mark.parametrize("where", ["own-py", "imported-py", "imported-script", "grew-since"]) +def test_a_file_over_the_read_limit_says_the_walk_stopped(tmp_path, monkeypatch, where): + from fi.alk.harness import credentials + + big = b"import common.web\n" + b"# " + b"x" * 1_000_100 + b"\n" + repo = _tree(tmp_path / "repo", {"common/__init__.py": "", "common/web.py": ""}).resolve() + (repo / "agents" / "a").mkdir(parents=True) + if where == "own-py": + (repo / "agents" / "a" / "main.py").write_bytes(big) + elif where == "imported-py": + (repo / "agents" / "a" / "main.py").write_text("import common.big\n") + (repo / "common" / "big.py").write_bytes(big) + elif where == "imported-script": + (repo / "agents" / "a" / "index.js").write_text('import "./big.js";\n') + (repo / "agents" / "a" / "big.js").write_bytes(b'import "../../common/web.py";\n' + b"//" + b"x" * 1_000_100 + b"\n") + else: + (repo / "agents" / "a" / "main.py").write_bytes(big) + real = credentials._plain_file + # Small when looked at, past the limit once read. + monkeypatch.setattr( + credentials, "_plain_file", lambda path, unread=None: path.name == "main.py" or real(path, unread) + ) + + found, stopped = _save_walk(repo) + assert stopped + + +def test_a_large_data_file_a_script_imports_doesnt_say_stopped(tmp_path): + repo = _tree(tmp_path / "repo", {"agents/a/index.js": 'import data from "./data.json";\n'}).resolve() + (repo / "agents" / "a" / "data.json").write_bytes(b"[" + b"1," * 600_000 + b"1]") + + _, stopped = _save_walk(repo) + assert not stopped + + +@pytest.mark.parametrize("where", ["own-link", "imported-link"]) +def test_a_file_link_to_code_says_the_walk_stopped(tmp_path, where): + repo = _tree( + tmp_path / "repo", + {"shared/main.py": "import common.web\n", "common/__init__.py": "", "common/web.py": "", "common/real.py": "import common.web\n"}, + ).resolve() + (repo / "agents" / "a").mkdir(parents=True) + if where == "own-link": + (repo / "agents" / "a" / "main.py").symlink_to(repo / "shared" / "main.py") + else: + (repo / "agents" / "a" / "main.py").write_text("import common.helpers\n") + (repo / "common" / "helpers.py").symlink_to("real.py") + + found, stopped = _save_walk(repo) + # The link isn't followed (Python would follow it), so the walk says what it left. + assert stopped and repo / "common" / "web.py" not in found + + +@pytest.mark.parametrize( + "case", ["python-parts", "python-long", "python-deep-base", "script-long", "script-deep", "own-file-deep"] +) +def test_a_name_dropped_for_its_length_or_depth_says_the_walk_stopped(tmp_path, case): + from fi.alk.harness import credentials + + repo = (tmp_path / "repo").resolve() + agent = repo / "agents" / "a" + agent.mkdir(parents=True) + deep = credentials._DEEPEST_FOLLOWED + if case == "python-parts": + (agent / "main.py").write_text("import " + ".".join(["a"] * (deep + 1)) + "\n") + elif case == "python-long": + (agent / "main.py").write_text("import " + "a" * (credentials._LONGEST_NAME + 1) + "\n") + elif case == "python-deep-base": + # From a folder four parts below the root the module path is past the cut, while from the + # root's src/ even its __init__.py is within it: only the base's depth cuts it. + agent = repo / "agents" / "x" / "y" / "a" + agent.mkdir(parents=True) + (agent / "main.py").write_text("import " + ".".join(["a"] * (deep - 2)) + "\n") + elif case == "script-long": + (agent / "index.js").write_text('import "./' + "a" * (credentials._LONGEST_NAME + 1) + '";\n') + elif case == "script-deep": + (agent / "index.js").write_text('import "./' + "/".join(["a"] * (deep + 2)) + '";\n') + else: + (agent.joinpath(*["d"] * (deep + 1))).mkdir(parents=True) + (agent.joinpath(*["d"] * (deep + 1)) / "x.py").write_text("import os\n") + (agent / "main.py").write_text("import os\n") + + _, stopped = _save_walk(repo, str(agent.relative_to(repo))) + assert stopped + + +def test_a_name_whose_file_resolves_past_the_depth_cut_says_bundlings_walk_stopped(tmp_path): + from fi.alk.harness import credentials + from fi.alk.harness.credentials import imported_by_agent + + # Bundling looks files up as the kernel does, so a link can take a shallow name deep. + repo = (tmp_path / "repo").resolve() + far = repo.joinpath(*["d"] * (credentials._DEEPEST_FOLLOWED + 1)) + far.mkdir(parents=True) + (far / "mod.py").write_text("") + (repo / "agents" / "a").mkdir(parents=True) + (repo / "agents" / "a" / "main.py").write_text("import lib.mod\n") + (repo / "agents" / "a" / "lib").symlink_to(far) + stopped: list[bool] = [] + + imported_by_agent(repo, "agents/a", stopped=stopped) + assert stopped == [True] + + +def test_a_link_target_past_so_many_parts_says_the_walk_stopped(tmp_path): + from fi.alk.harness import credentials + + repo = _tree(tmp_path / "repo", {"common/__init__.py": "", "common/web.py": "", "x/keep": ""}).resolve() + (repo / "agents" / "a").mkdir(parents=True) + (repo / "agents" / "a" / "main.py").write_text("import shared.web\n") + target = "/".join(["..", "..", *(["x", ".."] * (credentials._LONGEST_TARGET // 2)), "common"]) + (repo / "agents" / "a" / "shared").symlink_to(target) + + found, stopped = _save_walk(repo) + assert stopped and repo / "common" / "web.py" not in found + + +def test_a_chain_past_the_link_limit_doesnt_say_the_walk_stopped(tmp_path): + from fi.alk.harness import credentials + + # The kernel refuses such a chain too, so nothing runs behind it. + repo = _tree(tmp_path / "repo", {"real/__init__.py": "", "real/web.py": ""}).resolve() + links = credentials._LINKS_FOLLOWED + 1 + for n in range(links): + (repo / f"c{n}").symlink_to(f"c{n + 1}" if n < links - 1 else "real") + (repo / "agents" / "a").mkdir(parents=True) + (repo / "agents" / "a" / "main.py").write_text("import c0.web\n") + + _, stopped = _save_walk(repo) + assert not stopped + + +def test_a_spent_decode_total_says_the_walk_stopped(tmp_path, monkeypatch): + from fi.alk.harness import python_source + + # The import is spelled in escapes, so read without decoding the file names nothing. + repo = _behind(tmp_path, _ESCAPED_LINE + b"\\x69mport common.web\n") + monkeypatch.setattr(python_source, "LARGEST_ESCAPED_TOTAL", 10) + + _, stopped = _save_walk(repo) + assert stopped + + +@pytest.mark.parametrize( + "helper", + [ + b"x = 1\n", + b"# no imports here, though the word import is in this comment\nx = 1\n", + b"# coding: utf-7\nimport os\n# +\xff\n", + b"import os\nimport json\n", + ], + ids=["no-import", "parses-empty", "python-refuses", "stdlib-name"], +) +def test_files_that_lead_nowhere_dont_say_the_walk_stopped(tmp_path, helper): + repo = _behind(tmp_path, helper) + + _, stopped = _save_walk(repo) + assert not stopped + + +def test_a_folder_where_a_module_would_be_doesnt_say_the_walk_stopped(tmp_path): + repo = _tree(tmp_path / "repo", {"agents/a/main.py": "import common.web\n", "common/__init__.py": ""}).resolve() + (repo / "common" / "web.py").mkdir() + + _, stopped = _save_walk(repo) + assert not stopped + + +@pytest.mark.parametrize("extra, stopped_too", [(0, False), (1, True)]) +def test_the_save_walk_reads_its_own_total(tmp_path, caplog, extra, stopped_too): + from fi.alk.harness import agent_type_check + + # What the walk charges, main.py included, comes to exactly 512 KB, or one byte more. + main = "".join(f"import common.m{n}\n" for n in range(5)) + total = 512 * 1024 + extra + each, rest = divmod(total - len(main), 5) + files = {"agents/a/main.py": main, "common/__init__.py": ""} + for n in range(5): + size = each + (rest if n == 4 else 0) + head = "import os\n" + files[f"common/m{n}.py"] = head + "#" * (size - len(head) - 1) + "\n" + repo = _tree(tmp_path / "repo", files).resolve() + assert sum(len(v) for k, v in files.items() if k != "common/__init__.py") == total + + with caplog.at_level("WARNING", logger=agent_type_check.__name__): + agent_type_check._walked(repo, "agents/a") + assert ("stopped at its limits" in caplog.text) is stopped_too + + +def test_the_save_walk_total_is_its_own(tmp_path, monkeypatch): + from fi.alk.harness import agent_type_check + from fi.alk.harness.python_source import LARGEST_PARSED_TOTAL, ParseBudget + + totals: list[object] = [] + + class Recorded(ParseBudget): + def __init__(self, total=None): + totals.append(total) + super().__init__(total) + + monkeypatch.setattr(agent_type_check, "ParseBudget", Recorded) + repo = _tree(tmp_path / "repo", {"agents/a/main.py": "import os\n"}).resolve() + + agent_type_check.browser_tools_in_agent(repo, "agents/a") + assert agent_type_check._WALKED_PARSED == 2 * LARGEST_PARSED_TOTAL + assert totals.count(agent_type_check._WALKED_PARSED) == 1 and None in totals + + +def test_a_second_reading_the_parse_total_refuses_says_the_walk_stopped(tmp_path): + from fi.alk.harness.credentials import imported_by_agent + from fi.alk.harness.python_source import ( + LARGEST_PARSED_MODULE, + EscapeBudget, + ParseBudget, + decoded_first, + ) + + # Both readings parse, but the total has room only for the first. + helper = _ESCAPED_LINE + b"import common.web # see \\r\n" + repo = _behind(tmp_path, helper) + # Nothing after it to spend the total: the second reading's refusal is the only one. + (repo / "common" / "web.py").write_text("") + main = (repo / "agents" / "a" / "main.py").read_text() + taken, text, _ = decoded_first(helper, EscapeBudget()) + first = len(main) + len(text.replace("\r\n", "\n").replace("\r", "\n").encode()) + stopped: list[bool] = [] + + found = imported_by_agent( + repo, "agents/a", parse_limit=LARGEST_PARSED_MODULE, parse_budget=ParseBudget(first), stopped=stopped + ) + assert taken and repo / "common" / "web.py" in found and stopped == [True] diff --git a/tests/harness/test_python_source_codecs.py b/tests/harness/test_python_source_codecs.py index 7fc63bd46..e096cfa4c 100644 --- a/tests/harness/test_python_source_codecs.py +++ b/tests/harness/test_python_source_codecs.py @@ -539,12 +539,11 @@ def test_the_walk_keeps_the_decoded_reading_beside_an_escaped_cr(tmp_path): assert repo / "common" / "web.py" in imported_by_agent(repo, "agents/a") -def test_a_file_read_two_ways_is_charged_once(tmp_path): +def test_a_file_read_two_ways_is_charged_for_each_reading(tmp_path): pad = b"# " + b"x" * 76 + b"\n" - repo = _shared( - tmp_path, _ESCAPED + b'SEP = "\\r"\nimport common.web\n' + pad * 1200 - ) - _write(repo / "agents" / "a", b"import common.tools\n" + pad * 1200, "b.py") + agent = _ESCAPED + b'SEP = "\\r"\nimport common.web\n' + pad * 600 + repo = _shared(tmp_path, agent) + _write(repo / "agents" / "a", b"import common.tools\n" + pad * 600, "b.py") _write(repo / "common", b"", "tools.py") stopped: list[bool] = [] budget = python_source.ParseBudget() @@ -555,13 +554,16 @@ def test_a_file_read_two_ways_is_charged_once(tmp_path): parse_budget=budget, stopped=stopped, ) - assert repo / "common" / "tools.py" in found and stopped == [False] - # Charged once, as the larger reading (the UTF-8 one, where `\\r` is two characters). - parsed = [repo / "agents" / "a" / n for n in ("agent.py", "b.py")] - parsed.append(repo / "common" / "web.py") - assert budget.left == python_source.LARGEST_PARSED_TOTAL - sum( - p.stat().st_size for p in parsed + # Each reading is charged as it's parsed: the decoded one, where `\\r` is one character, and + # the UTF-8 one, where it's two. + assert repo / "common" / "tools.py" in found + others = [repo / "agents" / "a" / "b.py", repo / "common" / "web.py"] + assert budget.left == python_source.LARGEST_PARSED_TOTAL - ( + 2 * len(agent) - 1 + sum(p.stat().st_size for p in others) ) + # The decoded reading, the one Python runs, is refused: its CR splits the string. The walk + # can't be sure it followed what Python runs, so it says it stopped. + assert stopped == [True] def test_the_union_is_for_unicode_escape_only(tmp_path): diff --git a/tests/test_harness_credentials.py b/tests/test_harness_credentials.py index 76143628a..bc7420560 100644 --- a/tests/test_harness_credentials.py +++ b/tests/test_harness_credentials.py @@ -1056,3 +1056,16 @@ def test_an_odd_script_import_does_not_stop_the_files_other_imports(tmp_path): ) assert (root / "common/util.js").resolve() in imported_by_agent(root, "agents/a") + + +def test_a_walk_that_cant_parse_a_file_marks_the_manifest_truncated(tmp_path): + root = tmp_path.resolve() + (root / "agents" / "a").mkdir(parents=True) + (root / "agents" / "a" / "main.py").write_text("import common.helpers\n") + (root / "common").mkdir() + (root / "common" / "__init__.py").write_text("") + (root / "common" / "helpers.py").write_text("from common import web\nNAME = 'a" + "\u0301" * 40 + "'\n") + (root / "common" / "web.py").write_text("import os\nKEY = os.environ['WEB_API_KEY']\n") + + manifest = discover_credentials(root, agent_subfolder="agents/a") + assert manifest.truncated From 20a3b6ba6a3a91f1b9ab707343126f9f5e253d9e Mon Sep 17 00:00:00 2001 From: khushalsonawat Date: Wed, 11 Nov 2026 05:49:39 +0530 Subject: [PATCH 3/3] fix(harness): answer the change review of TH-9543's walk (TH-9543) The change review cleared with nothing above Medium. Its findings: - A folder link in the agent's folder that leads inside the repository says the walk stopped: the walk doesn't enter it, and the agent may run code there. One that leads out of the repository, such as a data folder, doesn't. - Tests for a refused link target before a candidate's last folder, a reading sized with its line ends folded (decoded and UTF-8), a script with more names than the walk tries, and a decode total spent exactly. - The save's warning says the walk stopped at its limits or left code unread; two test comments say which Python versions they describe. Signed-off-by: khushalsonawat --- src/fi/alk/harness/agent_type_check.py | 3 +- src/fi/alk/harness/credentials.py | 14 ++- tests/harness/test_agent_type_refusal.py | 110 +++++++++++++++++++++-- 3 files changed, 117 insertions(+), 10 deletions(-) diff --git a/src/fi/alk/harness/agent_type_check.py b/src/fi/alk/harness/agent_type_check.py index b7ed80fe6..8ab8e1c63 100644 --- a/src/fi/alk/harness/agent_type_check.py +++ b/src/fi/alk/harness/agent_type_check.py @@ -654,7 +654,8 @@ def _walked(root: Path, subfolder: str) -> frozenset[Path]: ) if stopped and stopped[0]: logger.warning( - "the save's import walk for %s stopped at its limits; %d files reached", + "the save's import walk for %s stopped at its limits or left code unread; " + "%d files reached", subfolder, len(found), ) diff --git a/src/fi/alk/harness/credentials.py b/src/fi/alk/harness/credentials.py index 2585be9d1..2a31c61ef 100644 --- a/src/fi/alk/harness/credentials.py +++ b/src/fi/alk/harness/credentials.py @@ -1216,15 +1216,25 @@ def _followed( return real -def _code_in(folder: Path) -> Iterator[Path]: +def _code_in( + folder: Path, root: Path | None = None, unread: list[bool] | None = None +) -> Iterator[Path]: """The Python and script files in ``folder``, outside the folders no one's code is in. Those folders aren't entered, and each file costs its name: building a parent for each part, as 3.12 does to make a path relative, cost a crafted folder of 31,000 nested packages 2.8 s. + A folder link isn't entered either; given ``root`` and ``unread``, one that leads inside + ``root`` sets ``unread``, since code the agent runs may be there. """ suffixes = {".py", *_SCRIPT_SUFFIXES} for directory, subdirectories, files in os.walk(folder): subdirectories[:] = [name for name in subdirectories if name not in _IGNORED_PARTS] + if unread is not None and root is not None and not unread[0]: + for name in subdirectories: + place = os.path.join(directory, name) + if os.path.islink(place) and is_within(Path(os.path.realpath(place)), root): + unread[0] = True + break for name in files: if name not in _IGNORED_PARTS and os.path.splitext(name)[1] in suffixes: yield Path(directory, name) @@ -1482,7 +1492,7 @@ def script_module(path: Path, name: str) -> Path | None: # included (at most this much more parsing per walk). escapes = EscapeBudget() seen: set[Path] = set() - pending = deque(sorted(_code_in(folder))) + pending = deque(sorted(_code_in(folder, root, unread))) own = set(pending) # Read once already, so importing one doesn't queue it again. read = bytes_read = imports_tried = 0 import_limit = _MAX_WALK_IMPORTS if import_limit is None else import_limit diff --git a/tests/harness/test_agent_type_refusal.py b/tests/harness/test_agent_type_refusal.py index 39d7495fb..e8c11ed95 100644 --- a/tests/harness/test_agent_type_refusal.py +++ b/tests/harness/test_agent_type_refusal.py @@ -5146,7 +5146,7 @@ def test_a_stopped_save_walk_logs_how_many_files_it_reached(tmp_path, caplog, mo with caplog.at_level("WARNING", logger=agent_type_check.__name__): found = agent_type_check._walked(repo, "agents/a") - assert f"stopped at its limits; {len(found)} files reached" in caplog.text + assert f"stopped at its limits or left code unread; {len(found)} files reached" in caplog.text assert len(found) >= 1 @@ -5812,10 +5812,12 @@ def test_a_file_past_the_parse_total_says_the_walk_stopped(tmp_path, monkeypatch @pytest.mark.parametrize( "helper, stopped_too, reached", [ - # Decoded, the CR ends the comment and Python runs the import; read as UTF-8 the comment - # runs on and the file parses, with only os. + # Decoded, the CR ends the comment and Python 3.12 runs the import (3.11 doesn't: a false + # cut there, in the safe direction); read as UTF-8 the comment runs on and the file parses, + # with only os. (_ESCAPED_LINE + b"import os # c\\r; import common.web\n", True, False), - # Decoded, it's two imports; read as UTF-8, a stray backslash: only the decoded parses. + # Decoded, the walk reads two imports (Python refuses the bare CR, so following it only + # over-reaches); read as UTF-8, a stray backslash: only the decoded reading parses. (_ESCAPED_LINE + b"import os\\rimport common.web\n", False, True), # Both readings parse. (_ESCAPED_LINE + b"import common.web # see \\r\n", False, True), @@ -5946,7 +5948,8 @@ def test_a_name_whose_file_resolves_past_the_depth_cut_says_bundlings_walk_stopp (far / "mod.py").write_text("") (repo / "agents" / "a").mkdir(parents=True) (repo / "agents" / "a" / "main.py").write_text("import lib.mod\n") - (repo / "agents" / "a" / "lib").symlink_to(far) + # At the root, outside the agent's folder, so only the resolved depth can cut. + (repo / "lib").symlink_to(far) stopped: list[bool] = [] imported_by_agent(repo, "agents/a", stopped=stopped) @@ -5959,8 +5962,9 @@ def test_a_link_target_past_so_many_parts_says_the_walk_stopped(tmp_path): repo = _tree(tmp_path / "repo", {"common/__init__.py": "", "common/web.py": "", "x/keep": ""}).resolve() (repo / "agents" / "a").mkdir(parents=True) (repo / "agents" / "a" / "main.py").write_text("import shared.web\n") - target = "/".join(["..", "..", *(["x", ".."] * (credentials._LONGEST_TARGET // 2)), "common"]) - (repo / "agents" / "a" / "shared").symlink_to(target) + # At the root, outside the agent's folder, so only the link's own refusal can cut. + target = "/".join([*(["x", ".."] * (credentials._LONGEST_TARGET // 2)), "common"]) + (repo / "shared").symlink_to(target) found, stopped = _save_walk(repo) assert stopped and repo / "common" / "web.py" not in found @@ -6080,3 +6084,95 @@ def test_a_second_reading_the_parse_total_refuses_says_the_walk_stopped(tmp_path repo, "agents/a", parse_limit=LARGEST_PARSED_MODULE, parse_budget=ParseBudget(first), stopped=stopped ) assert taken and repo / "common" / "web.py" in found and stopped == [True] + + +def test_a_folder_link_in_the_agents_folder_says_the_walk_stopped(tmp_path): + from fi.alk.harness.credentials import imported_by_agent + + # The walk doesn't enter a folder link; one that leads inside the repository may hold code + # the agent runs. + repo = _tree( + tmp_path / "repo", + {"shared/handlers/run.py": "import common.web\n", "common/__init__.py": "", "common/web.py": "", "agents/a/main.py": "x = 1\n"}, + ).resolve() + (repo / "agents" / "a" / "handlers").symlink_to("../../shared/handlers") + stopped: list[bool] = [] + + _, at_save = _save_walk(repo) + imported_by_agent(repo, "agents/a", stopped=stopped) + assert at_save and stopped == [True] + + +def test_a_folder_link_out_of_the_repository_doesnt_say_the_walk_stopped(tmp_path): + outside = _tree(tmp_path / "data", {"rows.csv": "a,b\n"}).resolve() + repo = _tree(tmp_path / "repo", {"agents/a/main.py": "x = 1\n"}).resolve() + (repo / "agents" / "a" / "data").symlink_to(outside) + + _, stopped = _save_walk(repo) + assert not stopped + + +def test_a_long_link_target_before_the_last_folder_says_the_walk_stopped(tmp_path): + from fi.alk.harness import credentials + + # The refused link is the first of the candidate's folders: the walk stops there and says so. + repo = _tree(tmp_path / "repo", {"common/sub/__init__.py": "", "common/sub/web.py": "", "x/keep": ""}).resolve() + (repo / "agents" / "a").mkdir(parents=True) + (repo / "agents" / "a" / "main.py").write_text("import shared.sub.web\n") + target = "/".join([*(["x", ".."] * (credentials._LONGEST_TARGET // 2)), "common"]) + (repo / "shared").symlink_to(target) + + found, stopped = _save_walk(repo) + assert stopped and repo / "common" / "sub" / "web.py" not in found + + +@pytest.mark.parametrize("case", ["decoded", "utf8"]) +def test_a_reading_is_sized_with_its_line_ends_folded(tmp_path, case): + from fi.alk.harness.credentials import imported_by_agent + from fi.alk.harness.python_source import LARGEST_PARSED_MODULE, ParseBudget + + # Each file is over the per-file parse limit as written and under it once its line ends are + # folded, as Python folds them: sized as written, the reading would be refused. + root = tmp_path.resolve() + (root / "a").mkdir() + (root / "a" / "m.py").write_text("x = 1\n") + if case == "decoded": + # The decoded reading: each escaped CRLF is two characters, one once folded. + body = _ESCAPED_LINE + b"import m\\r\\n" * 13_200 + else: + # The UTF-8 reading, with real CRLFs: only it imports m. + pad = b"#" + b"p" * 40 + b"\r\n" + body = _ESCAPED_LINE.replace(b"\n", b"\r\n") + b"s = '\\r'; import m\r\n" + pad * 3_100 + assert len(body) > LARGEST_PARSED_MODULE + (root / "a" / "main.py").write_bytes(body) + stopped: list[bool] = [] + + found = imported_by_agent( + root, "a", parse_limit=LARGEST_PARSED_MODULE, parse_budget=ParseBudget(512 * 1024), stopped=stopped + ) + assert root / "a" / "m.py" in found + if case == "decoded": + assert stopped == [False] + + +def test_a_script_with_more_names_than_the_walk_tries_says_it_stopped(tmp_path): + from fi.alk.harness.credentials import imported_by_agent + + # None of them is there: the third is left untried with nothing else left to follow. + repo = _tree(tmp_path / "repo", {"agents/a/index.js": 'import "./b";\nimport "./c";\nimport "./d";\n'}) + stopped: list[bool] = [] + + imported_by_agent(repo, "agents/a", import_limit=2, stopped=stopped) + assert stopped == [True] + + +def test_a_decode_total_spent_exactly_doesnt_say_the_walk_stopped(tmp_path, monkeypatch): + from fi.alk.harness import python_source + + helper = _ESCAPED_LINE + b"\\x69mport common.web\n" + repo = _behind(tmp_path, helper) + # The walk's decode total has room for this file and no byte more. + monkeypatch.setattr(python_source, "LARGEST_ESCAPED_TOTAL", len(helper)) + + found, stopped = _save_walk(repo) + assert not stopped and repo / "common" / "web.py" in found