@yolk-sdk/codemode 0.1.0-canary.96 → 0.1.0-canary.97

package/src/node.ts

@@ -8,6 +8,7 @@ import {
   type CodemodeTool,
   type CodemodeWasmModule
 } from '@earendil-works/pi-codemode'
+import { codeModeCallLabels } from './catalog.ts'
 import type {
   CodeModeError,
   CodeModeExecuteOptions,
@@ -96,6 +97,215 @@ const wrapperPrefix = 'async function __yolkCodeMode() {\n'
 
 const wrapperSuffix = '\n}'
 
+/** Most values one argument may hold; a larger argument is rejected. */
+const maxArgumentValues = 100_000
+
+/** Deepest object/array nesting one argument may have; deeper arguments are rejected. */
+const maxArgumentDepth = 64
+
+/**
+ * Evaluated in the VM before the script runs: `(tools, labels) => guardedTools`, with `labels` the
+ * `[name, callLabel]` pairs of the tools in order (`codeModeCallLabels`). Each call walks its first
+ * argument once, the way `JSON.stringify` reads it (each own enumerable string key and each array
+ * index read once, so getters run once; a callable `toJSON` is called once with the key, `''` at
+ * the root, and its result is walked at the same path), and builds a fresh copy (null-prototype
+ * objects and arrays) that is what the call sends. It rejects with a `TypeError` naming the call
+ * label and path when JSON would change the value instead of carrying it: non-finite numbers,
+ * invalid `Date`s, `undefined`/function/symbol array items and holes, function/symbol/bigint
+ * values, cycles, and objects that are neither plain (prototype `Object.prototype` or `null`) nor
+ * arrays and have no `toJSON`. An argument with more than `maxArgumentValues` values or nested more
+ * than `maxArgumentDepth` levels is rejected too, never sent unchecked. `undefined`-valued keys are
+ * left out (absent, as JSON drops them); the same object twice (not a cycle) is copied twice. An
+ * error thrown by a getter or `toJSON` rejects the call with that error. A rejected call never
+ * reaches the host. Every member of pi's `tools` is wrapped; unknown members fall through to pi's
+ * proxy (close-match suggestions). Not a security boundary (`globalThis.tools` is unguarded):
+ * arguments are still decoded by the host.
+ */
+const toolsGuardSource = `(function (tools, labels) {
+  // Captured before the script runs, so a script that changes built-ins cannot change the checks.
+  const { apply, getPrototypeOf, setPrototypeOf } = Reflect
+  const isFinite = Number.isFinite
+  const isArray = Array.isArray
+  const { keys, create, defineProperty } = Object
+  const getTime = Date.prototype.getTime
+  const TypeErrorCtor = TypeError
+  const SetCtor = Set
+  const setAdd = Set.prototype.add
+  const setHas = Set.prototype.has
+  const setDelete = Set.prototype.delete
+  const ObjectProto = Object.prototype
+  const { trunc, min } = Math
+  const stringify = JSON.stringify
+  const plainJson = 'pass plain JSON (objects, arrays, strings, finite numbers, booleans, null)'
+  const maxValues = ${maxArgumentValues}
+  const maxDepth = ${maxArgumentDepth}
+  class Rejected {
+    constructor(message) {
+      defineProperty(this, 'message', { value: message })
+    }
+  }
+  const keyPath = (path, key) =>
+    /^[A-Za-z_$][\\w$]*$/.test(key)
+      ? (path === '' ? key : path + '.' + key)
+      : path + '[' + stringify(key) + ']'
+  const typeName = value => {
+    const proto = getPrototypeOf(value)
+    const name = proto && typeof proto.constructor === 'function' ? proto.constructor.name : ''
+    return name ? (/^[AEIOU]/.test(name) ? 'an ' : 'a ') + name : 'a non-plain object'
+  }
+  // The JSON copy of a value already read from its holder at key (undefined: absent), or throws
+  // Rejected with why JSON would not carry it unchanged.
+  const snapshot = (read, key, path, slot, walk) => {
+    const at = path === '' ? 'is ' : 'at ' + path + ' is '
+    if (++walk.values > maxValues) {
+      throw new Rejected('has more than ' + maxValues + ' values; split the work across calls')
+    }
+    let value = read
+    const kind = typeof value
+    if ((kind === 'object' && value !== null) || kind === 'function' || kind === 'bigint') {
+      // The intrinsic getTime reads a Date's internal time (and throws for anything else).
+      let time
+      try {
+        time = apply(getTime, value, [])
+      } catch {
+        time = undefined
+      }
+      if (time !== undefined && !isFinite(time)) {
+        throw new Rejected(at + 'an invalid Date; pass a valid Date or an ISO string')
+      }
+      const toJSON = value.toJSON
+      if (typeof toJSON === 'function') value = apply(toJSON, value, [key])
+    }
+    switch (typeof value) {
+      case 'string':
+      case 'boolean':
+        return value
+      case 'number':
+        if (isFinite(value)) return value
+        throw new Rejected(
+          at + value + '; pass a finite number' + (slot === 'key' ? ' or omit the key' : '')
+        )
+      case 'undefined':
+        if (slot === 'item') throw new Rejected(at + 'undefined; arrays cannot hold undefined')
+        return undefined
+      case 'bigint':
+      case 'function':
+      case 'symbol':
+        throw new Rejected(at + 'a ' + typeof value + '; ' + plainJson)
+    }
+    if (value === null) return null
+    if (apply(setHas, walk.ancestors, [value])) {
+      throw new Rejected(at + 'a circular reference; ' + plainJson)
+    }
+    if (walk.depth >= maxDepth) {
+      throw new Rejected('at ' + path + ' is nested more than ' + maxDepth + ' levels deep')
+    }
+    const array = isArray(value)
+    if (!array) {
+      const proto = getPrototypeOf(value)
+      if (proto !== ObjectProto && proto !== null) {
+        throw new Rejected(at + typeName(value) + '; ' + plainJson)
+      }
+    }
+    apply(setAdd, walk.ancestors, [value])
+    walk.depth++
+    let copy
+    if (array) {
+      // LengthOfArrayLike, read once, as JSON.stringify does.
+      // Unary plus is ToNumber: one coercion, and it throws for a bigint as JSON.stringify does.
+      const raw = +value.length
+      const length = raw > 0 ? min(trunc(raw), 9007199254740991) : 0
+      // No prototype: no inherited toJSON or index setter can change the checked copy when pi
+      // serializes it (still an array to JSON.stringify).
+      copy = []
+      setPrototypeOf(copy, null)
+      for (let i = 0; i < length; i++) {
+        copy[i] = snapshot(value[i], '' + i, path + '[' + i + ']', 'item', walk)
+      }
+    } else {
+      copy = create(null)
+      // Indexed over the fresh key array: no iterator a script could replace.
+      const names = keys(value)
+      for (let i = 0; i < names.length; i++) {
+        const name = names[i]
+        const item = snapshot(value[name], name, keyPath(path, name), 'key', walk)
+        if (item !== undefined) copy[name] = item
+      }
+    }
+    walk.depth--
+    apply(setDelete, walk.ancestors, [value])
+    return copy
+  }
+  const wrap = (label, call) => async (...args) => {
+    let json
+    try {
+      json = snapshot(args[0], '', '', 'root', { values: 0, depth: 0, ancestors: new SetCtor() })
+    } catch (error) {
+      if (error instanceof Rejected) throw new TypeErrorCtor(label + ': argument ' + error.message)
+      throw error
+    }
+    return call(json)
+  }
+  // Each tool's function sits at tools[name] unless an earlier tool holds that key, in which case
+  // that earlier function is already wrapped.
+  const wrappers = new Map()
+  for (const [name, label] of labels) {
+    const call = name in tools ? tools[name] : undefined
+    if (typeof call === 'function' && !wrappers.has(call)) wrappers.set(call, wrap(label, call))
+  }
+  const guarded = Object.create(null)
+  for (const key of Object.keys(tools)) {
+    const call = tools[key]
+    if (!wrappers.has(call)) wrappers.set(call, wrap('tools[' + JSON.stringify(key) + ']', call))
+    guarded[key] = wrappers.get(call)
+  }
+  return new Proxy(Object.freeze(guarded), {
+    get: (target, key) => (key in target ? target[key] : tools[key])
+  })
+})`
+
+// pi runs `(async (tools, console) => {<code>\n})`. The script runs in an inner function of the
+// same shape whose `tools` is guarded; the prefix shares line 1 with the script (only line-1
+// columns shift) and the guard follows the script, so its line numbers are unchanged.
+const guardPrefix = 'return (async (tools, console) => {'
+
+/** JSON as a JavaScript expression (U+2028/U+2029 escaped for older parsers). */
+const jsonSource = (value: unknown) =>
+  JSON.stringify(value).replace(/[\u2028\u2029]/g, char =>
+    char === '\u2028' ? '\\u2028' : '\\u2029'
+  )
+
+const guardSuffix = (tools: ReadonlyArray<CodeModeExecutorTool>) => {
+  const names = tools.map(tool => tool.name)
+  const labels = codeModeCallLabels(names)
+
+  // `makeCodeModeTool` supplies each label; other callers get pi's binding rule.
+  const pairs = tools.map((tool, index) => [
+    tool.name,
+    tool.callLabel ?? labels[index] ?? `tools[${JSON.stringify(tool.name)}]`
+  ])
+
+  return `\n})(${toolsGuardSource}(tools, ${jsonSource(pairs)}), console)`
+}
+
+const lineTerminators = /\r\n?|[\n\u2028\u2029]/g
+
+/** Lines of the script as QuickJS counts them (at least; never fewer). */
+const scriptLineCount = (code: string) => (code.match(lineTerminators)?.length ?? 0) + 1
+
+const scriptFrame = /codemode\.js:(\d+):\d+/
+
+/** Drops stack frames of the wrapper and guard, which sit below the script's last line. */
+const withoutWrapperFrames = (stack: string, scriptLines: number) =>
+  stack
+    .split('\n')
+    .filter(line => {
+      const frame = scriptFrame.exec(line)
+
+      return frame === null || Number(frame[1]) <= scriptLines
+    })
+    .join('\n')
+
 const ErrorWithCode = Schema.Struct({ code: Schema.String })
 
 const errorCode = (error: unknown) =>
@@ -205,7 +415,7 @@ type MutableExecutionResult = {
   storeWrites?: CodeModeStoreWrites
 }
 
-const fromPiResult = (result: CodemodeResult): CodeModeExecutionResult => {
+const fromPiResult = (result: CodemodeResult, scriptLines: number): CodeModeExecutionResult => {
   if (result.ok) {
     const storeWrites = decodeStoreWrites(result.storeWrites)
 
@@ -235,7 +445,7 @@ const fromPiResult = (result: CodemodeResult): CodeModeExecutionResult => {
     message: named ? `${name}: ${message}` : message
   }
 
-  if (stack !== undefined) error.stack = stack
+  if (stack !== undefined) error.stack = withoutWrapperFrames(stack, scriptLines)
 
   return { ok: false, error, output: result.output }
 }
@@ -245,7 +455,8 @@ const describeError = (error: unknown) => (error instanceof Error ? error.messag
 /**
  * A `CodeModeExecutor` on `@earendil-works/pi-codemode`: every execution gets a fresh QuickJS VM
  * (WebAssembly) in a fresh worker thread, closed when the execution ends. Applies the timeout,
- * heap cap, abort signal, and store, strips TypeScript annotations first, and caps concurrent
+ * heap cap, abort signal, and store, strips TypeScript annotations first, rejects tool arguments
+ * that JSON would silently change inside the script (see `CodeModeExecutorTool`), and caps concurrent
  * executions per executor; create one executor per process (module scope) so the cap is
  * per process.
  *
@@ -299,11 +510,12 @@ export const makePiCodeModeExecutor = (
 
     try {
       return fromPiResult(
-        await sandbox.execute(stripped.code, {
+        await sandbox.execute(`${guardPrefix}${stripped.code}${guardSuffix(execution.tools)}`, {
           signal: execution.signal,
           store: execution.store,
           timeoutMs
-        })
+        }),
+        scriptLineCount(stripped.code)
       )
     } finally {
       await sandbox.close()

package/src/tool.ts

@@ -150,7 +150,9 @@ type NestedResolution =
   | { readonly ok: false; readonly message: string }
 
 /** What a nested call resolves to: `structuredContent` for tools with an output schema (when
- * present), text content otherwise; error results reject with their text.
+ * present), text content otherwise; error results reject with their text. A tool that declares an
+ * output schema but returns no `structuredContent` resolves to its text: that is a tool bug, kept
+ * readable rather than turned into a failure.
  */
 const resolveNestedResult = (tool: CodeModeCatalogTool, result: ToolResult): NestedResolution => {
   const text = nestedText(result.content)
@@ -317,13 +319,13 @@ const runScript = <Context>(input: RunInput<Context>): Effect.Effect<ToolResult,
             (tool: CodeModeCatalogTool): CodeModeExecutorTool['execute'] =>
             (args, { signal }) => {
               if (signal.aborted) {
-                return Promise.reject(new Error(`tools.${tool.identifier} was cancelled.`))
+                return Promise.reject(new Error(`${tool.callLabel} was cancelled.`))
               }
 
               if (records.length >= limits.maxNestedCalls) {
                 return Promise.reject(
                   new Error(
-                    `Nested call limit reached: a script may make at most ${limits.maxNestedCalls} tool calls. Batch the work or return partial results.`
+                    `${tool.callLabel}: Nested call limit reached: a script may make at most ${limits.maxNestedCalls} tool calls. Batch the work or return partial results.`
                   )
                 )
               }
@@ -375,8 +377,8 @@ const runScript = <Context>(input: RunInput<Context>): Effect.Effect<ToolResult,
                     reject(
                       new Error(
                         Cause.hasInterruptsOnly(exit.cause)
-                          ? `tools.${tool.identifier} was cancelled.`
-                          : `tools.${tool.identifier} failed unexpectedly.`
+                          ? `${tool.callLabel} was cancelled.`
+                          : `${tool.callLabel} failed unexpectedly.`
                       )
                     )
 
@@ -388,7 +390,8 @@ const runScript = <Context>(input: RunInput<Context>): Effect.Effect<ToolResult,
                   if (resolution.ok) {
                     resolve(resolution.value)
                   } else {
-                    reject(new Error(resolution.message))
+                    // The script sees which call failed; the nested-call record keeps the raw text.
+                    reject(new Error(`${tool.callLabel}: ${resolution.message}`))
                   }
                 })
               })
@@ -399,6 +402,7 @@ const runScript = <Context>(input: RunInput<Context>): Effect.Effect<ToolResult,
             description: tool.description,
             inputSchema: tool.inputSchema,
             outputSchema: tool.outputSchema,
+            callLabel: tool.callLabel,
             execute: callTool(tool)
           }))