Преглед изворни кода

fix(code-runtime-python): bound three child-side walks by depth, not width

Three separate paths in the CPython child allocated state proportional to a
value's width or a string's length, so a legitimate input the byte budgets
admit could die as the program's own MemoryError.

`_lossless_json_violation` enqueued one traversal tuple per member while
running, in `dispatch`, over MODEL-CONSTRUCTED binding arguments that no
child-side byte budget bounds first. It now uses the same (kind, container,
iterator) cursor the other two walks already had, checking dict keys as the
cursor pulls each entry. Measured over `[0] * 6_000_000` (~17 MB of JSON):
459.1 MiB of traversal tuples before, 0.0 MiB after.

`_decode_json_plain` matched JSON strings with a `(?:[^"\\]|\\.)*` repetition,
which makes CPython's engine retain backtracking state proportional to the
string's width: 146 MiB for a 1 MiB string, 557.8 MiB for 4 MiB. A legitimate
multi-megabyte binding reply raised MemoryError inside `_pump_replies`, and
because that pump is the only settler of the call's future, the run stranded
until the wall clock reported `timeout`. Strings now scan chunk-to-chunk over a
character class, which the engine matches without backtracking state; the same
4 MiB decode peaks at the 4.0 MiB result.

`_check_done_value` charged strings and dict keys what
`_dump_string(...).encode()` returned, building the escaped copy plus its
encode to MEASURE it -- ~6x the original each for control-heavy text, so
metering a value the budget then rejects could itself breach RLIMIT_AS and
report `exception` where the seam promises `output-limit`. The new
`_json_str_cost` counts instead, reusing `_json_string_cost`'s C-level passes
and reproducing `_dump_string`'s exact surrogate rules (fold spelled-out pairs,
charge six ASCII bytes per lone surrogate). Identical values, 228.9 MiB -> 19.1
MiB of peak on a 20M-NUL string.

Each fix ships a regression test. The two RLIMIT_AS repros are Linux-only:
Darwin does not apply the limit, so the peaks above are measured directly and
recorded in the test comments.
Chinesezjc пре 2 недеља
родитељ
комит
8f7d9121d1

+ 1 - 0
.gitignore

@@ -31,6 +31,7 @@ python/sdk-runtime/src/deepseek_harness_runtime/runtime/deepseek-harness-sdk-run
 python/sdk-runtime/src/deepseek_harness_runtime/runtime/node/
 python/**/__pycache__/
 python/**/.pytest_cache/
+packages/**/__pycache__/
 apps/web/dist/
 .artifacts/
 .dsh-build/

+ 113 - 18
packages/code-runtime/code-runtime-python/py/bootstrap.py

