diff --git a/packages/code-runtime/code-runtime-python/src/protocol.ts b/packages/code-runtime/code-runtime-python/src/protocol.ts index 29c790eaec..10d13a211c 100644 --- a/packages/code-runtime/code-runtime-python/src/protocol.ts +++ b/packages/code-runtime/code-runtime-python/src/protocol.ts @@ -202,7 +202,10 @@ function scalarJson(current: unknown): string { * number (non-finite, negative zero) is caught only when the value fits the * budget — an over-budget value is rejected regardless, so the distinction is * moot. Same JSON-plain precondition and traversal shape as - * {@link encodeJsonPlain}; per-scalar encoding delegates to `JSON.stringify`. + * {@link encodeJsonPlain}; per-scalar byte length is measured through + * {@link scalarJson} (matching the encoder, so a beyond-safe-range integer + * meters its exact BigInt digits, not `JSON.stringify`'s rounded spelling) and + * `JSON.stringify` for strings. * @param value - a JSON-plain value (e.g. straight from `JSON.parse`). * @param maxBytes - the completion-value budget in bytes. * @returns `{ ok: true, bytes }` with the exact serialized size, or @@ -235,15 +238,22 @@ export function checkDoneValue(value: unknown, maxBytes: number): { ok: true; by for (const item of current) stack.push(item) } else if (typeof current === 'object' && current !== null) { const record = current as Record - // Count own keys WITHOUT Object.entries/Object.keys: either would - // allocate one slot (entries: one pair array) per member before the - // bound below could run, recreating the spike the bound exists to stop. + // Count own keys WITHOUT Object.entries/Object.keys (either allocates one + // slot per member up front), AND bail mid-count the instant the minimum + // encoding exceeds the budget: braces (+2), each entry a quoted key + // (>= 2 bytes) + colon + >= 1-byte value (>= 4 bytes), and a comma per + // gap. A forged wide object with millions of keys and a small cap must + // fail in O(cap), not walk its whole breadth first. `bytes` still holds + // the pre-object total throughout this loop. let count = 0 - for (const key in record) if (Object.hasOwn(record, key)) count += 1 + for (const key in record) { + if (!Object.hasOwn(record, key)) continue + count += 1 + if (bytes + 2 + count * 4 + (count - 1) > maxBytes) return { ok: false, reason: 'over-budget' } + } + // The loop's final iteration already proved the whole object's lower + // bound fits, so no separate post-count check is needed here. bytes += 2 + (count > 1 ? count - 1 : 0) - // Same pre-enqueue bound: each entry contributes its quoted key (>= 2 - // bytes), the colon, and a >= 1-byte value. - if (bytes + count * 4 > maxBytes) return { ok: false, reason: 'over-budget' } for (const key in record) { if (!Object.hasOwn(record, key)) continue // The same string lower bound, before escaping the key. diff --git a/packages/code-runtime/code-runtime-python/tests/protocol.spec.ts b/packages/code-runtime/code-runtime-python/tests/protocol.spec.ts index dc0a01d47d..f0e06af723 100644 --- a/packages/code-runtime/code-runtime-python/tests/protocol.spec.ts +++ b/packages/code-runtime/code-runtime-python/tests/protocol.spec.ts @@ -207,6 +207,21 @@ describe('checkDoneValue', () => { const wide: Record = {} for (let i = 0; i < 10; i++) wide[`k${i}`] = i expect(checkDoneValue(wide, 12)).toEqual({ ok: false, reason: 'over-budget' }) + // A forged object with millions of keys and a small cap must reject in + // O(cap): the key COUNT loop itself bails once the running minimum encoding + // (braces + 4 bytes/entry + commas) crosses the budget, rather than walking + // the whole breadth before checking. Observable as a bounded key subset: + // build a Proxy whose ownKeys would yield far more than the cap admits and + // assert the metered walk never enumerates past it. + let enumerated = 0 + const millionKeys = new Proxy({}, { + ownKeys() { return Array.from({ length: 2_000_000 }, (_unused, i) => `k${i}`) }, + getOwnPropertyDescriptor() { enumerated += 1; return { enumerable: true, configurable: true, value: 0 } }, + }) + expect(checkDoneValue(millionKeys, 64)).toEqual({ ok: false, reason: 'over-budget' }) + // With cap 64, at most ~16 entries (4 bytes each) can fit before the bound + // trips, so the walk enumerates far fewer than the 2,000,000 declared keys. + expect(enumerated).toBeLessThan(1000) }) it('rejects an over-budget string on its length before escaping it', () => {