fix(code-runtime): drain late worker pipe output

2026-07-22 01:27:03 +08:00
parent e1c5832d21
commit 865d7de858
6 changed files with 82 additions and 25 deletions
@@ -23,7 +23,7 @@ Every field is validated and defaulted; `maxOutputBytes` is a safe integer of at
 - **The port assumes a hostile peer** — model code can reach `parentPort` and forge traffic, so every inbound message is shape-validated and REBUILT before anything reads it (`null`, primitives, junk types, and malformed payloads drop without a throw; forged extra fields never ride along), the host answers each call id at most once, resolves binding names as OWN properties only (a forged `constructor` cannot walk a prototype chain), drops post-settlement replies, and validates every binding resolution and completion as lossless JSON. Forged `log`/`done` messages cannot bypass the outer cap: the host repeats validation and accounts every admitted log plus the completion or diagnostic. Worker-side namespaces are null-prototype with `defineProperty`, so `__proto__`-shaped binding names are ordinary keys.
 - **Two independent budgets, because the peer is hostile** — `computeMs` meters the worker's MEASURED busy time (`worker.performance.eventLoopUtilization()` polling): a hot loop cannot hide behind a pending decoy dispatch, and a program awaiting a slow tool accrues nothing. `maxWallMs` backstops what busy time cannot see (awaiting a promise nobody resolves). Both funnel into `worker.terminate()`, which ends hot synchronous loops too; heap overflow surfaces as the worker's OOM exit (`kind: 'worker-exit'`).
 - **Intermediate binding values are complete JSON** — binding arguments and resolutions cross by structured clone after lossless-JSON validation and have no byte cap. They never enter the outer-output ledger or model context; provider/executor acquisition bounds and process/worker memory remain the limits.
- **Logs stream eagerly into one outer ledger** — console/stdout/stderr text crosses the port in emission order, so a timed-out or killed program still shows what it printed. `maxOutputBytes` accounts the JSON serialization of the outer `logs` array plus the completion value or failure diagnostic. At or below the cap the exact value returns; a lossy completion is `invalid-output`, and a combined overflow is `output-limit` rather than a substituted inspected string. The failure retains the fitting captured prefix and later follows the normal outer `run_code` spill policy.
+- **Logs stream eagerly into one outer ledger** — console/stdout/stderr text crosses the port in emission order, so a timed-out or killed program still shows what it printed. Native writes that bypass the patched stream slots arrive on pipes independent of the completion port; settlement therefore continues bounded pipe capture until worker termination completes before materializing the result. `maxOutputBytes` accounts the JSON serialization of the outer `logs` array plus the completion value or failure diagnostic. At or below the cap the exact value returns; a lossy completion is `invalid-output`, and a combined overflow is `output-limit` rather than a substituted inspected string. The failure retains the fitting captured prefix and later follows the normal outer `run_code` spill policy.
 - **Empty environment** — the worker gets `env: {}` and `execArgv: []`: no ambient credentials (stronger than the scrubbed-env rule for spawned commands) and no inherited loader flags.
 - **Dispose to quiescence** — teardown fails in-flight runs as `abort` and AWAITS each worker's exit before resolving.

@@ -8,6 +8,7 @@

 import { Worker } from 'node:worker_threads'
 import { stripTypeScriptTypes } from 'node:module'
+import type { Readable } from 'node:stream'
 import { fileURLToPath } from 'node:url'
 import { Context } from 'cordis'
 import z from 'schemastery'
@@ -106,6 +107,25 @@ function messageOf(error: unknown): string {
  return error instanceof Error ? error.message : String(error)
 }

