Merge branch 'code-runtime-worker' into code-mode-tools

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 }
 }
 

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) } } : {},
         })
       }
 

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<void>(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<void>(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', () => {

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({