ソースを参照

fix(code-runtime): validate arguments before worker dispatch

Tianyi Cui 2 ヶ月 前
コミット
2623ddbee4

+ 31 - 15
packages/code-runtime/code-runtime-worker/src/bootstrap.ts

@@ -188,6 +188,12 @@ export class ToolCallError extends Error {
   }
 }
 
+/** Create the namespace-specific rejection for one lossy binding argument. */
+function bindingArgumentFailure(global: string, name: string): Error {
+  const message = 'binding arguments must be lossless JSON'
+  return global === 'tools' ? new ToolCallError(name, message) : new Error(message)
+}
+
 /**
  * Route host replies into the pending-call map: each reply settles its call
  * at most once, and a reply for an unknown id (stray, or a duplicate answer
@@ -211,7 +217,8 @@ export function wireReplies(port: BootstrapPort, pending: Map<number, PendingCal
  * Build the binding namespace objects the program sees: one null-prototype global per
  * namespace, each declared name an own enumerable async function that bridges over the port
  * (`__proto__`/`constructor`/`toString` are ordinary keys, never prototype collisions).
- * Non-cloneable arguments and host failure replies reject only the corresponding call.
+ * Lossy arguments reject before posting; clone failures and host failure
+ * replies reject only the corresponding call.
  *
  * @param data - the boot payload's namespace declarations (globals + names).
  * @param port - the port binding calls are posted to.
@@ -230,22 +237,31 @@ export function makeNamespaces(
     for (const name of names) {
       Object.defineProperty(namespace, name, {
         enumerable: true,
-        value: (args: unknown): Promise<unknown> => new Promise((resolve, reject) => {
-          const id = nextId.value++
-          pending.set(id, {
-            resolve,
-            reject: (error) => {
-              reject(global === 'tools' ? new ToolCallError(name, error.message) : error)
-            },
-          })
+        value: (args: unknown): Promise<unknown> => {
+          let detached: unknown
           try {
-            port.postMessage({ type: 'call', id, global, name, args })
-          } catch (error: unknown) {
-            pending.delete(id)
-            const message = `binding arguments must be structured-cloneable: ${error instanceof Error ? error.message : String(error)}`
-            reject(global === 'tools' ? new ToolCallError(name, message) : new Error(message))
+            detached = snapshotCodeJsonValue(args)
+          } catch {
+            detached = undefined
           }
-        }),
+          if (detached === undefined) return Promise.reject(bindingArgumentFailure(global, name))
+          return new Promise((resolve, reject) => {
+            const id = nextId.value++
+            pending.set(id, {
+              resolve,
+              reject: (error) => {
+                reject(global === 'tools' ? new ToolCallError(name, error.message) : error)
+              },
+            })
+            try {
+              port.postMessage({ type: 'call', id, global, name, args: detached })
+            } catch (error: unknown) {
+              pending.delete(id)
+              const message = `binding arguments must be structured-cloneable: ${error instanceof Error ? error.message : String(error)}`
+              reject(global === 'tools' ? new ToolCallError(name, message) : new Error(message))
+            }
+          })
+        },
       })
     }
     return namespace

+ 23 - 8
packages/code-runtime/code-runtime-worker/src/index.ts

@@ -15,6 +15,7 @@ import { CodeRuntime } from '@deepseek-ai/dsh-code-runtime'
 import type { CodeBindingFunction, CodeJsonValue, CodeRunFailure, CodeRunRequest, CodeRunResult } from '@deepseek-ai/dsh-code-runtime'
 import { snapshotJsonValue } from '@deepseek-ai/dsh-session'
 import type { ReplyMessage, WorkerBootData, WorkerToHost } from './protocol.ts'
+import { truncateJsonStringBytes } from './output-json.ts'
 
 /** Plugin config: every execution cap, changeable from `cordis.yml` (no hardcoded tunables). */
 export interface Config {
@@ -174,16 +175,29 @@ class OutputLedger {
     return { logs, error }
   }
 
-  /** Build the explicit output-limit failure while retaining the fitting log prefix. */
+  /** Build the explicit output-limit failure while retaining a fitting prefix of the final log. */
   limit(logs: string[]): CodeRunResult {
     const fullMessage = `outer output exceeded ${this.maxBytes} bytes`
-    let retainedBytes = this.bytes
     const messageBytes = Buffer.byteLength(JSON.stringify(fullMessage), 'utf8')
-    while (logs.length > 0 && retainedBytes + messageBytes > this.maxBytes) {
-      const removed = logs.pop()
+    const retained = [...logs]
+    let retainedBytes = jsonBytes(retained)
+    const logBudget = this.maxBytes - messageBytes
+    while (retained.length > 0 && retainedBytes > logBudget) {
+      const removed = retained.pop()
       /* v8 ignore next -- the while guard proves pop cannot return undefined. */
       if (removed === undefined) throw new Error('output ledger lost its final log entry')
-      retainedBytes -= Buffer.byteLength(JSON.stringify(removed), 'utf8') + (logs.length > 0 ? 1 : 0)
+      const separatorBytes = retained.length > 0 ? 1 : 0
+      retainedBytes -= Buffer.byteLength(JSON.stringify(removed), 'utf8') + separatorBytes
+      const prefix = truncateJsonStringBytes(removed, logBudget - retainedBytes - separatorBytes)
+      if (prefix.length > 0) {
+        retained.push(prefix)
+        retainedBytes += Buffer.byteLength(JSON.stringify(prefix), 'utf8') + separatorBytes
+        break
+      }
+    }
+    if (logBudget < 2) {
+      retained.length = 0
+      retainedBytes = 2
     }
     const availableMessageBytes = this.maxBytes - retainedBytes
     // This fixed diagnostic is ASCII with no JSON escapes, so two bytes are
@@ -191,7 +205,7 @@ class OutputLedger {
     const message = messageBytes <= availableMessageBytes
       ? fullMessage
       : fullMessage.slice(0, availableMessageBytes - 2)
-    return { logs, error: { kind: 'output-limit', message } }
+    return { logs: retained, error: { kind: 'output-limit', message } }
   }
 }
 
@@ -326,7 +340,8 @@ export class WorkerCodeRuntime extends CodeRuntime {
       // a chunk flushing after settlement mutates only the discarded buffers,
       // and the ledger bounds that growth until the pipes close.
       const captureStray = (chunk: Buffer): void => {
-        if (!settled && !output.admit(chunk.toString('utf8'), strayLogs)) finish(output.limit([...logs, ...strayLogs]))
+        const text = chunk.toString('utf8')
+        if (!settled && !output.admit(text, strayLogs)) finish(output.limit([...logs, ...strayLogs, text]))
       }
       worker.stdout.on('data', captureStray)
       worker.stderr.on('data', captureStray)
@@ -416,7 +431,7 @@ 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]))
+          finish(output.limit([...logs, ...strayLogs, message.text]))
           return
         }
         if (message.type === 'output-limit' && !settled) {

+ 36 - 0
packages/code-runtime/code-runtime-worker/src/output-json.ts

@@ -0,0 +1,36 @@
+/** JSON string-prefix accounting for the outer-output ledger. @module @deepseek-ai/dsh-code-runtime-worker/output-json */
+
+/** Control characters with a two-byte short JSON escape instead of `\u00XX`. */
+const SHORT_ESCAPE_CODES = new Set([0x08, 0x09, 0x0a, 0x0c, 0x0d])
+
+/** Serialized bytes contributed by one complete Unicode code point inside JSON quotes. */
+function serializedCharacterBytes(character: string): number {
+  if (character.length === 2) return 4
+  if (character === '"' || character === '\\') return 2
+  const code = character.charCodeAt(0)
+  if (code >= 0xd800 && code <= 0xdfff) return 6
+  if (code < 0x20) return SHORT_ESCAPE_CODES.has(code) ? 2 : 6
+  return Buffer.byteLength(character, 'utf8')
+}
+
+/**
+ * Return the longest code-point-aligned prefix whose JSON string encoding,
+ * including its surrounding quotes, fits `maxBytes`.
+ *
+ * @param text - the candidate string.
+ * @param maxBytes - serialized JSON-string bytes available.
+ * @returns the fitting prefix, or an empty string when even useful content cannot fit.
+ */
+export function truncateJsonStringBytes(text: string, maxBytes: number): string {
+  if (maxBytes < 2) return ''
+  if (Buffer.byteLength(JSON.stringify(text), 'utf8') <= maxBytes) return text
+  let bytes = 2
+  let end = 0
+  for (const character of text) {
+    const cost = serializedCharacterBytes(character)
+    if (bytes + cost > maxBytes) break
+    bytes += cost
+    end += character.length
+  }
+  return text.slice(0, end)
+}

+ 1 - 0
packages/code-runtime/code-runtime-worker/src/worker-json.ts

@@ -33,6 +33,7 @@ export function snapshotCodeJsonValue(value: unknown): CodeJsonValue | undefined
 
     if (Array.isArray(candidate)) {
       if (Object.getPrototypeOf(candidate) !== Array.prototype) return undefined
+      if (Reflect.ownKeys(candidate).length !== candidate.length + 1) return undefined
       return within(candidate, () => {
         const result: CodeJsonValue[] = []
         for (let index = 0; index < candidate.length; index++) {

+ 35 - 4
packages/code-runtime/code-runtime-worker/tests/bootstrap.spec.ts

@@ -190,7 +190,7 @@ describe('makeNamespaces', () => {
     await expect(tools['toString']?.({})).resolves.toBe('toString-ok')
   })
 
-  it('rejects a non-cloneable argument without leaking the pending entry', async () => {
+  it('rejects a postMessage clone failure without leaking the pending entry', async () => {
     let firstCall = true
     const throwingPort: BootstrapPort = {
       // First call throws an Error (the real DataCloneError shape), the
@@ -203,8 +203,8 @@ describe('makeNamespaces', () => {
     }
     const pending = new Map<number, PendingCall>()
     const [tools] = makeNamespaces({ namespaces: [{ global: 'tools', names: ['x'] }] }, throwingPort, pending, { value: 1 }) as [Record<string, (args: unknown) => Promise<unknown>>]
-    const first = await rejectionOf(tools.x?.(() => 1) ?? Promise.resolve())
-    const second = await rejectionOf(tools.x?.(() => 1) ?? Promise.resolve())
+    const first = await rejectionOf(tools.x?.({ first: true }) ?? Promise.resolve())
+    const second = await rejectionOf(tools.x?.({ second: true }) ?? Promise.resolve())
     expect(first).toMatchObject({ name: 'ToolCallError', toolName: 'x' })
     expect(second).toMatchObject({ name: 'ToolCallError', toolName: 'x' })
     expect(first).toBeInstanceOf(ToolCallError)
@@ -214,6 +214,32 @@ describe('makeNamespaces', () => {
     expect(pending.size).toBe(0)
   })
 
+  it('rejects lossy arguments before posting or allocating a call id', async () => {
+    let posts = 0
+    const port: BootstrapPort = { postMessage: () => { posts += 1 }, on: () => {} }
+    const pending = new Map<number, PendingCall>()
+    const nextId = { value: 1 }
+    const [tools] = makeNamespaces(
+      { namespaces: [{ global: 'tools', names: ['x'] }] }, port, pending, nextId,
+    ) as [Record<string, (args: unknown) => Promise<unknown>>]
+    const decorated = [1]
+    Object.defineProperty(decorated, 'extra', { value: true })
+    const throwing = Object.defineProperty({}, 'value', {
+      enumerable: true,
+      get: () => { throw new Error('getter exploded') },
+    })
+
+    for (const value of [() => 1, new Date(), decorated, throwing]) {
+      const failure = await rejectionOf(tools.x?.(value) ?? Promise.resolve())
+      expect(failure).toMatchObject({
+        name: 'ToolCallError', toolName: 'x', message: 'binding arguments must be lossless JSON',
+      })
+    }
+    expect(posts).toBe(0)
+    expect(pending.size).toBe(0)
+    expect(nextId.value).toBe(1)
+  })
+
   it('uses ordinary Error for non-tools namespace failures', async () => {
     const deniedPort = new FakePort()
     deniedPort.respond = message => message.type === 'call'
@@ -226,9 +252,14 @@ describe('makeNamespaces', () => {
     expect(denied).toBeInstanceOf(Error)
     expect(denied).not.toBeInstanceOf(ToolCallError)
 
+    const invalid = await rejectionOf(helpers.x?.(() => 1) ?? Promise.resolve())
+    expect(invalid).toBeInstanceOf(Error)
+    expect(invalid).not.toBeInstanceOf(ToolCallError)
+    expect((invalid as Error).message).toBe('binding arguments must be lossless JSON')
+
     const clonePort: BootstrapPort = { postMessage: () => { throw new Error('clone failed') }, on: () => {} }
     const [cloneHelpers] = makeNamespaces({ namespaces: [{ global: 'helpers', names: ['x'] }] }, clonePort, new Map(), { value: 1 }) as [Record<string, (args: unknown) => Promise<unknown>>]
-    const cloneFailure = await rejectionOf(cloneHelpers.x?.(() => 1) ?? Promise.resolve())
+    const cloneFailure = await rejectionOf(cloneHelpers.x?.({}) ?? Promise.resolve())
     expect(cloneFailure).toBeInstanceOf(Error)
     expect(cloneFailure).not.toBeInstanceOf(ToolCallError)
   })

+ 18 - 0
packages/code-runtime/code-runtime-worker/tests/output-json.spec.ts

@@ -0,0 +1,18 @@
+import { describe, expect, it } from 'vitest'
+import { truncateJsonStringBytes } from '../src/output-json.ts'
+
+describe('truncateJsonStringBytes', () => {
+  it('returns a fitting string whole and rejects budgets without JSON quotes', () => {
+    expect(truncateJsonStringBytes('fits', 6)).toBe('fits')
+    expect(truncateJsonStringBytes('x', 1)).toBe('')
+  })
+
+  it('accounts every JSON escape and cuts only between complete code points', () => {
+    const prefix = '"\\\b\t\n\f\r\u0000😀\ud800€a'
+    const text = `${prefix}z`
+    const budget = Buffer.byteLength(JSON.stringify(prefix), 'utf8')
+
+    expect(truncateJsonStringBytes(text, budget)).toBe(prefix)
+    expect(Buffer.byteLength(JSON.stringify(truncateJsonStringBytes(text, budget)), 'utf8')).toBe(budget)
+  })
+})

+ 41 - 1
packages/code-runtime/code-runtime-worker/tests/runtime.spec.ts

@@ -218,6 +218,19 @@ describe('WorkerCodeRuntime — budgets and containment (real workers)', () => {
     expect(Buffer.byteLength(JSON.stringify(result.logs), 'utf8')).toBeLessThan(300)
   })
 
+  it('retains a fitting prefix when one oversized log is the first output', async () => {
+    const { runtime } = await setup({ maxOutputBytes: 96 })
+    const result = await runtime.run({
+      program: 'console.log(`start-${`😀"\\\\\\n`.repeat(100)}`); return null',
+      bindings: [],
+    })
+    expect(result.error).toEqual({ kind: 'output-limit', message: 'outer output exceeded 96 bytes' })
+    expect(result.logs).toHaveLength(1)
+    expect(result.logs[0]?.startsWith('start-')).toBe(true)
+    expect(Buffer.byteLength(JSON.stringify(result.logs), 'utf8')
+      + Buffer.byteLength(JSON.stringify(result.error?.message), 'utf8')).toBeLessThanOrEqual(96)
+  })
+
   it('fails an oversized return value without substituting a string', async () => {
     const { runtime } = await setup({ maxOutputBytes: 64 })
     const result = await runtime.run({ program: 'return "y".repeat(10_000)', bindings: [] })
@@ -304,7 +317,8 @@ describe('WorkerCodeRuntime — budgets and containment (real workers)', () => {
     })
     expect(result.error?.kind).toBe('output-limit')
     expect(result.logs).toContain('a'.repeat(20))
-    expect(result.logs).not.toContain('b'.repeat(100))
+    expect(result.logs[1]?.length).toBeGreaterThan(0)
+    expect('b'.repeat(100).startsWith(result.logs[1] ?? '')).toBe(true)
   }, 15_000)
 })
 
@@ -409,6 +423,32 @@ describe('WorkerCodeRuntime — hostile programs (real workers)', () => {
     expect(result.value).toEqual({ name: 'ToolCallError', toolName: 'bad', message: 'binding resolution must be lossless JSON' })
   })
 
+  it('rejects lossy binding arguments in the worker before invoking the host binding', async () => {
+    const { runtime } = await setup()
+    let calls = 0
+    const result = await runtime.run({
+      program: `
+        const decorated = [1]; Object.defineProperty(decorated, 'extra', { value: true });
+        const values = [new Date(), decorated, () => 1];
+        const failures = [];
+        for (const value of values) {
+          try { await tools.never(value) } catch (error) {
+            failures.push({ typed: error instanceof ToolCallError, name: error.name, toolName: error.toolName, message: error.message });
+          }
+        }
+        return failures;
+      `,
+      bindings: tools({ never: async () => { calls += 1; return null } }),
+    })
+    expect(calls).toBe(0)
+    expect(result.value).toEqual(new Array(3).fill({
+      typed: true,
+      name: 'ToolCallError',
+      toolName: 'never',
+      message: 'binding arguments must be lossless JSON',
+    }))
+  })
+
   it('contains throwing getters while snapshotting binding resolutions', async () => {
     const { runtime } = await setup()
     const result = await runtime.run({

+ 9 - 0
packages/code-runtime/code-runtime-worker/tests/worker-json.spec.ts

@@ -60,12 +60,21 @@ describe('snapshotCodeJsonValue', () => {
     class ExoticArray extends Array<number> {}
     const cyclic: Record<string, unknown> = {}
     cyclic.self = cyclic
+    const decorated = [1]
+    Object.defineProperty(decorated, 'extra', { value: true })
+    const compensatedSparse = new Array(1)
+    Object.defineProperty(compensatedSparse, 'extra', { value: true })
+    const symbolDecorated = [1]
+    Object.defineProperty(symbolDecorated, Symbol('extra'), { value: true })
 
     for (const value of [
       new ExoticObject(),
       new Map([['value', 1]]),
       new ExoticArray(1),
       new Array(1),
+      decorated,
+      compensatedSparse,
+      symbolDecorated,
       cyclic,
       [undefined],
       { value: undefined },