+/** Resolve after a worker pipe emits all queued data, or closes/errors during termination. */
+function waitForPipeDrain(stream: Readable): Promise<void> {
+  if (stream.readableEnded || stream.destroyed) return Promise.resolve()
+  return new Promise((resolve) => {
+    const done = (): void => {
+      stream.off('end', done)
+      stream.off('close', done)
+      stream.off('error', done)
+      resolve()
+    }
+    stream.once('end', done)
+    stream.once('close', done)
+    stream.once('error', done)
+    // Close the event-registration race if termination finished between the
+    // initial state check and the listeners above.
+    if (stream.readableEnded || stream.destroyed) done()
+  })
+}
+
 /**
 * Runtime shape gate for inbound port traffic. The peer runs MODEL CODE and
 * can post anything — `null`, primitives, objects with poisoned fields — so
@@ -335,13 +355,20 @@ export class WorkerCodeRuntime extends CodeRuntime {
      const logs: string[] = []
      const strayLogs: string[] = []
      const output = new OutputLedger(this.config.maxOutputBytes)
+      let terminalOverride: CodeRunResult | undefined

-      // No settled guard: `finish` snapshots the arrays when it resolves, so
-      // a chunk flushing after settlement mutates only the discarded buffers,
-      // and the ledger bounds that growth until the pipes close.
+      // Pipe and message-port delivery are independent. Continue bounded pipe
+      // capture after a terminal message while worker termination drains bytes
+      // that were already queued; `finish` materializes the result only after
+      // termination completes.
      const captureStray = (chunk: Buffer): void => {
+        if (terminalOverride !== undefined) return
        const text = chunk.toString('utf8')
-        if (!settled && !output.admit(text, strayLogs)) finish(output.limit([...logs, ...strayLogs, text]))
+        if (!output.admit(text, strayLogs)) {
+          const limited = output.limit([...logs, ...strayLogs, text])
+          terminalOverride = limited
+          finish(() => limited)
+        }
      }
      worker.stdout.on('data', captureStray)
      worker.stderr.on('data', captureStray)
@@ -350,14 +377,20 @@ export class WorkerCodeRuntime extends CodeRuntime {
      // logs captured before timeout, abort, or failure remain in the result.
      let finishResolve!: () => void
      const finished = new Promise<void>((done) => { finishResolve = done })
-      const finish = (result: CodeRunResult): void => {
+      const finish = (finalize: () => CodeRunResult): void => {
        if (settled) return
        settled = true
        clearInterval(eluTimer)
        clearTimeout(wallTimer)
        request.signal?.removeEventListener('abort', onAbort)
        this.live.delete(live)
-        void worker.terminate().then(() => {
+        // Let the poll phase deliver pipe bytes already queued independently
+        // of the terminal port message before termination closes the streams.
+        void new Promise<void>((resume) => { setImmediate(resume) }).then(async () => {
+          const stdoutDrained = waitForPipeDrain(worker.stdout)
+          const stderrDrained = waitForPipeDrain(worker.stderr)
+          await Promise.all([worker.terminate(), stdoutDrained, stderrDrained])
+          const result = terminalOverride ?? finalize()
          finishResolve()
          resolve(result)
        })
@@ -365,22 +398,24 @@ export class WorkerCodeRuntime extends CodeRuntime {

      const onDone = (message: WorkerToHost): void => {
        if (message.type !== 'done') return
-        const captured = [...logs, ...strayLogs]
        if (message.error) {
-          finish(output.failure(captured, message.error))
+          const error = message.error
+          finish(() => output.failure([...logs, ...strayLogs], error))
          return
        }
        if (message.value === undefined) {
-          finish(output.success(captured))
+          finish(() => output.success([...logs, ...strayLogs]))
          return
        }
        // The worker-thread boundary has already structured-cloned this
        // hostile value, so accessors and proxies cannot survive to throw
        // during the lossless-JSON snapshot.
        const value = snapshotJsonValue(message.value) as CodeJsonValue | undefined
-        finish(value === undefined
-          ? output.failure(captured, { kind: 'invalid-output', message: 'program completion must be lossless JSON' })
-          : output.success(captured, value))
+        if (value === undefined) {
+          finish(() => output.failure([...logs, ...strayLogs], { kind: 'invalid-output', message: 'program completion must be lossless JSON' }))
+        } else {
+          finish(() => output.success([...logs, ...strayLogs], value))
+        }
      }

      const onCall = (message: WorkerToHost): void => {
@@ -438,21 +473,25 @@ export class WorkerCodeRuntime extends CodeRuntime {
        const message = parseWorkerMessage(raw)
        if (!message) return
        if (message.type === 'log' && !settled && !output.admit(message.text, logs)) {
-          finish(output.limit([...logs, ...strayLogs, message.text]))
+          const limited = output.limit([...logs, ...strayLogs, message.text])
+          terminalOverride = limited
+          finish(() => limited)
          return
        }
        if (message.type === 'output-limit' && !settled) {
-          finish(output.limit([...logs, ...strayLogs]))
+          const limited = output.limit([...logs, ...strayLogs])
+          terminalOverride = limited
+          finish(() => limited)
          return
        }
        onCall(message)
        onDone(message)
      })
      worker.on('error', (error: Error) => {
-        finish(output.failure([...logs, ...strayLogs], { kind: 'worker-exit', message: `worker error: ${error.message}` }))
+        finish(() => output.failure([...logs, ...strayLogs], { kind: 'worker-exit', message: `worker error: ${error.message}` }))
      })
      worker.on('exit', (exitCode: number) => {
-        finish(output.failure([...logs, ...strayLogs], { kind: 'worker-exit', message: `worker exited with code ${exitCode} before completing` }))
+        finish(() => output.failure([...logs, ...strayLogs], { kind: 'worker-exit', message: `worker exited with code ${exitCode} before completing` }))
      })

      // The compute budget reads the worker's own measured busy time, so a
@@ -461,21 +500,21 @@ export class WorkerCodeRuntime extends CodeRuntime {
      const eluTimer = setInterval(() => {
        const elu = worker.performance.eventLoopUtilization()
        if (elu.active > this.config.computeMs) {
-          finish(output.failure([...logs, ...strayLogs], { kind: 'timeout', message: `compute budget exhausted (${this.config.computeMs}ms busy)` }))
+          finish(() => output.failure([...logs, ...strayLogs], { kind: 'timeout', message: `compute budget exhausted (${this.config.computeMs}ms busy)` }))
        }
      }, ELU_POLL_INTERVAL_MS)
      const wallTimer = setTimeout(() => {
-        finish(output.failure([...logs, ...strayLogs], { kind: 'timeout', message: `wall-clock ceiling reached (${this.config.maxWallMs}ms)` }))
+        finish(() => output.failure([...logs, ...strayLogs], { kind: 'timeout', message: `wall-clock ceiling reached (${this.config.maxWallMs}ms)` }))
      }, this.config.maxWallMs)
      const onAbort = (): void => {
-        finish(output.failure([...logs, ...strayLogs], { kind: 'abort', message: String(request.signal?.reason) }))
+        finish(() => output.failure([...logs, ...strayLogs], { kind: 'abort', message: String(request.signal?.reason) }))
      }
      request.signal?.addEventListener('abort', onAbort, { once: true })

      const live: LiveRun = {
        worker,
        finished,
-        settle: (failure: CodeRunFailure) => { finish(output.failure([...logs, ...strayLogs], failure)) },
+        settle: (failure: CodeRunFailure) => { finish(() => output.failure([...logs, ...strayLogs], failure)) },
      }
      this.live.add(live)
    })
@@ -333,6 +333,24 @@ describe('WorkerCodeRuntime — budgets and containment (real workers)', () => {
    expect(result.logs[1]?.length).toBeGreaterThan(0)
    expect('b'.repeat(100).startsWith(result.logs[1] ?? '')).toBe(true)
  }, 15_000)
+
+  it('drains pipe output queued before terminal worker teardown completes', async () => {
+    const { runtime } = await setup({ maxOutputBytes: 200_000 })
+    const payload = `late-pipe-${'x'.repeat(100_000)}`
+    const result = await runtime.run({
+      program: `
+        const { parentPort } = await import('node:worker_threads');
+        const write = (text) => Object.getPrototypeOf(process.stdout).write.call(process.stdout, text);
+        write('late-pipe-' + 'x'.repeat(100_000));
+        parentPort.postMessage({ type: 'done', value: 'done' });
+        for (;;) {}
+      `,
+      bindings: [],
+    })
+    expect(result.error).toBeUndefined()
+    expect(result.value).toBe('done')
+    expect(result.logs.join('') === payload).toBe(true)
+  }, 15_000)
 })

 describe('WorkerCodeRuntime — hostile programs (real workers)', () => {