Explorar o código

fix(code-runtime-python): correct the open-merge cap arithmetic on both sides

The review's arithmetic checks: the closing-frame walk used cap
logBudget - openCost, so a compliant merged entry (58-byte wire cost under a
64-byte budget) could see a negative cap and truncate; the first-fragment cap
used logBudget instead of the ledger's logBudget - 1, so an open frame costing
63 was admitted with a bill of 64, pushing the ledger negative and letting a
subsequent empty frame ride in one byte past the configured cap; and the child
billed a closing frame as a fresh entry (quotes+separator again) instead of the
merged tail, truncating an exact-fit 30+30 entry.

Fixes: first-fragment cap logBudget - 1 (matching admit), continuation and
closing-frame cap logBudget + 2 (billed without quotes), jsonStringCostUpTo
returns undefined below 2 bytes, and the child's split billing keys off
_open_started alone (a closing frame pays content only) with the cheaper bound
len(text) while a merge is open. Regression cases cover all three arithmetic
paths.
Chinesezjc hai 1 mes
pai
achega
3001cc23be

+ 8 - 6
packages/code-runtime/code-runtime-python/py/bootstrap.py

@@ -167,7 +167,7 @@ class LogBuffer:
         # above the budget truncates without ever encoding it — the full encode
         # would allocate a second equally large string and could turn a
         # truncatable log into an RLIMIT_AS death.
-        if (len(text) + 3 if not open or not self._open_started else len(text) + 1) > self._remaining:
+        if (len(text) + 3 if not self._open_started else len(text)) > self._remaining:
             self._truncated = True
             self._sink(log_truncation_marker(self._max_bytes), truncated=True)
             return
@@ -193,11 +193,13 @@ class LogBuffer:
         # instead of emitting the truncation marker. The +1 also floors an empty
         # entry above zero, so a flood of blank ``print()`` lines exhausts the
         # budget instead of emitting unbounded zero-cost log frames.
-        # Split billing for an `open` entry: the first fragment pays the full
-        # JSON-string cost plus the separator; each continuation pays only its
-        # content (the quotes and the separator were billed on the first
-        # fragment). A closed entry pays the full cost as before.
-        if open and self._open_started:
+        # Split billing for a merged entry: the FIRST fragment pays the full
+        # JSON-string cost plus the separator; every later fragment — a
+        # continuation OR the closing frame (it is the merged entry's tail, not
+        # a new entry) — pays only its content, since the quotes and separator
+        # were billed on the first fragment. A standalone closed entry (no open
+        # in progress) pays the full cost as before.
+        if self._open_started:
             cost = _json_string_cost(raw) - 2
             if cost < 0:
                 cost = 0

+ 13 - 11
packages/code-runtime/code-runtime-python/src/index.ts

