fix(util): last complete object wins when skipping JSON preamble (#718)

`createPartialJsonStream({ skipPreamble: true })` latched onto the FIRST
syntactically valid `{…}` in the preamble and permanently ignored the real
payload:

    input : '<think>the schema is {"type":"object"} so I will emit</think>{"name":"Bob"}'
    finish: {"type":"object"}   complete: true

`consumeSeek` opened a candidate at the first `{`; when it closed,
`afterValue()` set `Mode.Done` and `consume()` returned on every later chunk.
The `fail()` recovery only rescued candidates that produced NO data
(`rootIsEmpty()`), so a well-formed prose object was unrecoverable — and
`complete` reported `true`, so nothing downstream signalled a problem.

Both `skipPreamble` call sites (HFT, TF-MediaPipe) are exactly the local /
thinking models that restate the schema or a few-shot example before
answering. A model emitting `{"name":"Alice"}` as an illustration inside
`<think>` and then the real `{"name":"Bob"}` yielded `finish.data.object =
{"name":"Alice"}` — schema-valid, so `StructuredGenerationTask`'s
re-validation accepted it and the WRONG record was persisted.

A closed root is now provisional under `skipPreamble`: it is retained as
`lastComplete` and the parser returns to `Mode.Seek`, so the next `{` at
depth 0 starts a fresh candidate that supersedes it. Nothing is re-scanned —
a restart begins at the `{` that triggered it — so the parser stays O(total
input) and the live-root aliasing contract is untouched. `finish()` returns
the live root when it carries data (a genuinely truncated real payload),
else `lastComplete`, else `{}`.

Also in this change:

- `finish()` is declared `JsonValue` (new exported type) instead of
  `Record<string, unknown>`. It always returned arrays and scalars for
  array-/scalar-rooted documents; the declared type was a lie, so
  `Object.keys(stream.finish())` on an array root type-checked and produced
  `["0","1","2"]`. New `finishObject()` returns `{}` for a non-object root;
  all 11 provider run-fns migrated to it, so an array root now fails
  required-key validation loudly.
- `complete` no longer reports `true` for a root that only closed because
  `finish()` repaired a truncated token (`push('"abc')`), which contradicted
  its own JSDoc.
- `resolveRef.ts`: the per-container `hasRef` probe inside `walk` re-scanned
  each node's whole subtree, making `resolveOutput` O(n·depth); `containsCycle`
  added a second full traversal, and both recursed once per level so a deeply
  nested value raised a `RangeError` out of cache resolution. One iterative
  pass now answers all three questions up front (which containers hold a
  reachable ref, whether any ref exists, whether there is a back-edge).
  Measured: 4k-deep chain 590ms -> 30ms; a 50k-deep ref-free value no longer
  overflows.

Co-authored-by: Claude <noreply@anthropic.com>

Author: sroussey

.claude/CLAUDE.md

@@ -247,11 +247,30 @@ sequence of partial deltas cannot supply.
 Do **not** accumulate the JSON text to produce it. Feed the deltas to
 `createPartialJsonStream()` (`@workglow/util/worker`, or `/schema` off-worker):
 `push(chunk)` is O(chunk) and returns the partial object to emit as an
-`object-delta`, and `finish()` returns the value for `finish.data.object`.
+`object-delta`, and `finishObject()` returns the value for `finish.data.object`.
 Re-parsing a growing buffer on every delta is O(n²) and blocks the worker
 thread. Use the `skipPreamble` option for providers that emit prose, a
 `<think>` block, or a code fence ahead of the JSON.
 
+Use `finishObject()`, not `finish()`. `finish()` is typed `JsonValue` because it
+honestly returns whatever the document had at its root — an array or a scalar
+for a malformed response — and `StructuredGenerationTask` requires an object.
+`finishObject()` yields `{}` for a non-object root, so validation fails loudly on
+the missing required keys instead of the task receiving a value whose "keys" are
+array indices.
+
+`skipPreamble` is **last-complete-wins**: a closed root is provisional, so a
+later `{` starts a fresh candidate that supersedes it and the LAST complete
+object is what `finish()` returns. A thinking model that restates the schema or
+shows a few-shot example before answering would otherwise lock onto the prose
+object — and a schema-shaped one passes re-validation, so the wrong record gets
+persisted with no error anywhere. The cost is that trailing prose containing its
+own complete object supersedes the payload, so keep asking for the JSON last;
+trailing prose with no `{` in it never restarts anything. Nothing is re-scanned
+(a restart begins at the `{` that triggered it), so the parser stays O(total
+input) — but it does keep scanning trailing text for the life of the stream
+rather than exiting early at the first close.
+
 `push()` returns the parser's **live** root, which later pushes mutate — that
 aliasing is what keeps it linear. It is safe for the `object-delta` path
 (`StreamEventAccumulator` / `StreamProcessor` use replace semantics for

packages/task-graph/src/cache/resolveRef.ts

@@ -114,100 +114,110 @@ export async function resolveOutput<T>(
   resolver: CacheRefResolver,
   options?: ResolveOutputOptions
 ): Promise<T> {
-  if (!hasRef(output, new WeakSet())) return output;
+  if (isCacheRef(output)) {
+    const limit = createLimiter(options?.concurrency);
+    return (await walk(output, resolver, limit, new WeakSet(), undefined, new WeakSet())) as T;
+  }
+  const scan = scanContainers(output);
+  if (!scan.hasRef) return output;
   const limit = createLimiter(options?.concurrency);
   // Acyclic values (the norm) memoize each container's resolution promise so a
   // subtree shared between two slots resolves ONCE and both slots receive the
   // same resolved copy — a plain visited-set would hand the second slot the
   // original, unresolved object. Cyclic values keep the conservative
   // visited-set behavior: cycles are returned by reference, unrewritten.
-  const memo = containsCycle(output) ? undefined : new WeakMap<object, Promise<unknown>>();
-  return (await walk(output, resolver, limit, new WeakSet(), memo)) as T;
+  const memo = scan.hasCycle ? undefined : new WeakMap<object, Promise<unknown>>();
+  return (await walk(output, resolver, limit, new WeakSet(), memo, scan.refBearing)) as T;
 }
 
-/**
- * Depth-first cycle probe over the same container vocabulary as the walker
- * (plain objects, Array, Map, Set; leaves are opaque). Gray/black coloring:
- * an object seen again while still on the current path is a back-edge.
- */
-function containsCycle(
-  value: unknown,
-  gray: WeakSet<object> = new WeakSet(),
-  black: WeakSet<object> = new WeakSet()
-): boolean {
-  if (value === null || typeof value !== "object" || isLeaf(value)) return false;
-  const obj = value as object;
-  if (black.has(obj)) return false;
-  if (gray.has(obj)) return true;
-  gray.add(obj);
-  const children: Iterable<unknown> = Array.isArray(value)
-    ? value
-    : value instanceof Map
-      ? value.values()
-      : value instanceof Set
-        ? value
-        : Object.values(value);
-  for (const child of children) {
-    if (containsCycle(child, gray, black)) return true;
-  }
-  gray.delete(obj);
-  black.add(obj);
-  return false;
+/** What one traversal of the output tells the walker. */
+interface ContainerScan {
+  /** Containers from which a {@link CacheRef} is reachable. */
+  readonly refBearing: WeakSet<object>;
+  /** Whether any ref is reachable at all — lets `resolveOutput` preserve identity. */
+  readonly hasRef: boolean;
+  /** Whether a back-edge exists, which disables resolution memoization. */
+  readonly hasCycle: boolean;
+}
+
+/** Children the walker would descend into, in walk order. */
+function childrenOf(value: object): readonly unknown[] {
+  if (Array.isArray(value)) return value;
+  if (value instanceof Map) return Array.from(value.values());
+  if (value instanceof Set) return Array.from(value);
+  return Object.values(value as Record<string, unknown>);
 }
 
 /**
- * Cheap pre-scan: returns `true` if any {@link CacheRef} is reachable inside
- * `value`. Lets `resolveOutput` short-circuit and preserve identity when
- * nothing needs resolving.
+ * Single iterative pass over the walker's container vocabulary (plain objects,
+ * Array, Map, Set; leaves are opaque) answering everything the walk needs to
+ * know up front: which containers hold a reachable {@link CacheRef}, whether
+ * any ref exists at all, and whether the graph has a back-edge.
  *
- * `visited` short-circuits cyclic and shared-subtree structures: revisiting an
- * already-seen object answers `false` instead of recursing forever. The
- * pre-scan is a containment check, so reporting `false` on a revisit is safe —
- * if a ref were reachable through that subtree, the FIRST visit would have
- * found it.
+ * One pass, not three: a per-container `hasRef` probe inside the walk re-scanned
+ * each node's whole subtree, making resolution O(n·depth). It is iterative for
+ * the same reason — a deeply nested model output (thousands of levels) blew the
+ * stack in the recursive probe, surfacing as a `RangeError` from cache
+ * resolution rather than a cache miss.
+ *
+ * Cycles keep the pre-existing approximation: a container reached again while
+ * still on the current path contributes nothing to its parent's answer. A ref
+ * reachable ONLY through a back-edge can therefore go unnoticed, which is
+ * consistent with the walker returning cycles by reference, unrewritten.
  */
-function hasRef(value: unknown, visited: WeakSet<object>): boolean {
-  if (isCacheRef(value)) return true;
-  if (value === null || value === undefined) return false;
-  if (isLeaf(value)) return false;
-  if (typeof value === "object") {
-    if (visited.has(value as object)) return false;
-    visited.add(value as object);
+function scanContainers(root: unknown): ContainerScan {
+  const refBearing = new WeakSet<object>();
+  if (root === null || typeof root !== "object" || isLeaf(root)) {
+    return { refBearing, hasRef: false, hasCycle: false };
   }
-  if (Array.isArray(value)) {
-    for (const v of value) {
-      if (hasRef(v, visited)) return true;
+  const done = new WeakSet<object>();
+  const onPath = new WeakSet<object>();
+  let hasCycle = false;
+
+  type Frame = { readonly node: object; readonly children: readonly unknown[]; index: number };
+  const stack: Frame[] = [{ node: root, children: childrenOf(root), index: 0 }];
+  onPath.add(root);
+
+  while (stack.length > 0) {
+    const frame = stack[stack.length - 1]!;
+    if (frame.index >= frame.children.length) {
+      stack.pop();
+      onPath.delete(frame.node);
+      done.add(frame.node);
+      const parent = stack[stack.length - 1];
+      if (parent !== undefined && refBearing.has(frame.node)) refBearing.add(parent.node);
+      continue;
     }
-    return false;
-  }
-  if (value instanceof Map) {
-    for (const v of value.values()) {
-      if (hasRef(v, visited)) return true;
+    const child = frame.children[frame.index++];
+    if (isCacheRef(child)) {
+      refBearing.add(frame.node);
+      continue;
     }
-    return false;
-  }
-  if (value instanceof Set) {
-    for (const v of value) {
-      if (hasRef(v, visited)) return true;
+    if (child === null || typeof child !== "object" || isLeaf(child)) continue;
+    const obj = child as object;
+    if (onPath.has(obj)) {
+      hasCycle = true;
+      continue;
     }
-    return false;
-  }
-  if (typeof value === "object") {
-    const source = value as Record<string, unknown>;
-    for (const k of Object.keys(source)) {
-      if (hasRef(source[k], visited)) return true;
+    if (done.has(obj)) {
+      // Shared subtree: its answer is already final, so reuse it.
+      if (refBearing.has(obj)) refBearing.add(frame.node);
+      continue;
     }
-    return false;
+    onPath.add(obj);
+    stack.push({ node: obj, children: childrenOf(obj), index: 0 });
   }
-  return false;
+
+  return { refBearing, hasRef: refBearing.has(root), hasCycle };
 }
 
 async function walk(
   value: unknown,
   resolver: CacheRefResolver,
   limit: Limiter,
   visited: WeakSet<object>,
-  memo: WeakMap<object, Promise<unknown>> | undefined
+  memo: WeakMap<object, Promise<unknown>> | undefined,
+  refBearing: WeakSet<object>
 ): Promise<unknown> {
   if (isCacheRef(value)) {
     return limit.run(() => resolver(value));
@@ -226,39 +236,44 @@ async function walk(
     // topology, including any unresolved refs the cycle contains.
     return value;
   }
-  if (!hasRef(value, new WeakSet())) return value;
+  // Answered by the single up-front scan; probing the subtree here made
+  // resolution O(n·depth) and recursed once per level.
+  if (!refBearing.has(obj)) return value;
   if (memo) {
-    const promise = walkContainer(value, resolver, limit, visited, memo);
+    const promise = walkContainer(value, resolver, limit, visited, memo, refBearing);
     memo.set(obj, promise);
     return promise;
   }
   visited.add(obj);
-  return walkContainer(value, resolver, limit, visited, memo);
+  return walkContainer(value, resolver, limit, visited, memo, refBearing);
 }
 
 async function walkContainer(
   value: unknown,
   resolver: CacheRefResolver,
   limit: Limiter,
   visited: WeakSet<object>,
-  memo: WeakMap<object, Promise<unknown>> | undefined
+  memo: WeakMap<object, Promise<unknown>> | undefined,
+  refBearing: WeakSet<object>
 ): Promise<unknown> {
   if (Array.isArray(value)) {
-    return Promise.all(value.map((v) => walk(v, resolver, limit, visited, memo)));
+    return Promise.all(value.map((v) => walk(v, resolver, limit, visited, memo, refBearing)));
   }
   if (value instanceof Map) {
     const out = new Map();
     const entries = Array.from(value.entries());
     const resolved = await Promise.all(
-      entries.map(async ([k, v]) => [k, await walk(v, resolver, limit, visited, memo)] as const)
+      entries.map(
+        async ([k, v]) => [k, await walk(v, resolver, limit, visited, memo, refBearing)] as const
+      )
     );
     for (const [k, v] of resolved) out.set(k, v);
     return out;
   }
   if (value instanceof Set) {
     const out = new Set();
     const resolved = await Promise.all(
-      Array.from(value).map((v) => walk(v, resolver, limit, visited, memo))
+      Array.from(value).map((v) => walk(v, resolver, limit, visited, memo, refBearing))
     );
     for (const v of resolved) out.add(v);
     return out;
@@ -273,7 +288,7 @@ async function walkContainer(
     // matches the input even though resolutions race.
     const keys = Object.keys(source);
     const resolvedValues = await Promise.all(
-      keys.map((k) => walk(source[k], resolver, limit, visited, memo))
+      keys.map((k) => walk(source[k], resolver, limit, visited, memo, refBearing))
     );
     for (let i = 0; i < keys.length; i++) out[keys[i]!] = resolvedValues[i];
     return out;

packages/task-graph/src/task-graph/__tests__/resolveOutput.test.ts

@@ -192,4 +192,53 @@ describe("resolveOutput", () => {
     expect(out.target instanceof URL).toBe(true);
     expect(resolver).not.toHaveBeenCalled();
   });
+
+  describe("deep values", () => {
+    /** A left-spine chain `{child:{child:…{bottom}}}`, `width` scalars per level. */
+    const chain = (
+      depth: number,
+      bottom: Record<string, unknown>,
+      width = 0
+    ): Record<string, unknown> => {
+      let node = bottom;
+      for (let i = 0; i < depth; i++) {
+        const next: Record<string, unknown> = { child: node };
+        for (let k = 0; k < width; k++) next[`f${k}`] = `v${k}`;
+        node = next;
+      }
+      return node;
+    };
+
+    it("scans a 50k-deep ref-free value without overflowing the stack", async () => {
+      // The pre-scan recursed once per level, so a deeply nested model output
+      // raised a RangeError out of cache resolution instead of answering "no
+      // refs here". Identity is still preserved: nothing needed resolving, and
+      // the walker is never entered.
+      const resolver = vi.fn<CacheRefResolver>();
+      const input = chain(50_000, { leaf: "bottom" });
+      const out = await resolveOutput(input, resolver);
+      expect(out).toBe(input);
+      expect(resolver).not.toHaveBeenCalled();
+    });
+
+    it("resolves a ref under a deep chain without re-scanning each level", async () => {
+      // A per-container pre-scan inside the walk re-read each node's whole
+      // subtree, making resolution O(n·depth). The bound is loose enough that
+      // only an algorithmic regression trips it: ~590ms before, ~30ms after.
+      const depth = 400;
+      const blob = new Blob([new Uint8Array([1])]);
+      const input = chain(depth, { payload: ref("cache://deep") as unknown as Blob }, 100);
+      const started = performance.now();
+      const out = (await resolveOutput(input, fakeResolver({ "cache://deep": blob }))) as Record<
+        string,
+        unknown
+      >;
+      const elapsed = performance.now() - started;
+
+      let cursor: Record<string, unknown> = out;
+      for (let i = 0; i < depth; i++) cursor = cursor.child as Record<string, unknown>;
+      expect(cursor.payload).toBe(blob);
+      expect(elapsed).toBeLessThan(250);
+    });
+  });
 });

packages/util/CHANGELOG.md

@@ -1,5 +1,40 @@
 # @workglow/util
 
+## Unreleased
+
+### Bug Fixes
+
+#### util
+
+- `createPartialJsonStream({ skipPreamble: true })`: **last complete object wins**.
+  A closed root is now provisional — it is retained and scanning resumes, so a
+  later `{` starts a fresh candidate that supersedes it. Previously the parser
+  latched onto the FIRST syntactically valid object in the preamble and ignored
+  the real payload, while still reporting `complete: true`. A thinking model that
+  restates the schema (`{"type":"object"}`) or shows a few-shot example
+  (`{"name":"Alice"}`) before answering therefore returned the wrong document —
+  and a schema-shaped one passes re-validation, so it was persisted silently.
+  Nothing is re-scanned, so the parser stays O(total input); the accepted
+  trade-off is that trailing prose containing a complete object now supersedes
+  the payload. `Mode.Done` is no longer reachable in this mode, so the parser
+  keeps scanning trailing text for the life of the stream (bounded, no buffering).
+- `PartialJsonStream.complete` no longer reports `true` for a root that only
+  closed because `finish()` repaired a truncated token (`push('"abc')`).
+
+### BREAKING
+
+#### util
+
+- `PartialJsonStream.finish()` is now declared as `JsonValue` (a new exported
+  type) instead of `Record<string, unknown>`. The runtime behavior is unchanged
+  — an array- or scalar-rooted document was always returned as such — but the
+  declared type was a lie, so `Object.keys(stream.finish())` on an array root
+  type-checked and silently produced `["0","1","2"]`. Callers feeding a consumer
+  that requires an object should switch to the new
+  `PartialJsonStream.finishObject()`, which yields `{}` for an array or scalar
+  root; all in-repo provider run-fns have been migrated. Third-party run-fns
+  calling `finish()` will see a compile break — intentionally.
+
 ## 0.3.38
 
 ### Features