fix(code-runtime-python): reject only an oversized FIRST frame before the join, not a multi-frame buffer

The pre-join check charged the whole unframed buffer, which legitimately holds
several frames each within FRAME_PARSE_CAP_BYTES: a first frame of exactly the
cap followed by a second frame crossed the counter and was misreported as a
worker-exit. The pre-join rejection now fires only while the held bytes are a
single unframed line (this chunk carries no newline); once a newline arrives,
a FIRST-FRAME check measures the bytes up to the first newline across the held
chunks (including sealed blocks) and rejects only that frame before the join —
keeping the peak at one copy of its wire bytes — while later frames in the
same buffer are handled by the restored per-line check. Regression cases: a
72 MiB newline-free buffer is rejected pre-join (fail-before: joining would
have doubled it); two within-cap frames whose combined buffer crosses the cap
both survive (fail-before: the unconditional counter check turns it red).

packages/code-runtime/code-runtime-python/src/index.ts

@@ -1320,13 +1320,17 @@ export class PythonCodeRuntime extends CodeRuntime {
         // defaults.
         //
         // The cap used HERE is FRAME_PARSE_CAP_BYTES, not the 256 MiB wire
-        // ceiling: a single frame between 64 MiB and the ceiling would otherwise
-        // be fully `Buffer.concat`-ed (a second copy of its bytes) and only then
-        // dropped in the line loop — the peak-memory doubling this pre-concat
-        // check exists to prevent, now for a frame the parser is guaranteed to
-        // discard. Dropping the oversized unframed buffer before the join keeps
-        // the peak at one copy of the wire bytes.
-        if (pendingBytes > FRAME_PARSE_CAP_BYTES) {
+        // ceiling, and ONLY when the held bytes are still a single unframed
+        // line (this chunk carries no newline, and earlier newline-bearing
+        // chunks were joined immediately): a frame between 64 MiB and the
+        // ceiling would otherwise be fully `Buffer.concat`-ed (a second copy
+        // of its bytes) and only then dropped in the line loop — the
+        // peak-memory doubling this pre-concat check exists to prevent.
+        // Dropping the oversized unframed buffer before the join keeps the
+        // peak at one copy of the wire bytes. When this chunk DOES carry a
+        // newline the buffer holds several frames, so the FIRST-FRAME check
+        // below (not this counter, which charges them all) decides.
+        if (pendingBytes > FRAME_PARSE_CAP_BYTES && !chunk.includes(0x0a)) {
           pendingChunks = []
           sealedBlocks = []
           pendingBytes = 0
@@ -1361,6 +1365,35 @@ export class PythonCodeRuntime extends CodeRuntime {
           pendingChunks = []
         }
         if (chunk.includes(0x0a)) {
+          // First-FRAME check before the join: measure the bytes up to the
+          // first newline across the held chunks. The byte counter cannot
+          // serve here — it charges the whole buffer, which legitimately
+          // holds several frames each within the cap. A first frame past the
+          // cap is dropped before the join (one copy of its wire bytes);
+          // later frames in the same buffer are handled by the per-line check
+          // in the loop below.
+          let firstFrameLen = 0
+          let sawNewline = false
+          // Sealed blocks hold newline-free prefixes only (a newline-bearing
+          // chunk is joined immediately), so they are entirely part of the
+          // first frame.
+          for (const b of sealedBlocks) firstFrameLen += b.length
+          for (const c of pendingChunks) {
+            const nl = c.indexOf(0x0a)
+            if (nl >= 0) {
+              firstFrameLen += nl
+              sawNewline = true
+              break
+            }
+            firstFrameLen += c.length
+          }
+          if (sawNewline && firstFrameLen > FRAME_PARSE_CAP_BYTES) {
+            pendingChunks = []
+            sealedBlocks = []
+            pendingBytes = 0
+            finish({ error: { kind: 'worker-exit', message: `protocol frame exceeded ${FRAME_PARSE_CAP_BYTES} bytes on fd 3` } })
+            return
+          }
           let buffered = Buffer.concat(sealedBlocks.length > 0 ? [...sealedBlocks, ...pendingChunks] : pendingChunks)
           sealedBlocks = []
           let newline: number
@@ -1369,10 +1402,10 @@ export class PythonCodeRuntime extends CodeRuntime {
             buffered = buffered.subarray(newline + 1)
             /* v8 ignore next -- an empty line comes only from a forged `\n\n` write. */
             if (line.length === 0) continue
-            // No per-line cap check here: the unframed-buffer counter above
-            // already guarantees every line is within FRAME_PARSE_CAP_BYTES
-            // before this join runs, so a cap check on the line would be
-            // dead code.
+            // A later frame in this buffer may still exceed the cap; drop that
+            // single line like any junk frame (the first frame was already
+            // bounded by the check above).
+            if (line.length > FRAME_PARSE_CAP_BYTES) continue
             const text = line.toString('utf8')
             // JSON.parse would silently ROUND an integer token outside the
             // safe range before validation could see it, so a forged frame

packages/code-runtime/code-runtime-python/tests/runtime.spec.ts

@@ -4847,35 +4847,25 @@ describe('PythonCodeRuntime — hostile peer', () => {
     expect(result.error?.message).toContain(`protocol frame exceeded ${ceiling} bytes`)
   }, 90_000)
 
-  it('rejects an over-ceiling fd-3 buffer without first joining it into one line', async () => {
-    // The ceiling has to be enforced on the byte COUNTER before Buffer.concat,
+  it('rejects an over-cap newline-free fd-3 buffer without first joining it into one line', async () => {
+    // The cap has to be enforced on the byte COUNTER before Buffer.concat,
     // not on the joined line afterwards: the join is a second copy of
     // everything held, so a program could force roughly twice the advertised
-    // 256 MiB of host memory before anything rejected it.
+    // 64 MiB of host memory before anything rejected it.
     //
-    // This program makes the two orders observably different rather than merely
-    // differently sized. It writes exactly the ceiling with no newline (at the
-    // limit, so nothing trips), then a newline followed by 8 MiB more. Checking
-    // the counter first sees more than the ceiling on the newline-bearing pipe
-    // chunk and rejects. Checking the joined line instead produced a FIRST LINE
-    // of exactly the ceiling — inside the per-line bound, so it passed as a junk
-    // frame — and left an 8 MiB residual well under the bound, so the breach was
-    // never reported: measured, the run settled as
-    // `python exited (code=0, signal=null) before completing` after the host had
-    // held the ceiling AND copied it, which is the doubling this check prevents.
+    // This program writes past the cap with no newline: the counter crosses on
+    // the 9th 8 MiB write (72 MiB) while the buffer is still a single unframed
+    // line, so the pre-join check rejects it without concat-ing a second copy.
+    // Checking the joined line instead would have produced a 72 MiB FIRST LINE
+    // that the per-line bound then dropped only after the doubling had happened.
     const ceiling = 64 * 1024 * 1024
     const { runtime } = await setup({ maxWallMs: 60_000, addressSpaceMb: 2048 })
     const result = await runtime.run({
       program: [
         'import os',
         'chunk = b"A" * (8 * 1024 * 1024)',
-        `for _ in range(${ceiling / (8 * 1024 * 1024)}):`,
+        `for _ in range(${ceiling / (8 * 1024 * 1024) + 1}):`,
         '    os.write(3, chunk)',
-        // One drain loop: a single os.write past the pipe buffer returns short,
-        // and a truncated tail would change which bytes cross the ceiling.
-        'view = memoryview(b"\\n" + b"B" * (8 * 1024 * 1024))',
-        'while view:',
-        '    view = view[os.write(3, view):]',
         'return "never"',
       ].join('\n'),
       bindings: [],
@@ -4885,4 +4875,30 @@ describe('PythonCodeRuntime — hostile peer', () => {
     expect(result.error?.message).toContain(`protocol frame exceeded ${ceiling} bytes`)
   }, 120_000)
 
+  it('keeps two within-cap frames whose combined buffer crosses the cap', async () => {
+    // The unframed byte counter charges the WHOLE buffer, which legitimately
+    // holds several frames each within FRAME_PARSE_CAP_BYTES. A first frame of
+    // exactly the cap followed by a second frame crosses the counter without
+    // either frame exceeding the cap; the first-frame check (not the counter)
+    // must let them through, or a legitimate near-cap frame plus a trailing
+    // frame would be misreported as a worker-exit.
+    const { runtime } = await setup({ maxWallMs: 60_000, addressSpaceMb: 2048 })
+    const result = await runtime.run({
+      program: [
+        'import os',
+        'chunk = b"A" * (8 * 1024 * 1024)',
+        // Exactly the cap, no newline — at the limit, so nothing trips.
+        'for _ in range(8):',
+        '    os.write(3, chunk)',
+        // A newline, then a small legitimate log frame.
+        'os.write(3, b"\\n{\\"type\\":\\"log\\",\\"text\\":\\"after-cap-frames\\"}\\n")',
+        'return "done"',
+      ].join('\n'),
+      bindings: [],
+    })
+    expect(result.error).toBeUndefined()
+    expect(result.value).toBe('done')
+    expect(result.logs).toContain('after-cap-frames')
+  }, 120_000)
+
 })