diff --git a/packages/code-runtime/code-runtime-worker/src/bootstrap.ts b/packages/code-runtime/code-runtime-worker/src/bootstrap.ts index 62374ccaa0..f2e0d343f3 100644 --- a/packages/code-runtime/code-runtime-worker/src/bootstrap.ts +++ b/packages/code-runtime/code-runtime-worker/src/bootstrap.ts @@ -98,6 +98,10 @@ export function makeConsoleShim(logs: LogBuffer): Record<(typeof CONSOLE_LEVELS) * Redirect a stream's `write` into the log buffer (the program-visible * `process.stdout`/`process.stderr` in the real worker), so raw writes land * in emission order alongside console output instead of racing down a pipe. + * The shim keeps Node's `write(chunk[, encoding][, callback])` contract: the + * callback fires asynchronously once the chunk is admitted (a program + * awaiting flush completion must complete, not sit until the wall timeout), + * even for writes the exhausted budget drops. * @param logs - the buffer captured writes are pushed into. * @param stream - the stream whose `write` slot is patched. * @param source - the log source the captured writes are attributed to. @@ -109,8 +113,14 @@ export function captureStreamWrites(logs: LogBuffer, stream: PatchableStream, so // detached, so the unbound-method concern does not apply. // eslint-disable-next-line @typescript-eslint/unbound-method const original = stream.write - stream.write = (chunk: unknown): boolean => { + stream.write = (chunk: unknown, ...rest: unknown[]): boolean => { logs.push({ source, text: typeof chunk === 'string' ? chunk : String(chunk) }) + // Node's optional-encoding shape: the callback is whichever of the next + // two positions holds a function (a non-function there is the encoding). + const callback = [rest[0], rest[1]].find( + (arg): arg is (error?: Error | null) => void => typeof arg === 'function', + ) + if (callback) queueMicrotask(() => { callback(null) }) return true } return () => { stream.write = original } @@ -119,6 +129,28 @@ export function captureStreamWrites(logs: LogBuffer, stream: PatchableStream, so /** Bounded inspect options: deep enough to be useful, bounded so a pathological value cannot explode the rendering. */ const INSPECT_OPTIONS = { depth: 4, maxArrayLength: 100, maxStringLength: 10_000 } as const +/** + * The longest prefix of `text` whose UTF-8 encoding fits `maxBytes`, cut at + * a code-point boundary (never mid-surrogate-pair). The byte caps are BYTE + * caps — `String.prototype.slice` counts UTF-16 code units, up to 3× smaller + * than what a multibyte string actually costs across the boundary. + * @param text - the string to bound. + * @param maxBytes - the UTF-8 byte budget the prefix must fit. + * @returns the prefix (all of `text` when it already fits). + */ +export function truncateUtf8Bytes(text: string, maxBytes: number): string { + if (Buffer.byteLength(text, 'utf8') <= maxBytes) return text + let bytes = 0 + let end = 0 + for (const char of text) { + const cost = Buffer.byteLength(char, 'utf8') + if (bytes + cost > maxBytes) break + bytes += cost + end += char.length + } + return text.slice(0, end) +} + /** * Prepare the program's completion value for the done message: a value whose * MEASURED cross-boundary size fits `maxValueBytes` crosses raw — exact @@ -126,9 +158,10 @@ const INSPECT_OPTIONS = { depth: 4, maxArrayLength: 100, maxStringLength: 10_000 * everything else, so a huge container whose BOUNDED inspect rendering * happens to be small cannot smuggle itself past the cap. Anything else * (non-cloneable, or oversized) is REPLACED by its bounded `util.inspect` - * rendering, truncated with an in-band marker — the seam contract's "a - * non-transferable value is replaced by a string rendering", extended to - * oversized ones so a huge return cannot flood the host. + * rendering, byte-truncated ({@link truncateUtf8Bytes}) with an in-band + * marker — the seam contract's "a non-transferable value is replaced by a + * string rendering", extended to oversized ones so a huge return cannot + * flood the host. * @param value - the program's completion value. * @param maxValueBytes - the byte cap for the value. * @returns the done-message fragment: `{}` for `undefined`, else `{ value }`. @@ -150,7 +183,9 @@ export function prepareValue(value: unknown, maxValueBytes: number): { value?: u if (size !== undefined && size <= maxValueBytes) return { value } } const rendered = typeof value === 'string' ? value : inspect(value, INSPECT_OPTIONS) - const capped = rendered.length > maxValueBytes ? `${rendered.slice(0, maxValueBytes)}… [truncated]` : rendered + const capped = Buffer.byteLength(rendered, 'utf8') > maxValueBytes + ? `${truncateUtf8Bytes(rendered, maxValueBytes)}… [truncated]` + : rendered return { value: capped } } diff --git a/packages/code-runtime/code-runtime-worker/src/index.ts b/packages/code-runtime/code-runtime-worker/src/index.ts index 7749fb0318..f78f06cb0b 100644 --- a/packages/code-runtime/code-runtime-worker/src/index.ts +++ b/packages/code-runtime/code-runtime-worker/src/index.ts @@ -18,7 +18,7 @@ import { Context } from 'cordis' import z from 'schemastery' import { CodeRuntime } from '@deepseek-ai/dsh-code-runtime' import type { CodeBindingFunction, CodeLogEntry, CodeRunFailure, CodeRunRequest, CodeRunResult } from '@deepseek-ai/dsh-code-runtime' -import { prepareValue } from './bootstrap.ts' +import { prepareValue, truncateUtf8Bytes } from './bootstrap.ts' import { logTruncationMarker } from './protocol.ts' import type { ReplyMessage, WorkerBootData, WorkerToHost } from './protocol.ts' @@ -166,9 +166,8 @@ function parseWorkerMessage(raw: unknown): WorkerToHost | undefined { /** * Headroom the host's value re-cap grants over `maxValueBytes`: exactly the * truncation suffix {@link prepareValue} appends, so a value the WORKER - * already capped passes through unchanged instead of being marked twice. - * (A multibyte rendering the worker sliced by characters can still exceed - * this and pick up a second marker — bounded and harmless.) + * already capped (byte-exact prefix + this marker) passes through unchanged + * instead of being marked twice. */ const VALUE_RENDER_SLACK = Buffer.byteLength('… [truncated]', 'utf8') @@ -357,7 +356,7 @@ export class WorkerCodeRuntime extends CodeRuntime { // unchanged (see VALUE_RENDER_SLACK); the error text is bounded too. finish({ ...prepareValue(message.value, this.config.maxValueBytes + VALUE_RENDER_SLACK), - ...message.error ? { error: { kind: 'exception' as const, message: message.error.message.slice(0, this.config.maxValueBytes) } } : {}, + ...message.error ? { error: { kind: 'exception' as const, message: truncateUtf8Bytes(message.error.message, this.config.maxValueBytes) } } : {}, }) } diff --git a/packages/code-runtime/code-runtime-worker/tests/bootstrap.spec.ts b/packages/code-runtime/code-runtime-worker/tests/bootstrap.spec.ts index 685c152b10..e41f4455bb 100644 --- a/packages/code-runtime/code-runtime-worker/tests/bootstrap.spec.ts +++ b/packages/code-runtime/code-runtime-worker/tests/bootstrap.spec.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from 'vitest' import { EventEmitter } from 'node:events' -import { LogBuffer, makeConsoleShim, makeNamespaces, captureStreamWrites, prepareValue, runWorkerMain, wireReplies } from '@deepseek-ai/dsh-code-runtime-worker/src/bootstrap.ts' +import { LogBuffer, makeConsoleShim, makeNamespaces, captureStreamWrites, prepareValue, runWorkerMain, truncateUtf8Bytes, wireReplies } from '@deepseek-ai/dsh-code-runtime-worker/src/bootstrap.ts' import type { BootstrapPort, PatchableStream, PendingCall } from '@deepseek-ai/dsh-code-runtime-worker/src/bootstrap.ts' import type { ReplyMessage, WorkerToHost } from '@deepseek-ai/dsh-code-runtime-worker/src/protocol.ts' import type { CodeLogEntry } from '@deepseek-ai/dsh-code-runtime' @@ -90,6 +90,27 @@ describe('captureStreamWrites', () => { expect(seen[0]).toMatchObject({ source: 'stdout' }) expect(underlying).toBe('after') }) + + it('invokes the write callback asynchronously, in both optional-encoding shapes', async () => { + const buffer = new LogBuffer(1_000, () => {}) + const stream: PatchableStream = { write: () => true } + captureStreamWrites(buffer, stream, 'stdout') + const calls: (Error | null | undefined)[] = [] + stream.write('two-arg', (error?: Error | null) => calls.push(error)) + stream.write('three-arg', 'utf8', (error?: Error | null) => calls.push(error)) + // Node's contract: the callback fires after the write call returns. + expect(calls).toEqual([]) + await new Promise(resolve => stream.write('awaited flush', resolve)) + expect(calls).toEqual([null, null]) + }) + + it('still fires the callback for a write the exhausted budget drops', async () => { + const buffer = new LogBuffer(4, () => {}) + const stream: PatchableStream = { write: () => true } + captureStreamWrites(buffer, stream, 'stdout') + stream.write('this write overflows the budget and is dropped') + await new Promise(resolve => stream.write('also dropped', resolve)) + }) }) describe('prepareValue', () => { @@ -118,6 +139,34 @@ describe('prepareValue', () => { expect(typeof value).toBe('string') expect(value).toContain('more items') }) + + it('caps a multibyte string by UTF-8 bytes, not UTF-16 length', () => { + // 4 code units but 12 UTF-8 bytes: a length-counting cap would pass the + // full string through untruncated. + expect(prepareValue('€€€€', 4)).toEqual({ value: '€… [truncated]' }) + }) + + it('caps a multibyte rendering by UTF-8 bytes too', () => { + // Wire size (24-byte string inside an array) exceeds the cap, so the + // value crosses as its rendering — whose truncation must also be + // byte-exact: "[ '" (3 bytes) + two € (6 bytes) = 9; a third € would + // overflow the 10-byte budget. + expect(prepareValue(['€€€€€€€€'], 10)).toEqual({ value: "[ '€€… [truncated]" }) + }) +}) + +describe('truncateUtf8Bytes', () => { + it('returns a fitting string whole', () => { + expect(truncateUtf8Bytes('fits', 4)).toBe('fits') + }) + + it('cuts at a code-point boundary, never mid-surrogate-pair', () => { + // Each 😀 is one code point, two code units, four UTF-8 bytes: a 5-byte + // budget fits exactly one — and never leaves a lone surrogate behind. + const cut = truncateUtf8Bytes('😀😀', 5) + expect(cut).toBe('😀') + expect(Buffer.byteLength(truncateUtf8Bytes('😀😀', 3), 'utf8')).toBe(0) + }) }) describe('makeNamespaces', () => { diff --git a/packages/code-runtime/code-runtime-worker/tests/runtime.spec.ts b/packages/code-runtime/code-runtime-worker/tests/runtime.spec.ts index 2d29dc143d..edc2bd1271 100644 --- a/packages/code-runtime/code-runtime-worker/tests/runtime.spec.ts +++ b/packages/code-runtime/code-runtime-worker/tests/runtime.spec.ts @@ -221,6 +221,29 @@ describe('WorkerCodeRuntime — budgets and containment (real workers)', () => { expect(result.value).toBe(`${'y'.repeat(64)}… [truncated]`) }) + it('caps a multibyte return value by UTF-8 bytes, not string length', async () => { + // 4 code units, 12 UTF-8 bytes: a length-counting cap would let the full + // string cross. The worker's byte-exact capped rendering then passes the + // host re-cap unchanged (cap + marker is exactly the granted slack). + const { runtime } = await setup({ maxValueBytes: 4 }) + const result = await runtime.run({ program: 'return "€€€€"', bindings: [] }) + expect(result.value).toBe('€… [truncated]') + }) + + it('completes a program that awaits its write callback, capturing the chunk', async () => { + // Node's write(chunk[, encoding][, callback]) contract: dropping the + // callback would leave this promise pending until the wall ceiling and + // misreport a completed program as a timeout. + const { runtime } = await setup({ maxWallMs: 2_000 }) + const result = await runtime.run({ + program: 'await new Promise(resolve => process.stdout.write("flushed", resolve)); return "done"', + bindings: [], + }) + expect(result.error).toBeUndefined() + expect(result.value).toBe('done') + expect(result.logs).toContainEqual({ source: 'stdout', text: 'flushed' }) + }) + it('caps a huge container whose bounded rendering is small (wire size, not rendering, is what counts)', async () => { const { runtime } = await setup() const result = await runtime.run({ program: 'return new Array(50_000).fill(7)', bindings: [] }) @@ -340,6 +363,21 @@ describe('WorkerCodeRuntime — hostile programs (real workers)', () => { expect(result.error).toEqual({ kind: 'exception', message: 'fake failure' }) }) + it('byte-bounds forged multibyte error text at the host', async () => { + // Forged error text bypasses the worker entirely; the host bound is a + // BYTE bound (two € = 6 bytes fit an 8-byte cap, a third would not). + const { runtime } = await setup({ maxValueBytes: 8 }) + const result = await runtime.run({ + program: ` + const { parentPort } = await import('node:worker_threads'); + parentPort.postMessage({ type: 'done', error: { message: '€'.repeat(1000) } }); + for (;;) {} + `, + bindings: [], + }) + expect(result.error).toEqual({ kind: 'exception', message: '€€' }) + }) + it('answers a binding whose resolution cannot be cloned with a failure reply', async () => { const { runtime } = await setup() const result = await runtime.run({