@@ -991,10 +991,22 @@ async def _pump_replies(
             continue
 
 
+# Non-string scalars only. The string form is scanned by hand in
+# :func:`_decode_json_plain` because a ``(?:[^"\\]|\\.)*`` repetition makes
+# CPython's backtracking engine retain per-repetition state proportional to the
+# string's WIDTH: measured at ~146 MiB of engine state for a 1 MiB string and
+# ~558 MiB for 4 MiB, so a legitimate multi-megabyte binding reply raised
+# MemoryError out of ``_pump_replies``, leaving its future unsettled until the
+# wall clock reported a timeout.
 _SCALAR_RE = re.compile(
-    r'"(?:[^"\\]|\\.)*"|-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?|true|false|null'
+    r'-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?|true|false|null'
 )
 
+# A run of ordinary string body characters. The star applies to a CHARACTER
+# CLASS, which the engine matches in one linear pass with no backtracking state,
+# so the scanner's cost is the number of escapes, not the string's width.
+_STRING_CHUNK_RE = re.compile(r'[^"\\]*')
+
 
 def _decode_json_plain(text: str) -> Any:
     """Parse one JSON document iteratively (no per-level recursion).
@@ -1017,7 +1029,26 @@ def _decode_json_plain(text: str) -> Any:
             i += 1
         return i
 
+    def scan_string(i: int) -> int:
+        # Walk chunk by chunk: each match consumes every character up to the next
+        # quote or backslash, so an escape costs one extra step and a plain body
+        # costs one pass. Returns the offset just past the closing quote.
+        j = i + 1
+        while True:
+            j = _STRING_CHUNK_RE.match(text, j).end()
+            if j >= length:
+                raise ValueError(f"unterminated string at offset {i}")
+            char = text[j]
+            if char == '"':
+                return j + 1
+            # text[j] is a backslash: skip it and the character it escapes. A
+            # trailing backslash runs j past `length`, caught on the next pass.
+            j += 2
+
     def scalar(i: int):
+        if i < length and text[i] == '"':
+            end = scan_string(i)
+            return json.loads(text[i:end]), end
         match = _SCALAR_RE.match(text, i)
         if match is None:
             raise ValueError(f"invalid JSON at offset {i}")
@@ -1271,6 +1302,39 @@ for _escaped_byte, _surcharge in _JSON_ESCAPE_SURCHARGES:
     _JSON_BYTE_COST[_escaped_byte[0]] = 1 + _surcharge
 
 
+def _json_str_cost(text: str) -> int:
+    """Byte length of ``text``'s JSON string form, WITHOUT building that form.
+
+    The str-side twin of :func:`_json_string_cost`, for the completion-value
+    meter. Measuring by materializing ``_dump_string(text).encode()`` allocates
+    the escaped copy plus its encode -- for a NUL-heavy string that is ~6x the
+    original each, so metering a value the budget would have REJECTED could
+    itself breach ``RLIMIT_AS`` and report ``exception`` where the contract
+    promises ``output-limit``.
+
+    The common case encodes once (~1x, well inside the load gate's envelope) and
+    counts escapes with the same C-level passes :func:`_json_string_cost` uses.
+    A string carrying surrogate code units has no UTF-8 form at all, so it takes
+    the exact path :func:`_dump_string` defines: fold each spelled-out high-low
+    pair into its astral character first (the host meters that as its raw 4-byte
+    form), then charge six ASCII bytes for every surviving lone surrogate and
+    count the rest from its encodable remainder.
+    @param text: the string to measure.
+    @return: the byte length of its JSON string form, quotes included.
+    """
+
+    try:
+        return _json_string_cost(text.encode("utf-8"))
+    except UnicodeEncodeError:
+        pass
+    folded = _SURROGATE_PAIR.sub(_combine_surrogate_pair, text)
+    lone = len(_SURROGATE.findall(folded))
+    # Six ASCII bytes per lone surrogate; the remainder is ordinary text whose
+    # own quotes are dropped here because the outer call adds them once.
+    without = _SURROGATE.sub("", folded)
+    return _json_string_cost(without.encode("utf-8")) + lone * 6
+
+
 def _json_string_cost(raw: bytes) -> int:
     """UTF-8 byte length of one string's JSON form, WITHOUT building that form.
 
@@ -1435,7 +1499,9 @@ def _check_done_value(value: Any, max_bytes: int):
             # The same string lower bound, before escaping the key.
             if total + len(key) + 3 > max_bytes:
                 return over_budget
-            total += len(_dump_scalar(key).encode("utf-8")) + 1
+            # Same counting rule as the string branch: a control-heavy KEY
+            # expands just as far, and `_dump_scalar` on a str is `_dump_string`.
+            total += _json_str_cost(key) + 1
             if total > max_bytes:
                 return over_budget
             stack.append(frame)
@@ -1453,8 +1519,12 @@ def _check_done_value(value: Any, max_bytes: int):
                 return over_budget
             # A lone surrogate has no UTF-8 form but a lossless JSON one — the
             # ASCII ``\uXXXX`` escape :func:`_dump_string` emits — so it is
-            # metered, not rejected, matching the shared seam.
-            total += len(_dump_string(current).encode("utf-8"))
+            # metered, not rejected, matching the shared seam. Metered by
+            # COUNTING, not by building the escaped form: that copy plus its
+            # encode is ~6x the original for a control-heavy string, so measuring
+            # a value the budget rejects could breach RLIMIT_AS and surface as
+            # `exception` instead of the promised `output-limit`.
+            total += _json_str_cost(current)
         elif type(current) is int:
             # The canonical boundary accepts every JS-double-exact value: an int
             # outside +-2**53-1 is fine IFF the double round-trip is exact.
@@ -1539,13 +1609,41 @@ def _lossless_json_violation(value: Any) -> str | None:
     # ancestor (a cycle) is detected without rejecting a legitimately shared
     # acyclic subtree.
     on_path: set[int] = set()
-    # Each frame is (value, is_leave): a leave frame pops its container off the path.
-    stack: list[tuple[Any, bool]] = [(value, False)]
+    # O(DEPTH) auxiliary space, not O(width), for the reason
+    # :func:`_check_done_value` documents: this walk runs in ``dispatch`` on
+    # MODEL-CONSTRUCTED binding arguments, which no child-side byte budget
+    # bounds first (the frame ceiling is the host's, and it applies after this
+    # returns). Enqueueing one frame per member would let a legitimate
+    # ``[0] * 6_000_000`` argument -- ~17 MB of JSON -- allocate ~366 MB of
+    # traversal tuples and die as the program's own MemoryError instead of
+    # round-tripping. A container therefore pushes ONE cursor frame holding its
+    # iterator; children are pulled one at a time.
+    exhausted = object()
+    visit, container_cursor = 0, 1
+    # A visit frame is (visit, value); a cursor frame is (cursor, container, iterator).
+    stack: list[tuple[int, Any, Any]] = [(visit, value, None)]
     while stack:
-        current, is_leave = stack.pop()
-        if is_leave:
-            on_path.discard(id(current))
+        kind = stack[-1][0]
+        if kind == container_cursor:
+            _, container, iterator = stack[-1]
+            child = next(iterator, exhausted)
+            if child is exhausted:
+                # Leaving the container: it is no longer on the current path, so
+                # a legitimately shared acyclic subtree is not mistaken for a cycle.
+                on_path.discard(id(container))
+                stack.pop()
+                continue
+            if type(container) is dict:
+                # The dict cursor yields (key, value): check the key as it is
+                # pulled. Only an EXACT str key survives -- int, float, None, and
+                # tuple keys coerce or raise, and a str subclass can carry
+                # overrides the encoder does not honor.
+                key, child = child
+                if type(key) is not str:
+                    return f"non-string dict key ({type(key).__name__})"
+            stack.append((visit, child, None))
             continue
+        _, current, _unused = stack.pop()
         if current is None or type(current) is bool:
             continue
         if type(current) is str:
@@ -1579,17 +1677,14 @@ def _lossless_json_violation(value: Any) -> str | None:
             if id(current) in on_path:
                 return "circular reference"
             on_path.add(id(current))
-            stack.append((current, True))
             if type(current) is dict:
-                for key in current:
-                    # Only an EXACT str key survives: int, float, None, and
-                    # tuple keys coerce or raise, and a str subclass can carry
-                    # overrides the encoder does not honor.
-                    if type(key) is not str:
-                        return f"non-string dict key ({type(key).__name__})"
-                stack.extend((child, False) for child in current.values())
+                # Keys are checked as the cursor pulls each entry, not in a
+                # separate pass: ``current.values()`` would need a second walk,
+                # and materializing ``items()`` up front allocates one tuple per
+                # member -- the very spike the cursor removes.
+                stack.append((container_cursor, current, iter(current.items())))
             else:
-                stack.extend((child, False) for child in current)
+                stack.append((container_cursor, current, iter(current)))
             continue
         return f"unsupported type ({type(current).__name__})"
     return None

+ 73 - 0
packages/code-runtime/code-runtime-python/tests/runtime.spec.ts

@@ -1079,6 +1079,28 @@ describe('PythonCodeRuntime — programs and bindings', () => {
     expect(result.error?.message).toContain('exceeded 64 bytes')
   })
 
+  it('meters a control-heavy completion value without materializing its escaped form', async () => {
+    // The child's lower bound admits a string by CHARACTER count, then the meter
+    // charged what `_dump_string(current).encode()` returned -- building the
+    // escaped copy plus its encode. Each NUL escapes to six bytes, so metering a
+    // value the budget then REJECTS allocated ~6x the original twice over:
+    // measured at 228.9 MiB of peak for a 20M-NUL string, against 19.1 MiB for
+    // the counting path that returns the identical 120,000,002 bytes. Past
+    // RLIMIT_AS the meter died as `exception: MemoryError`, inverting the
+    // `output-limit` this seam promises for an over-budget value.
+    //
+    // 8M NULs is 8,000,002 raw but 48,000,002 escaped: over the 16 MiB budget
+    // only when charged the escaped cost, so this also pins that the cheap
+    // character bound alone does not decide the verdict.
+    const { runtime } = await setup({ maxValueBytes: 16 * 1024 * 1024, maxWallMs: 60_000 })
+    const result = await runtime.run({
+      program: 'return "\\x00" * 8_000_000',
+      bindings: [],
+    })
+    expect(result.value).toBeUndefined()
+    expect(result.error?.kind).toBe('output-limit')
+  }, 90_000)
+
   it('rejects a wide completion as output-limit before materializing its traversal state', async () => {
     // `[0] * 2000000` sits far above maxValueBytes but well below the frame
     // ceiling. The folded checker must reject it via the pre-enqueue bound —
@@ -3715,6 +3737,57 @@ describe('PythonCodeRuntime — hostile peer', () => {
     expect((result.value as number[]).length).toBe(6_000_000)
   }, 90_000)
 
+  it('validates wide binding arguments in O(depth), not O(width)', async () => {
+    // The completion-value walks are budgeted; this one is not. `dispatch` runs
+    // `_lossless_json_violation` on the arguments the MODEL built, and no
+    // child-side byte budget bounds them first: the frame ceiling is the host's
+    // and applies only after this validation returns. A per-member traversal
+    // frame therefore turned a legitimate call into the program's own
+    // MemoryError. Measured with tracemalloc on the two walk shapes over this
+    // exact argument (JSON ~17 MB): the cursor peaks at 0.0 MiB of auxiliary
+    // state, the pre-fix `stack.extend` at 459.1 MiB -- past the 384 MiB
+    // configured below, so the discriminating failure is real. It is Linux-only:
+    // Darwin skips RLIMIT_AS, so this case round-trips there either way.
+    //
+    // The binding echoes its argument's length back, so the assertion proves the
+    // call actually round-tripped rather than merely avoiding a crash.
+    const { runtime } = await setup({ addressSpaceMb: 384, maxWallMs: 60_000 })
+    const result = await runtime.run({
+      program: 'return await tools.width([0] * 6_000_000)',
+      bindings: [{
+        global: 'tools',
+        functions: { width: async (items: unknown) => (items as number[]).length },
+      }],
+    })
+    expect(result.error).toBeUndefined()
+    expect(result.value).toBe(6_000_000)
+  }, 90_000)
+
+  it('decodes a multi-megabyte binding reply without regex backtracking state', async () => {
+    // The child parses every host reply with `_decode_json_plain`. Its scalar
+    // regex matched strings with a `(?:[^"\\]|\\.)*` repetition, which makes
+    // CPython's backtracking engine retain state proportional to the string's
+    // WIDTH -- measured at ~146 MiB of engine state for a 1 MiB string and
+    // ~558 MiB for 4 MiB. A legitimate multi-megabyte reply therefore raised
+    // MemoryError inside `_pump_replies`; because that pump is the only settler
+    // of the call's future, the run stranded until the wall clock reported a
+    // `timeout` instead of returning the value the binding produced.
+    //
+    // Strings now scan chunk-to-chunk over a character class (no backtracking
+    // state). Measured on this exact 4 MiB reply: the pre-fix regex peaks at
+    // 557.8 MiB, past the default 512 MiB address space, while the scanner peaks
+    // at the 4.0 MiB result itself. Linux-only, like the other RLIMIT_AS repros:
+    // Darwin does not apply the limit, so the spike is merely allocated there.
+    const reply = 'A'.repeat(4 * 1024 * 1024)
+    const { runtime } = await setup({ maxWallMs: 60_000 })
+    const result = await runtime.run({
+      program: 'value = await tools.big({})\nreturn len(value)',
+      bindings: [{ global: 'tools', functions: { big: async () => reply } }],
+    })
+    expect(result.error).toBeUndefined()
+    expect(result.value).toBe(reply.length)
+  }, 90_000)
+
   it('bounds a flood of zero-byte log lines through the per-entry separator charge', async () => {
     // Blank print() lines carry zero content bytes; without the +1 separator
     // charge they would bypass maxLogBytes entirely and grow the retained