@tt-a1i/openpi 0.3.1 → 0.5.0

package/extensions/workflows/sandbox-child.cjs

@@ -164,7 +164,17 @@ const BOOTSTRAP = String.raw`
         }
         return value;
       } catch (error) {
-        log("pipeline: item " + index + " dropped — " + ((error && error.message) || String(error)));
+        callHost(
+          "log",
+          JSON.stringify({
+            text:
+              "pipeline: item " +
+              index +
+              " dropped — " +
+              ((error && error.message) || String(error)),
+            kind: "pipeline-drop",
+          }),
+        );
         return null;
       }
     });
@@ -207,6 +217,11 @@ const BOOTSTRAP = String.raw`
       const value = raw && typeof raw === "object" ? raw[key] : undefined;
       return typeof value === "number" && Number.isFinite(value) ? value : 0;
     };
+    const readLimit = (key) => {
+      const limits = raw && typeof raw === "object" ? raw.limits : undefined;
+      const value = limits && typeof limits === "object" ? limits[key] : undefined;
+      return typeof value === "number" && Number.isFinite(value) ? value : 0;
+    };
     // Every field is coerced, so a missing or malformed snapshot reads as zero
     // rather than undefined: a script doing arithmetic on it gets 0, not NaN.
     return deepFreeze({
@@ -217,6 +232,12 @@ const BOOTSTRAP = String.raw`
       total: read("total"),
       cost: read("cost"),
       agents: read("agents"),
+      limits: {
+        concurrency: readLimit("concurrency"),
+        maxAgentCalls: readLimit("maxAgentCalls"),
+        callsUsed: readLimit("callsUsed"),
+        callsRemaining: readLimit("callsRemaining"),
+      },
     });
   }
 
@@ -298,8 +319,9 @@ process.on("message", (message) => {
     }
     initialized = true;
     token = message.token;
+    const seq = typeof message.seq === "number" ? message.seq : undefined;
     if (typeof message.usageJson === "string") usageJson = message.usageJson;
-    run(message.source, message.argsJson, message.maxConcurrency);
+    run(message.source, message.argsJson, message.maxConcurrency, seq);
     return;
   }
   if (message.token !== token || message.kind !== "agentResult") return;
@@ -307,6 +329,7 @@ process.on("message", (message) => {
   const pending = pendingAgents.get(message.id);
   if (!pending) return;
   pendingAgents.delete(message.id);
+  const seq = typeof message.seq === "number" ? message.seq : undefined;
   if (typeof message.resultJson === "string")
     pending.resolve(message.resultJson);
   else
@@ -315,9 +338,14 @@ process.on("message", (message) => {
         typeof message.error === "string" ? message.error : "Agent IPC failed",
       ),
     );
+  if (seq !== undefined) {
+    setImmediate(() => {
+      send({ kind: "idle", seq });
+    });
+  }
 });
 
-function run(source, argsJson, maxConcurrency) {
+function run(source, argsJson, maxConcurrency, initSeq) {
   try {
     const sandbox = Object.create(null);
     sandbox.__argsJson = argsJson;
@@ -389,6 +417,11 @@ function run(source, argsJson, maxConcurrency) {
       context,
       { timeout: 1000 },
     );
+    if (initSeq !== undefined) {
+      setImmediate(() => {
+        send({ kind: "idle", seq: initSeq });
+      });
+    }
     Promise.resolve(context.__workflowPromise)
       .then((resultJson) => {
         if (typeof resultJson !== "string")

package/extensions/workflows/sandbox.ts

@@ -1,9 +1,9 @@
+import { type ChildProcess, spawn } from "node:child_process";
 import { randomBytes } from "node:crypto";
-import { spawn, type ChildProcess } from "node:child_process";
 import * as path from "node:path";
 import { fileURLToPath } from "node:url";
 import { MAX_WORKFLOW_AGENT_CALLS } from "../shared/setup-config.ts";
-import { safeStringify, toSerializable } from "./serialization.ts";
+import { encodeCompleteJson, toSerializable } from "./serialization.ts";
 
 const MAX_SOURCE_BYTES = 512 * 1024;
 const MAX_ARGS_BYTES = 256 * 1024;
@@ -19,6 +19,11 @@ const MAX_LOG_MESSAGE_BYTES = 8 * 1024;
  * bypasses the controller entirely.
  */
 export const AGENT_CALL_BACKSTOP_MARGIN = 8;
+/**
+ * Maximum synchronous execution time allowed between async agent boundaries.
+ * Non-yielding code (e.g. while(true){}) after await is terminated by this timeout.
+ */
+const SANDBOX_SYNC_TIMEOUT_MS = 1_000;
 
 export interface SandboxAgentOptions {
   agent_type?: unknown;
@@ -53,7 +58,7 @@ export interface RunWorkflowSandboxOptions {
     signal: AbortSignal,
   ) => Promise<SandboxAgentResult>;
   onPhase: (title: string) => void;
-  onLog: (text: string) => void;
+  onLog: (text: string, kind?: "pipeline-drop") => void;
   /**
    * Cumulative run usage, read at send time so the child's `usage()` reflects
    * the agent that just settled rather than a value captured at launch.
@@ -147,13 +152,23 @@ export function runWorkflowSandbox(options: RunWorkflowSandboxOptions) {
       ),
     );
 
-  const argsJson = safeStringify(
+  const encodedArgs = encodeCompleteJson(
     { defined: options.args !== undefined, value: options.args },
-    { maxBytes: MAX_ARGS_BYTES, maxDepth: 16, maxNodes: 10_000 },
+    {
+      maxBytes: MAX_ARGS_BYTES,
+      maxDepth: 16,
+      maxNodes: 10_000,
+      maxStringBytes: MAX_ARGS_BYTES,
+    },
   );
-  if (byteLength(argsJson) > MAX_ARGS_BYTES) {
-    return Promise.reject(new Error("Workflow args exceed the IPC limit"));
+  if (!encodedArgs.ok) {
+    return Promise.reject(
+      new Error(
+        `Workflow args exceed the ${MAX_ARGS_BYTES}-byte IPC limit (${encodedArgs.limit} limit at ${encodedArgs.path})`,
+      ),
+    );
   }
+  const argsJson = encodedArgs.json;
 
   return new Promise<unknown>((resolve, reject) => {
     const workerPath = fileURLToPath(
@@ -182,6 +197,28 @@ export function runWorkflowSandbox(options: RunWorkflowSandboxOptions) {
     const activeAgentRequests = new Map<number, AbortController>();
     let requestCount = 0;
     let finished = false;
+    let executionSeq = 0;
+    let executionWatchdog: NodeJS.Timeout | undefined;
+
+    const armWatchdog = (timeoutMs = SANDBOX_SYNC_TIMEOUT_MS) => {
+      if (finished) return 0;
+      if (executionWatchdog) clearTimeout(executionWatchdog);
+      const currentSeq = ++executionSeq;
+      executionWatchdog = setTimeout(() => {
+        if (finished || currentSeq !== executionSeq) return;
+        finish(new Error("Script execution timed out"));
+      }, timeoutMs);
+      executionWatchdog.unref?.();
+      return currentSeq;
+    };
+
+    const disarmWatchdog = (seq?: number) => {
+      if (seq !== undefined && seq !== executionSeq) return;
+      if (executionWatchdog) {
+        clearTimeout(executionWatchdog);
+        executionWatchdog = undefined;
+      }
+    };
 
     // The child parses this and falls back to zeros if it is ever unusable, so
     // a broken snapshot degrades `usage()` to a zero reading instead of
@@ -195,6 +232,7 @@ export function runWorkflowSandbox(options: RunWorkflowSandboxOptions) {
     };
 
     const cleanup = () => {
+      disarmWatchdog();
       for (const abortController of activeAgentRequests.values()) {
         abortController.abort(new Error("Workflow stopped"));
       }
@@ -277,12 +315,21 @@ export function runWorkflowSandbox(options: RunWorkflowSandboxOptions) {
           if (!isRecord(payload) || typeof payload.text !== "string") {
             throw new Error("invalid text");
           }
-          options.onLog(payload.text);
+          if (payload.kind !== undefined && payload.kind !== "pipeline-drop") {
+            throw new Error("invalid log kind");
+          }
+          options.onLog(payload.text, payload.kind);
         } catch {
           finish(new Error("Workflow sandbox sent an invalid log line"));
         }
         return;
       }
+      if (raw.kind === "idle") {
+        if (typeof raw.seq === "number") {
+          disarmWatchdog(raw.seq);
+        }
+        return;
+      }
       if (raw.kind === "agent") {
         if (
           typeof raw.payloadJson !== "string" ||
@@ -335,10 +382,12 @@ export function runWorkflowSandbox(options: RunWorkflowSandboxOptions) {
               error: "Agent result exceeded the workflow IPC output limit",
             });
           }
+          const seq = armWatchdog();
           child.send({
             token,
             kind: "agentResult",
             id,
+            seq,
             resultJson,
             usageJson: usageJson(),
           });
@@ -385,6 +434,10 @@ export function runWorkflowSandbox(options: RunWorkflowSandboxOptions) {
       finish(new Error("Workflow sandbox sent an unknown IPC message"));
     });
 
+    // Arm the synchronous execution watchdog for the initial script invocation
+    // (covers non-yielding code before and after initial microtask yields).
+    const initSeq = armWatchdog();
+
     child.send(
       {
         kind: "init",
@@ -393,6 +446,7 @@ export function runWorkflowSandbox(options: RunWorkflowSandboxOptions) {
         argsJson,
         maxConcurrency: options.maxConcurrency,
         usageJson: usageJson(),
+        seq: initSeq,
       },
       (error) => {
         if (error) finish(error);

package/extensions/workflows/serialization.ts

@@ -8,6 +8,24 @@ export interface SerializationOptions {
   maxStringBytes?: number;
 }
 
+export type CompleteJsonLimit = "bytes" | "depth" | "nodes" | "string";
+
+export type CompleteJsonEncoding =
+  | {
+      readonly ok: true;
+      readonly json: string;
+      readonly bytes: number;
+      readonly nodes: number;
+    }
+  | {
+      readonly ok: false;
+      readonly limit: CompleteJsonLimit;
+      readonly path: string;
+      readonly maximum: number;
+      readonly observedAtLeast: number;
+      readonly nodes: number;
+    };
+
 const DEFAULT_MAX_BYTES = 1024 * 1024;
 const DEFAULT_MAX_DEPTH = 16;
 const DEFAULT_MAX_NODES = 20_000;
@@ -17,6 +35,269 @@ function byteLength(value: string) {
   return Buffer.byteLength(value, "utf8");
 }
 
+function boundedPath(parent: string, child: string) {
+  const childPreview =
+    child.length <= 128 ? child : `${child.slice(0, 128)}...[key truncated]`;
+  const path = `${parent}.${childPreview}`;
+  return byteLength(path) <= 256
+    ? path
+    : `${truncateUtf8(path, 240)}...[path truncated]`;
+}
+
+function configuredBudget(name: string, value: number, minimum: number) {
+  if (!Number.isSafeInteger(value) || !Number.isFinite(value)) {
+    throw new RangeError(`${name} must be a finite integer`);
+  }
+  if (value < minimum) {
+    throw new RangeError(`${name} must be at least ${minimum}`);
+  }
+  return value;
+}
+
+/**
+ * Encode a complete inert JSON value while bounding work and output during
+ * traversal. Unlike `safeStringify`, this interface never returns a truncated
+ * value: callers that persist authoritative data can fail closed instead.
+ */
+export function encodeCompleteJson(
+  value: unknown,
+  options: Required<SerializationOptions>,
+): CompleteJsonEncoding {
+  const maxBytes = configuredBudget("maxBytes", options.maxBytes, 1);
+  const maxDepth = configuredBudget("maxDepth", options.maxDepth, 0);
+  const maxNodes = configuredBudget("maxNodes", options.maxNodes, 1);
+  const maxStringBytes = configuredBudget(
+    "maxStringBytes",
+    options.maxStringBytes,
+    0,
+  );
+  const chunks: string[] = [];
+  const seen = new WeakMap<object, string>();
+  let bytes = 0;
+  let nodes = 0;
+  let failure: Exclude<CompleteJsonEncoding, { ok: true }> | undefined;
+
+  const stop = (
+    limit: CompleteJsonLimit,
+    path: string,
+    maximum: number,
+    observedAtLeast: number,
+  ) => {
+    failure ??= {
+      ok: false,
+      limit,
+      path: truncateUtf8(path, 256),
+      maximum,
+      observedAtLeast,
+      nodes,
+    };
+    return false;
+  };
+
+  const append = (text: string, path: string) => {
+    const addition = byteLength(text);
+    if (bytes + addition > maxBytes) {
+      return stop("bytes", path, maxBytes, bytes + addition);
+    }
+    chunks.push(text);
+    bytes += addition;
+    return true;
+  };
+
+  const encodedString = (text: string, path: string) => {
+    const sourceBytes = byteLength(text);
+    if (sourceBytes > maxStringBytes) {
+      stop("string", path, maxStringBytes, sourceBytes);
+      return undefined;
+    }
+    return JSON.stringify(text);
+  };
+
+  const safeErrorText = (error: unknown) => {
+    try {
+      return error instanceof Error ? error.message : String(error);
+    } catch {
+      return "unknown error";
+    }
+  };
+
+  const visit = (current: unknown, depth: number, path: string): boolean => {
+    if (failure) return false;
+    if (nodes >= maxNodes) {
+      return stop("nodes", path, maxNodes, nodes + 1);
+    }
+    nodes++;
+    if (depth > maxDepth) {
+      return stop("depth", path, maxDepth, depth);
+    }
+
+    if (current === null) return append("null", path);
+    if (typeof current === "boolean")
+      return append(current ? "true" : "false", path);
+    if (typeof current === "string") {
+      const encoded = encodedString(current, path);
+      return encoded === undefined ? false : append(encoded, path);
+    }
+    if (typeof current === "number") {
+      return append(
+        Number.isFinite(current)
+          ? JSON.stringify(current)
+          : JSON.stringify(`[number: ${String(current)}]`),
+        path,
+      );
+    }
+    if (typeof current === "bigint")
+      return append(JSON.stringify(`${current.toString()}n`), path);
+    if (typeof current === "undefined")
+      return append(JSON.stringify("[undefined]"), path);
+    if (typeof current === "symbol")
+      return append(
+        JSON.stringify(`[symbol: ${current.description ?? ""}]`),
+        path,
+      );
+    if (typeof current === "function")
+      return append(
+        JSON.stringify(`[function: ${current.name || "anonymous"}]`),
+        path,
+      );
+    if (typeof current !== "object")
+      return append(JSON.stringify(String(current)), path);
+
+    const prior = seen.get(current);
+    if (prior) return append(JSON.stringify(`[circular: ${prior}]`), path);
+    seen.set(current, path);
+
+    if (current instanceof Date) {
+      try {
+        return append(
+          JSON.stringify(
+            Number.isNaN(current.getTime())
+              ? "[date: invalid]"
+              : current.toISOString(),
+          ),
+          path,
+        );
+      } catch (error) {
+        return append(
+          JSON.stringify(`[unreadable date: ${safeErrorText(error)}]`),
+          path,
+        );
+      }
+    }
+
+    if (current instanceof Error) {
+      let name = "Error";
+      let message = "";
+      let stack: string | undefined;
+      try {
+        name = current.name;
+      } catch {}
+      try {
+        message = current.message;
+      } catch {}
+      try {
+        stack = current.stack;
+      } catch {}
+      return visit(
+        {
+          name,
+          message,
+          ...(stack ? { stack: truncateUtf8(stack, 16 * 1024) } : {}),
+        },
+        depth,
+        path,
+      );
+    }
+
+    if (Array.isArray(current)) {
+      let length: number;
+      try {
+        length = current.length;
+      } catch (error) {
+        return visit(
+          `[unreadable array: ${safeErrorText(error)}]`,
+          depth + 1,
+          path,
+        );
+      }
+      if (!append("[", path)) return false;
+      for (let index = 0; index < length; index++) {
+        if (nodes >= maxNodes) {
+          return stop(
+            "nodes",
+            boundedPath(path, String(index)),
+            maxNodes,
+            nodes + 1,
+          );
+        }
+        if (index > 0 && !append(",", path)) return false;
+        if (!(index in current)) {
+          nodes++;
+          if (!append("null", boundedPath(path, String(index)))) return false;
+          continue;
+        }
+        let item: unknown;
+        try {
+          item = current[index];
+        } catch (error) {
+          item = `[unreadable property: ${safeErrorText(error)}]`;
+        }
+        if (!visit(item, depth + 1, boundedPath(path, String(index)))) {
+          return false;
+        }
+      }
+      return append("]", path);
+    }
+
+    const objectChunkStart = chunks.length;
+    const objectByteStart = bytes;
+    if (!append("{", path)) return false;
+    let count = 0;
+    try {
+      // JavaScript has no interruptible own-key iterator: a Proxy's ownKeys
+      // trap may allocate before iteration begins. Once keys are available,
+      // the budgets below stop before reading any further property values.
+      for (const key in current) {
+        if (!Object.prototype.hasOwnProperty.call(current, key)) continue;
+        if (nodes >= maxNodes) {
+          return stop("nodes", boundedPath(path, key), maxNodes, nodes + 1);
+        }
+        const keyBytes = byteLength(key);
+        const childPath = boundedPath(path, key);
+        if (keyBytes > maxStringBytes) {
+          return stop("string", childPath, maxStringBytes, keyBytes);
+        }
+        const keyJson = JSON.stringify(key);
+        if (count > 0 && !append(",", path)) return false;
+        if (!append(keyJson, childPath) || !append(":", childPath)) {
+          return false;
+        }
+        let child: unknown;
+        try {
+          child = (current as Record<string, unknown>)[key];
+        } catch (error) {
+          child = `[unreadable property: ${safeErrorText(error)}]`;
+        }
+        if (!visit(child, depth + 1, childPath)) return false;
+        count++;
+      }
+    } catch (error) {
+      chunks.length = objectChunkStart;
+      bytes = objectByteStart;
+      return visit(
+        `[unreadable object: ${safeErrorText(error)}]`,
+        depth + 1,
+        path,
+      );
+    }
+    return append("}", path);
+  };
+
+  visit(value, 0, "$root");
+  if (failure) return failure;
+  return { ok: true, json: chunks.join(""), bytes, nodes };
+}
+
 export function truncateUtf8(value: string, maxBytes: number) {
   if (maxBytes <= 0) return "";
   if (byteLength(value) <= maxBytes) return value;
@@ -72,9 +353,29 @@ export function toSerializable(
     seen.set(current, location);
 
     if (Array.isArray(current)) {
-      return current.map((item, index) =>
-        visit(item, depth + 1, `${location}[${index}]`),
-      );
+      const result: unknown[] = [];
+      const length = current.length;
+      for (let index = 0; index < length; index++) {
+        if (nodes >= maxNodes) {
+          result.push("[truncated: node limit]");
+          break;
+        }
+        if (!(index in current)) {
+          nodes++;
+          result.length++;
+          continue;
+        }
+        try {
+          result.push(
+            visit(current[index], depth + 1, `${location}[${index}]`),
+          );
+        } catch (error) {
+          result.push(
+            `[unreadable property: ${error instanceof Error ? error.message : String(error)}]`,
+          );
+        }
+      }
+      return result;
     }
 
     if (current instanceof Date) {
@@ -93,24 +394,31 @@ export function toSerializable(
     }
 
     const result: Record<string, unknown> = Object.create(null);
-    let keys: string[];
     try {
-      keys = Object.keys(current);
+      for (const key in current) {
+        if (!Object.prototype.hasOwnProperty.call(current, key)) continue;
+        if (nodes >= maxNodes) {
+          let marker = "[truncated: node limit]";
+          while (Object.prototype.hasOwnProperty.call(result, marker)) {
+            marker += ".";
+          }
+          result[marker] = true;
+          break;
+        }
+        try {
+          result[key] = visit(
+            (current as Record<string, unknown>)[key],
+            depth + 1,
+            `${location}.${key}`,
+          );
+        } catch (error) {
+          result[key] =
+            `[unreadable property: ${error instanceof Error ? error.message : String(error)}]`;
+        }
+      }
     } catch (error) {
       return `[unreadable object: ${error instanceof Error ? error.message : String(error)}]`;
     }
-    for (const key of keys) {
-      try {
-        result[key] = visit(
-          (current as Record<string, unknown>)[key],
-          depth + 1,
-          `${location}.${key}`,
-        );
-      } catch (error) {
-        result[key] =
-          `[unreadable property: ${error instanceof Error ? error.message : String(error)}]`;
-      }
-    }
     return result;
   };
 

package/extensions/workflows/tool-renderer.ts

@@ -0,0 +1,22 @@
+/** In-memory association between a workflow transcript projection and its child renderer. */
+
+import type { AgentToolRenderer } from "../shared/agent-tool-renderer.ts";
+import type { TranscriptEntry } from "./model.ts";
+
+const renderers = new WeakMap<
+  ReadonlyArray<TranscriptEntry>,
+  AgentToolRenderer
+>();
+
+export function bindWorkflowToolRenderer<
+  T extends ReadonlyArray<TranscriptEntry>,
+>(transcript: T, renderer: AgentToolRenderer): T {
+  renderers.set(transcript, renderer);
+  return transcript;
+}
+
+export function workflowToolRenderer(
+  transcript: ReadonlyArray<TranscriptEntry>,
+) {
+  return renderers.get(transcript);
+}