@@ -462,6 +462,7 @@ function serializedCharCost(code: number, character: string): number {
  * @returns the exact serialized byte cost, or `undefined` once it exceeds `maxBytes`.
  */
 function jsonStringCostUpTo(text: string, maxBytes: number): number | undefined {
+  if (maxBytes < 2) return undefined
   let bytes = 2 // the enclosing quotes
   for (const character of text) {
     bytes += serializedCharCost(character.codePointAt(0) as number, character)
@@ -1028,11 +1029,9 @@ export class PythonCodeRuntime extends CodeRuntime {
       // An unterminated line flushed with the `open` flag: the next log frame
       // appends to it (no fake newline between entries), and finish() pushes
       // the residual if the run ends with it still open. Held as a fragment
-      // ARRAY with an incrementally billed content cost, so k tiny open frames
-      // cost O(k) — re-joining and re-walking the whole held text per frame
-      // would be O(k * budget) (jsonStringCostUpTo re-walks from the start).
+      // ARRAY, so k tiny open frames cost O(k) — re-joining and re-walking the
+      // whole held text per frame would be O(k * budget).
       let openParts: string[] = []
-      let openCost = 0
 
       // One host-side ledger covers normal frames, forged frames, and stray stdout bytes.
       // The ledger starts one byte below maxLogBytes: each entry is charged its
@@ -1515,20 +1514,24 @@ export class PythonCodeRuntime extends CodeRuntime {
               // its content (jsonStringCostUpTo includes the two quotes), and
               // the closing frame only its own content — the merged entry's
               // wire cost is billed exactly once, split across the fragments.
+              // Caps: the first fragment's exact-cost walk uses logBudget - 1
+              // (the ledger's reserved byte, matching admit), a continuation's
+              // logBudget + 2 (a continuation is billed WITHOUT quotes, so its
+              // billed cost cost - 2 fits exactly when the walk's cost is at
+              // most logBudget + 2).
               if (!logsTruncated) {
-                const cost = jsonStringCostUpTo(message.text, logBudget - openCost)
+                const cap = openParts.length === 0 ? logBudget - 1 : logBudget + 2
+                const cost = jsonStringCostUpTo(message.text, cap)
                 if (cost === undefined) {
                   logsTruncated = true
                   logs.push(logTruncationMarker(this.config.maxLogBytes))
                   clearStray(strayOut)
                   clearStray(strayErr)
                   openParts = []
-                  openCost = 0
                 } else {
                   const bill = openParts.length === 0 ? cost + 1 : Math.max(cost - 2, 0)
                   logBudget -= bill
                   openParts.push(message.text)
-                  openCost += bill
                 }
               }
               return
@@ -1536,12 +1539,13 @@ export class PythonCodeRuntime extends CodeRuntime {
             if (openParts.length > 0) {
               // Closing frame: the held fragments are already billed; bill only
               // this frame's own content (the quotes and separator ride on the
-              // first fragment) and push the merged entry once.
+              // first fragment) and push the merged entry once. Cap is
+              // logBudget + 2 for the same reason as a continuation.
               /* v8 ignore next -- logsTruncated is an invariant false here: an open
                * frame that would trip the ledger resets openParts, so a non-empty
                * hold implies the ledger never truncated. The guard is defensive. */
               if (!logsTruncated) {
-                const cost = jsonStringCostUpTo(message.text, logBudget - openCost)
+                const cost = jsonStringCostUpTo(message.text, logBudget + 2)
                 if (cost === undefined) {
                   logsTruncated = true
                   logs.push(logTruncationMarker(this.config.maxLogBytes))
@@ -1553,7 +1557,6 @@ export class PythonCodeRuntime extends CodeRuntime {
                 }
               }
               openParts = []
-              openCost = 0
               return
             }
             admit(message.text)
@@ -1942,7 +1945,6 @@ export class PythonCodeRuntime extends CodeRuntime {
           logs.push(openParts.join(''))
         }
         openParts = []
-        openCost = 0
         if (child.pid === undefined) {
           settle(result)
           return

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

@@ -1938,6 +1938,60 @@ describe('PythonCodeRuntime — programs and bindings', () => {
     expect(result.logs).toEqual(['x'.repeat(16)])
   }, 15_000)
 
+  it('admits a compliant merged entry whose closing frame fits the remaining budget', async () => {
+    // The review's arithmetic check: print('a'*30, flush); print('b'*25) under
+    // maxLogBytes: 64 has a merged wire cost of 2 quotes + 55 content + 1
+    // separator = 58 <= 63, so it MUST be admitted as one entry. The earlier
+    // cap math (logBudget - openCost) made the closing frame's walk see a
+    // negative cap and truncate a compliant entry.
+    const { runtime } = await setup({ maxLogBytes: 64 })
+    const result = await runtime.run({
+      program: [
+        "print('a' * 30, end='', flush=True)",
+        "print('b' * 25)",
+        'return "done"',
+      ].join('\n'),
+      bindings: [],
+    })
+    expect(result.error).toBeUndefined()
+    expect(result.logs).toEqual(['a'.repeat(30) + 'b'.repeat(25)])
+  }, 15_000)
+
+  it('rejects an open frame that would overflow the ledger by one byte', async () => {
+    // The review's arithmetic check: an open frame whose full JSON cost is 63
+    // (maxLogBytes: 64 -> ledger 63) must be rejected by the first-fragment
+    // cap logBudget - 1 (62), not admitted with a bill of 64 that pushes the
+    // ledger negative.
+    const { runtime } = await setup({ maxLogBytes: 64 })
+    const result = await runtime.run({
+      program: [
+        "print('x' * 61, end='', flush=True)",
+        'return "done"',
+      ].join('\n'),
+      bindings: [],
+    })
+    expect(result.error).toBeUndefined()
+    expect(result.logs).toEqual([logTruncationMarker(64)])
+  }, 15_000)
+
+  it('bills the closing frame as the merged tail under an exact-fit budget', async () => {
+    // The child's split billing: a 30-char open + a 30-char closing frame cost
+    // 2 + 60 + 1 = 63 = ledger 63 exactly; the closing frame must be billed as
+    // the merged tail (content only), not as a fresh entry (which would
+    // double-charge the quotes+separator and truncate an exact-fit entry).
+    const { runtime } = await setup({ maxLogBytes: 64 })
+    const result = await runtime.run({
+      program: [
+        "print('a' * 30, end='', flush=True)",
+        "print('b' * 30)",
+        'return "done"',
+      ].join('\n'),
+      bindings: [],
+    })
+    expect(result.error).toBeUndefined()
+    expect(result.logs).toEqual(['a'.repeat(30) + 'b'.repeat(30)])
+  }, 15_000)
+
   it('truncates when the closing frame of a merged entry overflows the budget', async () => {
     // The merged entry's billed-once cost: an open fragment that nearly
     // exhausts the budget, then a closing frame whose content no longer fits —