diff --git a/packages/task-graph/src/cache/CacheRef.ts b/packages/task-graph/src/cache/CacheRef.ts index 20222a48e..eeffffd28 100644 --- a/packages/task-graph/src/cache/CacheRef.ts +++ b/packages/task-graph/src/cache/CacheRef.ts @@ -4,6 +4,14 @@ * SPDX-License-Identifier: Apache-2.0 */ +/** + * Brand value for {@link CacheRef}. A literal string (not a Symbol) so the brand + * survives JSON serialization across queue rows / IPC boundaries — a Symbol-based + * brand would be erased by `JSON.stringify` and the resulting object would no + * longer be identifiable as a cache reference on the receiving side. + */ +export const CACHE_REF_KIND = "task-graph/CacheRef" as const; + /** * A reference to bytes that live in the configured cache backing rather than * inline in a task `Output`. Emitted by `TaskRunner` for binary output ports @@ -14,24 +22,53 @@ * it back into bytes. `size` and `mime` are best-effort hints populated when * known at finish time; absent values do not imply unknown failure. * + * The `kind` brand discriminates a cache ref from other `{$ref: string}` + * shapes (e.g. JSON-Schema references) so the resolver never walks an + * untrusted `$ref` string into the cache. The brand is a literal so it survives + * JSON round-trip across queue boundaries. + * * Resolution is best-effort: the cache backing's TTL is the lifetime contract, * and `resolveOutputRef` returns `undefined` when the underlying entry has * been evicted. */ -export type CacheRef = { +export interface ICacheRef { + readonly kind: typeof CACHE_REF_KIND; readonly $ref: string; readonly size?: number; readonly mime?: string; -}; +} + +export type CacheRef = ICacheRef; /** - * Narrow an unknown value to {@link CacheRef}. The discriminator is a `$ref` - * property of type `string`; other fields are optional and not inspected. + * Narrow an unknown value to {@link CacheRef}. Discriminates on the literal + * {@link CACHE_REF_KIND} brand AND a string `$ref`; shape-only `{$ref: string}` + * objects (JSON-Schema refs, user metadata) do NOT match. */ export function isCacheRef(value: unknown): value is CacheRef { if (typeof value !== "object" || value === null) return false; - const candidate = value as { readonly $ref?: unknown }; - return typeof candidate.$ref === "string"; + const candidate = value as { readonly kind?: unknown; readonly $ref?: unknown }; + return candidate.kind === CACHE_REF_KIND && typeof candidate.$ref === "string"; +} + +/** + * Construct a branded {@link CacheRef}. Cache backings MUST use this helper (or + * spread `{kind: CACHE_REF_KIND, ...}` themselves) so the resulting ref carries + * the brand. Helpers in {@link CacheCoordinator} / {@link RunPrivateCacheRepo} + * defensively re-wrap legacy backings whose `saveOutputStream` predates the + * brand and returns an unbranded `{$ref}` shape. + */ +export function makeCacheRef(raw: { + readonly $ref: string; + readonly size?: number; + readonly mime?: string; +}): CacheRef { + return { + kind: CACHE_REF_KIND, + $ref: raw.$ref, + ...(raw.size !== undefined && { size: raw.size }), + ...(raw.mime !== undefined && { mime: raw.mime }), + }; } /** diff --git a/packages/task-graph/src/cache/RunPrivateCacheRepo.ts b/packages/task-graph/src/cache/RunPrivateCacheRepo.ts index 33e7b8bf0..a82b56244 100644 --- a/packages/task-graph/src/cache/RunPrivateCacheRepo.ts +++ b/packages/task-graph/src/cache/RunPrivateCacheRepo.ts @@ -7,6 +7,7 @@ import { TaskOutputRepository } from "../storage/TaskOutputRepository"; import type { TaskInput, TaskOutput } from "../task/TaskTypes"; import type { CacheRef } from "./CacheRef"; +import { isCacheRef, makeCacheRef } from "./CacheRef"; export interface RunPrivateCacheRepoOptions { backing: TaskOutputRepository; @@ -111,7 +112,7 @@ export class RunPrivateCacheRepo extends TaskOutputRepository { * via the wrapped `taskType`). Resolvers calling `getOutputByRef` on this * wrapper forward to the backing, which decodes its own `$ref`. */ - public override saveOutputStream( + public override async saveOutputStream( taskType: string, inputs: TaskInput, chunks: AsyncIterable, @@ -119,14 +120,16 @@ export class RunPrivateCacheRepo extends TaskOutputRepository { ): Promise { const fn = this.backing.saveOutputStream; if (typeof fn !== "function") { - return Promise.reject( - new Error( - `RunPrivateCacheRepo: backing repository does not implement saveOutputStream. ` + - `Call supportsStreaming() before saveOutputStream.` - ) + throw new Error( + `RunPrivateCacheRepo: backing repository does not implement saveOutputStream. ` + + `Call supportsStreaming() before saveOutputStream.` ); } - return fn.call(this.backing, this.ns(taskType), inputs, chunks, metadata); + // Re-wrap the backing's CacheRef so legacy backings that pre-date the + // `kind` brand still produce a discriminator-bearing ref through this + // wrapper. Branded refs pass through unchanged. + const raw = await fn.call(this.backing, this.ns(taskType), inputs, chunks, metadata); + return isCacheRef(raw) ? raw : makeCacheRef(raw); } /** diff --git a/packages/task-graph/src/task-graph/StreamPump.ts b/packages/task-graph/src/task-graph/StreamPump.ts index f38c039bb..ab3a2c509 100644 --- a/packages/task-graph/src/task-graph/StreamPump.ts +++ b/packages/task-graph/src/task-graph/StreamPump.ts @@ -8,7 +8,12 @@ import type { ResourceScope, ServiceRegistry } from "@workglow/util"; import type { TaskOutputRepository } from "../storage/TaskOutputRepository"; import type { ITask } from "../task/ITask"; import type { StreamEvent, StreamMode } from "../task/StreamTypes"; -import { edgeNeedsAccumulation, getOutputStreamMode, getStreamingPorts } from "../task/StreamTypes"; +import { + DEFAULT_BINARY_HIGH_WATER_BYTES, + edgeNeedsAccumulation, + getOutputStreamMode, + getStreamingPorts, +} from "../task/StreamTypes"; import type { TaskInput } from "../task/TaskTypes"; import { TaskStatus } from "../task/TaskTypes"; import { Dataflow, DATAFLOW_ALL_PORTS } from "./Dataflow"; @@ -354,15 +359,31 @@ export class StreamPump { task: ITask, binaryPortId: string | undefined, sink: (chunks: AsyncIterable) => Promise, - signal?: AbortSignal - ): { promise: Promise; detach: () => void } { + signal?: AbortSignal, + options?: { readonly highWaterMarkBytes?: number } + ): { + promise: Promise; + detach: () => void; + backpressure: () => Promise; + } { const queue: Uint8Array[] = []; + let bufferedBytes = 0; + const highWaterMarkBytes = Math.max( + 1, + options?.highWaterMarkBytes ?? DEFAULT_BINARY_HIGH_WATER_BYTES + ); let done = false; - let notify: (() => void) | undefined; + let chunkNotify: (() => void) | undefined; + let drainNotify: (() => void) | undefined; - const wake = () => { - const n = notify; - notify = undefined; + const wakeChunk = () => { + const n = chunkNotify; + chunkNotify = undefined; + n?.(); + }; + const wakeDrain = () => { + const n = drainNotify; + drainNotify = undefined; n?.(); }; @@ -370,13 +391,17 @@ export class StreamPump { if (event.type === "binary-delta") { if (binaryPortId === undefined || event.port === binaryPortId) { queue.push(event.binaryDelta); - wake(); + bufferedBytes += event.binaryDelta.byteLength; + wakeChunk(); } } }; const onEnd = () => { done = true; - wake(); + wakeChunk(); + // Release any cooperative-backpressure awaiter that was parked at + // high-water; without this an abort-while-parked would leak. + wakeDrain(); }; // Abort/error termination: StreamProcessor never emits `stream_end` on these // paths, so without this the iterable would await forever. Terminate the @@ -384,7 +409,8 @@ export class StreamPump { // promise settles — the source's own abort/error already surfaces to the run. const onTerminate = () => { done = true; - wake(); + wakeChunk(); + wakeDrain(); }; task.on("stream_chunk", onChunk); @@ -407,21 +433,48 @@ export class StreamPump { async function* chunkIterable(): AsyncIterable { while (true) { while (queue.length > 0) { - yield queue.shift()!; + const chunk = queue.shift()!; + bufferedBytes -= chunk.byteLength; + if (drainNotify && bufferedBytes < highWaterMarkBytes) wakeDrain(); + yield chunk; } if (done) return; await new Promise((resolve) => { - notify = resolve; + chunkNotify = resolve; }); } } + /** + * Cooperative backpressure hook. Resolves immediately while buffered + * bytes stay under the high-water mark; otherwise parks until the + * consumer drains the queue, or the stream is closed. + * + * The EventEmitter delivery path cannot apply mandatory backpressure + * (the listener fires synchronously), so this is opt-in: a task can + * `await ctx.binaryBackpressure()` between large yields. Tasks that + * never call it pay nothing. + */ + const backpressure = (): Promise => { + if (done) return Promise.resolve(); + if (bufferedBytes < highWaterMarkBytes) return Promise.resolve(); + return new Promise((resolve) => { + const prev = drainNotify; + drainNotify = prev + ? () => { + prev(); + resolve(); + } + : resolve; + }); + }; + // Discard the sink's return value (helper signals completion only; callers // wanting a CacheRef should hold the sink-returning promise themselves). const promise = sink(chunkIterable()) .finally(detach) .then(() => undefined); - return { promise, detach }; + return { promise, detach, backpressure }; } /** diff --git a/packages/task-graph/src/task/CacheCoordinator.ts b/packages/task-graph/src/task/CacheCoordinator.ts index c72c2c1a0..5a3d4b564 100644 --- a/packages/task-graph/src/task/CacheCoordinator.ts +++ b/packages/task-graph/src/task/CacheCoordinator.ts @@ -8,14 +8,14 @@ import { getPortCodec } from "@workglow/util"; import type { DataPortSchema } from "@workglow/util/schema"; import { type CachePolicy, isPolicyCached, isPolicyPrivate } from "../cache/CachePolicy"; import type { CacheRef } from "../cache/CacheRef"; -import { isCacheRef } from "../cache/CacheRef"; +import { isCacheRef, makeCacheRef } from "../cache/CacheRef"; import type { CacheRegistry } from "../cache/CacheRegistry"; import { RunPrivateCacheRepo } from "../cache/RunPrivateCacheRepo"; import type { TaskOutputRepository } from "../storage/TaskOutputRepository"; -import type { BinaryRefSink } from "./StreamProcessor"; -import { getBinaryPortFormat, getBinaryPortId, getStreamingPorts } from "./StreamTypes"; import type { ITask } from "./ITask"; +import type { BinaryRefSink } from "./StreamProcessor"; import type { StreamEvent } from "./StreamTypes"; +import { assertBinaryFormat, getBinaryPortId, getStreamingPorts } from "./StreamTypes"; import { Task } from "./Task"; import type { TaskRunContext } from "./TaskRunContext"; import type { TaskInput, TaskOutput } from "./TaskTypes"; @@ -235,8 +235,13 @@ export class CacheCoordinator - cache.saveOutputStream!(taskType, keyInputs, chunks, {}); + // Re-wrap the backing's CacheRef so legacy `saveOutputStream` implementations + // that pre-date the `kind` brand still produce a discriminator-bearing ref. + // Branded refs pass through unchanged (preserving size/mime hints). + const sink: BinaryRefSink = async (chunks) => { + const raw = await cache.saveOutputStream!(taskType, keyInputs, chunks, {}); + return isCacheRef(raw) ? raw : makeCacheRef(raw); + }; return new Map([[port, sink]]); } @@ -286,7 +291,7 @@ export class CacheCoordinator= referenceThresholdBytes) return undefined; const blob = await cache.getOutputByRef!(value); if (blob === undefined) return undefined; - const format = getBinaryPortFormat(outputSchema, port); + const format = assertBinaryFormat(outputSchema, port); const inlined = format === "binary" ? await blob.arrayBuffer() : blob; return { port, inlined }; }) diff --git a/packages/task-graph/src/task/ITask.ts b/packages/task-graph/src/task/ITask.ts index d1b93c7f1..84df24b3b 100644 --- a/packages/task-graph/src/task/ITask.ts +++ b/packages/task-graph/src/task/ITask.ts @@ -68,6 +68,16 @@ export interface IExecuteContext { * did not provide it. */ resourceScope?: ResourceScope; + /** + * Optional cooperative backpressure hook for streaming tasks that emit very + * large binary outputs by direct event emission (rather than through the + * StreamProcessor's `await router.push(...)` path). Tasks may `await` this + * between yields/emits to give downstream sinks a chance to drain. + * + * Defaults to a no-op when the runtime does not install a real backpressure + * source — tasks can call it unconditionally without paying a cost. + */ + binaryBackpressure?: () => Promise; } export type IExecutePreviewContext = Pick; @@ -125,6 +135,18 @@ export interface IRunConfig { */ referenceThresholdBytes?: number; + /** + * High-water mark (bytes) for the streaming runtime's per-port binary + * router buffer. When the buffered (un-consumed) byte total reaches or + * exceeds this threshold, the producer (`executeStream`) is parked between + * `binary-delta` yields until the cache sink drains the buffer back below + * the mark. Bounds peak memory for fast-producer / slow-sink scenarios. + * + * Defaults to {@link DEFAULT_BINARY_HIGH_WATER_BYTES} (8 MiB) when omitted + * or set to a non-positive value. + */ + binaryHighWaterBytes?: number; + /** * Optional callback invoked whenever a task's progress changes during execution. * @param task - The task whose progress changed. diff --git a/packages/task-graph/src/task/StreamProcessor.ts b/packages/task-graph/src/task/StreamProcessor.ts index a121c17b4..0ed04ab43 100644 --- a/packages/task-graph/src/task/StreamProcessor.ts +++ b/packages/task-graph/src/task/StreamProcessor.ts @@ -10,7 +10,8 @@ import type { Taskish } from "../task-graph/Conversions"; import type { ITask } from "./ITask"; import type { StreamEvent, StreamMode } from "./StreamTypes"; import { - getBinaryPortFormat, + assertBinaryFormat, + DEFAULT_BINARY_HIGH_WATER_BYTES, getOutputStreamMode, getStreamingPorts, materializeBinary, @@ -18,6 +19,7 @@ import { import { TaskAbortedError, TaskError } from "./TaskError"; import type { TaskRunContext } from "./TaskRunContext"; import type { TaskInput, TaskOutput } from "./TaskTypes"; +import { TaskStatus } from "./TaskTypes"; /** * Consumer for a port's binary-delta stream. The processor exposes chunks as @@ -29,7 +31,6 @@ import type { TaskInput, TaskOutput } from "./TaskTypes"; * once it knows the cache key. */ export type BinaryRefSink = (chunks: AsyncIterable) => Promise; -import { TaskStatus } from "./TaskTypes"; /** * Per-call run-state inputs shared by StreamProcessor.run. Bundles facade @@ -60,6 +61,14 @@ export interface StreamProcessorDeps { * Ports without a sink follow the normal accumulation path. */ readonly binaryRefSinks?: ReadonlyMap; + /** + * High-water mark (bytes) for the per-port binary stream router buffer. When + * the buffered (un-consumed) byte total reaches or exceeds this value, + * `BinaryStreamRouter.push()` returns a Promise that resolves only after the + * consumer drains the buffer back below the mark. Defaults to + * {@link DEFAULT_BINARY_HIGH_WATER_BYTES} when omitted. + */ + readonly binaryHighWaterBytes?: number; } /** @@ -99,14 +108,16 @@ export class StreamProcessor const accumulatedObjects = ctx.shouldAccumulate ? new Map | unknown[]>() : undefined; - const accumulatedBinary = ctx.shouldAccumulate - ? new Map() - : undefined; + const accumulatedBinary = ctx.shouldAccumulate ? new Map() : undefined; // Per-port routers: lazily created on the first binary-delta whose port has // a sink in `deps.binaryRefSinks`. Routes chunks to the sink instead of // accumulating in memory; at finish, awaits the sink's returned CacheRef // and writes it into the output at the port slot. const sinks = deps.binaryRefSinks; + const highWaterMark = + deps.binaryHighWaterBytes !== undefined && deps.binaryHighWaterBytes > 0 + ? deps.binaryHighWaterBytes + : DEFAULT_BINARY_HIGH_WATER_BYTES; const routers = new Map(); const ensureRouter = (port: string): BinaryStreamRouter | undefined => { if (!sinks) return undefined; @@ -114,7 +125,7 @@ export class StreamProcessor if (!sink) return undefined; let r = routers.get(port); if (!r) { - r = new BinaryStreamRouter(sink); + r = new BinaryStreamRouter(sink, highWaterMark); routers.set(port, r); } return r; @@ -125,6 +136,22 @@ export class StreamProcessor this.task.emit("stream_start"); + // Cooperative backpressure hook for executeStream() implementations that + // emit through a side channel (not StreamProcessor's awaited `push`). When + // any port has a router (we'd be applying byte-bounded backpressure on the + // direct `binary-delta` path anyway), `await ctx.binaryBackpressure()` + // waits until ALL active routers are at-or-below their high-water mark. + // Without a router this is a cheap no-op. + const binaryBackpressure = async (): Promise => { + if (routers.size === 0) return; + const waits: Promise[] = []; + for (const r of routers.values()) { + if (r._bufferedBytes >= r._highWaterMarkBytes) waits.push(r._awaitDrain()); + } + if (waits.length === 0) return; + await Promise.all(waits); + }; + const stream = this.task.executeStream!(input, { signal: ctx.abortController.signal, updateProgress: deps.onProgress, @@ -132,213 +159,218 @@ export class StreamProcessor registry: deps.registry, resourceScope: deps.resourceScope, inputStreams: deps.inputStreams, + binaryBackpressure, }); try { - for await (const event of stream) { - // For snapshot events, update runOutputData BEFORE emitting stream_chunk - // so listeners see the latest snapshot when they handle the event. - if (event.type === "snapshot") { - this.task.runOutputData = event.data as Output; - } - - switch (event.type) { - case "phase": { - // Phase events are metadata: emit for observability, translate to a - // progress event with optional progress + message, do NOT mutate - // accumulators or runOutputData, do NOT flip status to STREAMING. - this.task.emit("stream_chunk", event as StreamEvent); - await deps.onProgress(event.progress, event.message); - break; + for await (const event of stream) { + // For snapshot events, update runOutputData BEFORE emitting stream_chunk + // so listeners see the latest snapshot when they handle the event. + if (event.type === "snapshot") { + this.task.runOutputData = event.data as Output; } - case "text-delta": { - if (!streamingStarted) { - streamingStarted = true; - this.task.status = TaskStatus.STREAMING; - this.task.emit("status", this.task.status); - } - if (accumulated) { - accumulated.set(event.port, (accumulated.get(event.port) ?? "") + event.textDelta); + + switch (event.type) { + case "phase": { + // Phase events are metadata: emit for observability, translate to a + // progress event with optional progress + message, do NOT mutate + // accumulators or runOutputData, do NOT flip status to STREAMING. + this.task.emit("stream_chunk", event as StreamEvent); + await deps.onProgress(event.progress, event.message); + break; } - this.task.emit("stream_chunk", event as StreamEvent); - break; - } - case "object-delta": { - if (!streamingStarted) { - streamingStarted = true; - this.task.status = TaskStatus.STREAMING; - this.task.emit("status", this.task.status); + case "text-delta": { + if (!streamingStarted) { + streamingStarted = true; + this.task.status = TaskStatus.STREAMING; + this.task.emit("status", this.task.status); + } + if (accumulated) { + accumulated.set(event.port, (accumulated.get(event.port) ?? "") + event.textDelta); + } + this.task.emit("stream_chunk", event as StreamEvent); + break; } - if (accumulatedObjects) { - const existing = accumulatedObjects.get(event.port); - if (Array.isArray(event.objectDelta)) { - // Array delta: upsert items by `id` into accumulated array - const arr: unknown[] = Array.isArray(existing) ? [...existing] : []; - for (const item of event.objectDelta) { - const itemObj = item as Record; - if (itemObj && typeof itemObj === "object" && "id" in itemObj) { - const idx = arr.findIndex( - (e) => (e as Record).id === itemObj.id - ); - if (idx >= 0) arr[idx] = item; - else arr.push(item); - } else { - arr.push(item); + case "object-delta": { + if (!streamingStarted) { + streamingStarted = true; + this.task.status = TaskStatus.STREAMING; + this.task.emit("status", this.task.status); + } + if (accumulatedObjects) { + const existing = accumulatedObjects.get(event.port); + if (Array.isArray(event.objectDelta)) { + // Array delta: upsert items by `id` into accumulated array + const arr: unknown[] = Array.isArray(existing) ? [...existing] : []; + for (const item of event.objectDelta) { + const itemObj = item as Record; + if (itemObj && typeof itemObj === "object" && "id" in itemObj) { + const idx = arr.findIndex( + (e) => (e as Record).id === itemObj.id + ); + if (idx >= 0) arr[idx] = item; + else arr.push(item); + } else { + arr.push(item); + } } + accumulatedObjects.set(event.port, arr); + } else { + // Non-array (e.g. structured generation): replace semantics + accumulatedObjects.set(event.port, event.objectDelta); } - accumulatedObjects.set(event.port, arr); - } else { - // Non-array (e.g. structured generation): replace semantics - accumulatedObjects.set(event.port, event.objectDelta); } + // Update runOutputData with accumulated state so listeners see growing state + this.task.runOutputData = { + ...this.task.runOutputData, + [event.port]: accumulatedObjects?.get(event.port) ?? event.objectDelta, + } as Output; + this.task.emit("stream_chunk", event as StreamEvent); + break; } - // Update runOutputData with accumulated state so listeners see growing state - this.task.runOutputData = { - ...this.task.runOutputData, - [event.port]: accumulatedObjects?.get(event.port) ?? event.objectDelta, - } as Output; - this.task.emit("stream_chunk", event as StreamEvent); - break; - } - case "binary-delta": { - if (!streamingStarted) { - streamingStarted = true; - this.task.status = TaskStatus.STREAMING; - this.task.emit("status", this.task.status); - } - // Tee: when both a router AND an accumulator exist - // for this port (graph context where the cache can stream but a - // downstream edge needs the materialized value), push to BOTH — - // router writes to the cache for the small ref-bearing Output, - // accumulator drives the enriched finish event so edge consumers - // still receive a Blob/ArrayBuffer. - const router = ensureRouter(event.port); - if (router) router.push(event.binaryDelta); - if (accumulatedBinary) { - const arr = accumulatedBinary.get(event.port) ?? []; - arr.push(event.binaryDelta); - accumulatedBinary.set(event.port, arr); + case "binary-delta": { + if (!streamingStarted) { + streamingStarted = true; + this.task.status = TaskStatus.STREAMING; + this.task.emit("status", this.task.status); + } + // Tee: when both a router AND an accumulator exist + // for this port (graph context where the cache can stream but a + // downstream edge needs the materialized value), push to BOTH — + // router writes to the cache for the small ref-bearing Output, + // accumulator drives the enriched finish event so edge consumers + // still receive a Blob/ArrayBuffer. + // `await router.push(...)` here is where byte-bounded backpressure + // takes effect: the producer (executeStream) parks until the sink + // drains the router buffer back under the high-water mark, or + // until the router is closed (abort/error path). + const router = ensureRouter(event.port); + if (router) await router.push(event.binaryDelta); + if (accumulatedBinary) { + const arr = accumulatedBinary.get(event.port) ?? []; + arr.push(event.binaryDelta); + accumulatedBinary.set(event.port, arr); + } + this.task.emit("stream_chunk", event as StreamEvent); + break; } - this.task.emit("stream_chunk", event as StreamEvent); - break; - } - case "snapshot": { - if (!streamingStarted) { - streamingStarted = true; - this.task.status = TaskStatus.STREAMING; - this.task.emit("status", this.task.status); + case "snapshot": { + if (!streamingStarted) { + streamingStarted = true; + this.task.status = TaskStatus.STREAMING; + this.task.emit("status", this.task.status); + } + this.task.emit("stream_chunk", event as StreamEvent); + break; } - this.task.emit("stream_chunk", event as StreamEvent); - break; - } - case "finish": { - const hasEnrichment = - accumulated !== undefined || - accumulatedObjects !== undefined || - accumulatedBinary !== undefined || - routers.size > 0; - if (hasEnrichment) { - // Emit an enriched finish event: merge accumulated deltas into - // the finish payload so downstream dataflows get complete port data - // without needing to re-accumulate themselves. - const explicitPayload = (event.data || {}) as Record; - const merged: Record = { ...explicitPayload }; - if (accumulated) { - for (const [port, text] of accumulated) { - if (text.length > 0) merged[port] = text; + case "finish": { + const hasEnrichment = + accumulated !== undefined || + accumulatedObjects !== undefined || + accumulatedBinary !== undefined || + routers.size > 0; + if (hasEnrichment) { + // Emit an enriched finish event: merge accumulated deltas into + // the finish payload so downstream dataflows get complete port data + // without needing to re-accumulate themselves. + const explicitPayload = (event.data || {}) as Record; + const merged: Record = { ...explicitPayload }; + if (accumulated) { + for (const [port, text] of accumulated) { + if (text.length > 0) merged[port] = text; + } } - } - if (accumulatedObjects) { - for (const [port, obj] of accumulatedObjects) { - merged[port] = obj; + if (accumulatedObjects) { + for (const [port, obj] of accumulatedObjects) { + merged[port] = obj; + } } - } - if (accumulatedBinary) { - const outSchema = this.task.outputSchema(); - for (const [port, chunks] of accumulatedBinary) { - // Explicit binary finish payload wins. (Unlike text/object - // deltas above, which overwrite event.data, binary yields to - // an explicit payload — it's a whole artifact, not a partial.) - if (port in explicitPayload) continue; - const format = getBinaryPortFormat(outSchema, port); - merged[port] = materializeBinary(chunks, format); + if (accumulatedBinary) { + const outSchema = this.task.outputSchema(); + for (const [port, chunks] of accumulatedBinary) { + // Explicit binary finish payload wins. (Unlike text/object + // deltas above, which overwrite event.data, binary yields to + // an explicit payload — it's a whole artifact, not a partial.) + if (port in explicitPayload) continue; + const format = assertBinaryFormat(outSchema, port); + merged[port] = materializeBinary(chunks, format); + } } - } - // Close routers and collect refs. Explicit binary finish payload - // still wins for the OUTPUT slot (artifact precedence); the - // router's CacheRef is discarded in that case but the cache - // write already happened. - for (const router of routers.values()) router.end(); - const refs = new Map(); - for (const [port, router] of routers) { - if (port in explicitPayload) { - // Drain the promise so the sink doesn't leak; ignore the ref. - router.ref().catch(() => {}); - continue; + // Close routers and collect refs. Explicit binary finish payload + // still wins for the OUTPUT slot (artifact precedence); the + // router's CacheRef is discarded in that case but the cache + // write already happened. + for (const router of routers.values()) router.end(); + const refs = new Map(); + for (const [port, router] of routers) { + if (port in explicitPayload) { + // Drain the promise so the sink doesn't leak; ignore the ref. + router.ref().catch(() => {}); + continue; + } + refs.set(port, await router.ref()); } - refs.set(port, await router.ref()); - } - // For replace-mode streams, finish carries data: {} by convention. - // Fall back to the last snapshot (runOutputData) so the final output - // is not silently cleared when the finish payload is empty — - // overlaying router refs on top so cache-written bytes are not - // orphaned (the ref still lands in the OUTPUT slot). - if (streamMode === "replace" && Object.keys(merged).length === 0) { - const lastSnapshot = this.task.runOutputData; - if (lastSnapshot && Object.keys(lastSnapshot).length > 0) { - const snapshotWithRefs: Record = { ...lastSnapshot }; - for (const [port, ref] of refs) snapshotWithRefs[port] = ref; - finalOutput = snapshotWithRefs as Output; - this.task.emit("stream_chunk", { - type: "finish", - data: lastSnapshot, - } as StreamEvent); - break; + // For replace-mode streams, finish carries data: {} by convention. + // Fall back to the last snapshot (runOutputData) so the final output + // is not silently cleared when the finish payload is empty — + // overlaying router refs on top so cache-written bytes are not + // orphaned (the ref still lands in the OUTPUT slot). + if (streamMode === "replace" && Object.keys(merged).length === 0) { + const lastSnapshot = this.task.runOutputData; + if (lastSnapshot && Object.keys(lastSnapshot).length > 0) { + const snapshotWithRefs: Record = { ...lastSnapshot }; + for (const [port, ref] of refs) snapshotWithRefs[port] = ref; + finalOutput = snapshotWithRefs as Output; + this.task.emit("stream_chunk", { + type: "finish", + data: lastSnapshot, + } as StreamEvent); + break; + } + } + // The emitted finish event always carries the materialized payload + // (from accumulators) so edge consumers see Blob/ArrayBuffer. + // finalOutput diverges only when a router produced a ref for a + // port that wasn't already pinned by an explicit payload — that + // ref takes the slot in the return value so the queue/cache row + // stays small (the tee path). + this.task.emit("stream_chunk", { type: "finish", data: merged } as StreamEvent); + if (refs.size === 0) { + finalOutput = merged as unknown as Output; + } else { + const finalMerged: Record = { ...merged }; + for (const [port, ref] of refs) finalMerged[port] = ref; + finalOutput = finalMerged as unknown as Output; } - } - // The emitted finish event always carries the materialized payload - // (from accumulators) so edge consumers see Blob/ArrayBuffer. - // finalOutput diverges only when a router produced a ref for a - // port that wasn't already pinned by an explicit payload — that - // ref takes the slot in the return value so the queue/cache row - // stays small (the tee path). - this.task.emit("stream_chunk", { type: "finish", data: merged } as StreamEvent); - if (refs.size === 0) { - finalOutput = merged as unknown as Output; } else { - const finalMerged: Record = { ...merged }; - for (const [port, ref] of refs) finalMerged[port] = ref; - finalOutput = finalMerged as unknown as Output; - } - } else { - // No accumulation. For replace-mode streams the provider's finish - // event carries `data: {}` by convention — the snapshots already - // delivered the value, so the finish payload is intentionally - // empty. Fall back to `runOutputData` (set on every snapshot above) - // so we don't clobber the last snapshot with an empty object. This - // mirrors the same fallback in the accumulation branch. - const finishData = (event.data ?? {}) as Record; - if (streamMode === "replace" && Object.keys(finishData).length === 0) { - const lastSnapshot = this.task.runOutputData; - if (lastSnapshot && Object.keys(lastSnapshot).length > 0) { - finalOutput = lastSnapshot as Output; - this.task.emit("stream_chunk", { - type: "finish", - data: lastSnapshot, - } as StreamEvent); - break; + // No accumulation. For replace-mode streams the provider's finish + // event carries `data: {}` by convention — the snapshots already + // delivered the value, so the finish payload is intentionally + // empty. Fall back to `runOutputData` (set on every snapshot above) + // so we don't clobber the last snapshot with an empty object. This + // mirrors the same fallback in the accumulation branch. + const finishData = (event.data ?? {}) as Record; + if (streamMode === "replace" && Object.keys(finishData).length === 0) { + const lastSnapshot = this.task.runOutputData; + if (lastSnapshot && Object.keys(lastSnapshot).length > 0) { + finalOutput = lastSnapshot as Output; + this.task.emit("stream_chunk", { + type: "finish", + data: lastSnapshot, + } as StreamEvent); + break; + } } + finalOutput = event.data as Output; + this.task.emit("stream_chunk", event as StreamEvent); } - finalOutput = event.data as Output; - this.task.emit("stream_chunk", event as StreamEvent); + break; + } + case "error": { + throw event.error; } - break; - } - case "error": { - throw event.error; } } - } } catch (err) { // Surface the error to any in-flight router sinks so they reject // (rather than waiting forever on the producer). The original error is @@ -379,23 +411,27 @@ export class StreamProcessor * resolves). `fail(err)` causes the iterable to throw on the next read * (refPromise rejects). `end()` and `fail()` are idempotent. * - * Backpressure: there is none — the producer (`executeStream`) writes into - * `buffer` synchronously; if the sink consumes more slowly than the producer - * emits chunks, the buffer grows unbounded. This is acceptable for cache - * backings whose write throughput meets or exceeds the upstream source - * (the common case: cache is a local SSD/memory FS; source is bounded by - * network or compute). For genuinely slow backings (remote object stores, - * throttled FS), wrap the sink in a chunked-uploader that applies its own - * pacing — there is no signal we can send back into `binary-delta`. + * Backpressure: byte-bounded. `push()` returns a Promise; the producer is + * resolved immediately while the buffered (un-consumed) byte total stays + * below `highWaterMarkBytes`, and parks until the consumer drains the + * buffer back under the mark once the threshold is reached. `end()` and + * `fail()` BOTH release any parked producer so an abort mid-park does not + * leak the `push()` promise. */ -class BinaryStreamRouter { +export class BinaryStreamRouter { private readonly buffer: Uint8Array[] = []; + private bufferedBytes = 0; private finished = false; private failure: Error | undefined; - private notify: (() => void) | undefined; + /** Resolver for the consumer side (iterable awaiting next chunk). */ + private chunkNotify: (() => void) | undefined; + /** Resolver for the producer side parked waiting for drain. */ + private drainNotify: (() => void) | undefined; private readonly refPromise: Promise; + private readonly highWaterMarkBytes: number; - constructor(sink: BinaryRefSink) { + constructor(sink: BinaryRefSink, highWaterMarkBytes: number) { + this.highWaterMarkBytes = Math.max(1, highWaterMarkBytes); this.refPromise = sink(this.iterable()); // Observe rejection so an unawaited refPromise (e.g. after fail() in an // error path) doesn't surface as an unhandled rejection. Subsequent @@ -403,46 +439,111 @@ class BinaryStreamRouter { this.refPromise.catch(() => {}); } - push(chunk: Uint8Array): void { - if (this.finished) return; + /** + * Buffer one chunk and return a Promise the caller must await. The promise + * resolves immediately when buffered bytes remain under the high-water + * mark, and otherwise parks until the consumer drains the buffer (or until + * `end()` / `fail()` releases all parked callers). + */ + push(chunk: Uint8Array): Promise { + if (this.finished) return Promise.resolve(); this.buffer.push(chunk); - this.wake(); + this.bufferedBytes += chunk.byteLength; + this.wakeChunk(); + if (this.bufferedBytes < this.highWaterMarkBytes) return Promise.resolve(); + return new Promise((res) => { + // Chain resolvers so a long park doesn't lose earlier waiters. + const prev = this.drainNotify; + this.drainNotify = prev + ? () => { + prev(); + res(); + } + : res; + }); } end(): void { if (this.finished) return; this.finished = true; - this.wake(); + this.wakeChunk(); + // Release any producer parked at the high-water mark — abort mid-stream + // would otherwise orphan the parked Promise. + this.wakeDrain(); } fail(err: Error): void { if (this.finished) return; this.failure = err; this.finished = true; - this.wake(); + this.wakeChunk(); + this.wakeDrain(); } ref(): Promise { return this.refPromise; } - private wake(): void { - const n = this.notify; - this.notify = undefined; + /** @internal Test hook: current buffered byte count (consumer-unread). */ + public get _bufferedBytes(): number { + return this.bufferedBytes; + } + + /** @internal Test hook: high-water mark in effect. */ + public get _highWaterMarkBytes(): number { + return this.highWaterMarkBytes; + } + + /** + * @internal Used by {@link IExecuteContext.binaryBackpressure} so a task + * emitting via a side channel can park until the consumer drains. Resolves + * immediately when the buffer is already under the mark or the router has + * been closed. + */ + public _awaitDrain(): Promise { + if (this.finished) return Promise.resolve(); + if (this.bufferedBytes < this.highWaterMarkBytes) return Promise.resolve(); + return new Promise((res) => { + const prev = this.drainNotify; + this.drainNotify = prev + ? () => { + prev(); + res(); + } + : res; + }); + } + + private wakeChunk(): void { + const n = this.chunkNotify; + this.chunkNotify = undefined; + n?.(); + } + + private wakeDrain(): void { + const n = this.drainNotify; + this.drainNotify = undefined; n?.(); } private async *iterable(): AsyncIterable { while (true) { while (this.buffer.length > 0) { - yield this.buffer.shift()!; + const chunk = this.buffer.shift()!; + this.bufferedBytes -= chunk.byteLength; + // Wake any parked producer once we drop back below the mark. We + // resolve as soon as we cross the threshold rather than waiting for + // the buffer to drain fully — that keeps the producer pipelined. + if (this.drainNotify && this.bufferedBytes < this.highWaterMarkBytes) { + this.wakeDrain(); + } + yield chunk; } if (this.failure) throw this.failure; if (this.finished) return; await new Promise((res) => { - this.notify = res; + this.chunkNotify = res; }); } } } - diff --git a/packages/task-graph/src/task/StreamTypes.ts b/packages/task-graph/src/task/StreamTypes.ts index 2220a370e..8bff267bf 100644 --- a/packages/task-graph/src/task/StreamTypes.ts +++ b/packages/task-graph/src/task/StreamTypes.ts @@ -290,10 +290,32 @@ export function getBinaryPortId(schema: DataPortSchema): string | undefined { return undefined; } +/** + * Canonical vocabulary for the `format` annotation on a binary streaming output + * port. `"blob"` materializes chunks into a `Blob` (the default); `"binary"` + * materializes them into an `ArrayBuffer`. Any other value is rejected at + * registration time (see {@link assertBinaryFormat}) so a typo like `"Blob"` + * cannot silently fall through to the ArrayBuffer branch. + */ +export type BinaryFormat = "blob" | "binary"; + +/** + * Default high-water mark for the binary-stream router's producer buffer, in + * bytes. When the buffered (un-consumed) byte total reaches this threshold the + * producer awaits a drain signal from the consumer before pushing further + * chunks; below the threshold the producer is allowed to run free. 8 MiB lets + * even fast producers race ahead by a few chunks without stalling, while + * bounding worst-case memory growth when the sink (cache, disk, network) + * cannot keep up. Callers can override per-run via + * `IRunConfig.binaryHighWaterBytes`. + */ +export const DEFAULT_BINARY_HIGH_WATER_BYTES = 8 * 1024 * 1024; + /** * Reads the `format` annotation of a single output port from the task's output - * schema. Used to decide whether accumulated binary chunks materialize into a - * `Blob` (`format: "blob"` or absent) or an `ArrayBuffer` (`format: "binary"`). + * schema. Returns the raw string (or `undefined`) — callers needing the + * canonical {@link BinaryFormat} vocabulary should go through + * {@link assertBinaryFormat}, which rejects unknown values. */ export function getBinaryPortFormat(schema: DataPortSchema, port: string): string | undefined { if (typeof schema === "boolean") return undefined; @@ -302,27 +324,43 @@ export function getBinaryPortFormat(schema: DataPortSchema, port: string): strin return prop.format as string | undefined; } +/** + * Resolves the `format` annotation on a binary streaming port to a canonical + * {@link BinaryFormat}. `undefined` and `"blob"` both resolve to `"blob"`; + * `"binary"` resolves to `"binary"`. Anything else throws — a casing typo such + * as `"Blob"` or a leftover legacy value would otherwise be silently coerced + * to one branch and produce the wrong runtime type, so this is checked at + * task-registration time and again on the streaming hot paths. + */ +export function assertBinaryFormat(schema: DataPortSchema, port: string): BinaryFormat { + const f = getBinaryPortFormat(schema, port); + if (f === undefined || f === "blob") return "blob"; + if (f === "binary") return "binary"; + throw new Error( + `Port "${port}" has x-stream:"binary" but format:"${f}". Allowed: "blob" | "binary".` + ); +} + /** * Materializes ordered binary chunks into the value type declared by the - * output port's schema `format`: - * - `"binary"` → `ArrayBuffer` - * - `"blob"` (or absent) → `Blob` (the default) + * output port's canonical {@link BinaryFormat}: + * - `"blob"` → `Blob` (the default) + * - `"binary"` → `ArrayBuffer` * * Chunks are concatenated in arrival order. Callers MUST pass chunks in the - * order they were emitted. + * order they were emitted, and MUST resolve `format` through + * {@link assertBinaryFormat} so unknown values are rejected at registration + * rather than reinterpreted here. * * @param chunks - Ordered binary chunks to concatenate - * @param format - The output port's schema `format` (e.g. `"binary"` or `"blob"`) + * @param format - Canonical binary format selector * @returns The materialized `Blob` or `ArrayBuffer` */ export function materializeBinary( chunks: readonly Uint8Array[], - format: string | undefined + format: BinaryFormat ): Blob | ArrayBuffer { - if (format === "blob" || format === undefined) { - return new Blob(chunks as unknown as BlobPart[]); - } - // format === "binary" (and any other non-blob value) → ArrayBuffer + if (format === "blob") return new Blob(chunks as unknown as BlobPart[]); let total = 0; for (const c of chunks) total += c.byteLength; const merged = new Uint8Array(total); diff --git a/packages/task-graph/src/task/TaskRegistry.ts b/packages/task-graph/src/task/TaskRegistry.ts index 52307e813..c1f343fbe 100644 --- a/packages/task-graph/src/task/TaskRegistry.ts +++ b/packages/task-graph/src/task/TaskRegistry.ts @@ -14,6 +14,7 @@ import { } from "@workglow/util"; import { validateSchema } from "@workglow/util/schema"; import type { ITaskConstructor } from "./ITask"; +import { assertBinaryFormat, getStreamingPorts } from "./StreamTypes"; // eslint-disable-next-line @typescript-eslint/no-explicit-any type AnyTaskConstructor = ITaskConstructor; @@ -41,12 +42,32 @@ function registerTask(baseClass: AnyTaskConstructor): void { `Task type "${baseClass.type}" is already registered. Unregister it first to replace.` ); } + + // Validate every binary streaming output port's `format` against the + // canonical {@link BinaryFormat} vocabulary BEFORE adding to the registry. + // A typo like `format: "Blob"` would otherwise silently coerce to the + // ArrayBuffer branch in `materializeBinary`, producing the wrong runtime + // type for every consumer of that port. Fail at registration so the + // misconfiguration surfaces near the task definition site. + const outputSchema = baseClass.outputSchema(); + for (const { port, mode } of getStreamingPorts(outputSchema)) { + if (mode !== "binary") continue; + try { + assertBinaryFormat(outputSchema, port); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + throw new Error( + `Cannot register task "${baseClass.type}": invalid binary stream port. ${message}` + ); + } + } + taskConstructors.set(baseClass.type, baseClass); // Validate schemas at registration time (soft — warn only, don't throw) const schemas = [ { name: "inputSchema", schema: baseClass.inputSchema() }, - { name: "outputSchema", schema: baseClass.outputSchema() }, + { name: "outputSchema", schema: outputSchema }, ] as const; for (const { name, schema } of schemas) { diff --git a/packages/task-graph/src/task/TaskRunner.ts b/packages/task-graph/src/task/TaskRunner.ts index 68a8ff88a..8de5dcda3 100644 --- a/packages/task-graph/src/task/TaskRunner.ts +++ b/packages/task-graph/src/task/TaskRunner.ts @@ -270,6 +270,8 @@ export class TaskRunner< ) : undefined; + const binaryHighWaterBytes = + config.binaryHighWaterBytes ?? this.task.runConfig.binaryHighWaterBytes; outputs = isStreamable ? await this.streamProcessor.run(inputs, ctx, { registry: this.registry, @@ -278,6 +280,7 @@ export class TaskRunner< onProgress: this.handleProgress.bind(this), own: this.own, binaryRefSinks, + binaryHighWaterBytes, }) : await this.executeTask(inputs, ctx); diff --git a/packages/test/src/test/task-graph/CacheRef.test.ts b/packages/test/src/test/task-graph/CacheRef.test.ts index b736d7b51..430dd8286 100644 --- a/packages/test/src/test/task-graph/CacheRef.test.ts +++ b/packages/test/src/test/task-graph/CacheRef.test.ts @@ -4,30 +4,42 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { describe, expect, it } from "vitest"; +import type { CacheRef, IRunConfig } from "@workglow/task-graph"; import { + CACHE_REF_KIND, isCacheRef, + makeCacheRef, REFERENCE_THRESHOLD_BYTES_DEFAULT, resolveReferenceThreshold, } from "@workglow/task-graph"; -import type { CacheRef, IRunConfig } from "@workglow/task-graph"; +import { describe, expect, it } from "vitest"; describe("isCacheRef", () => { - it("accepts a minimal ref carrying only $ref", () => { - const ref: CacheRef = { $ref: "cache://k1" }; + it("accepts a minimal branded ref carrying only $ref", () => { + const ref: CacheRef = makeCacheRef({ $ref: "cache://k1" }); expect(isCacheRef(ref)).toBe(true); }); - it("accepts a ref carrying size and mime hints", () => { - const ref: CacheRef = { $ref: "cache://k2", size: 1024, mime: "audio/wav" }; + it("accepts a branded ref carrying size and mime hints", () => { + const ref: CacheRef = makeCacheRef({ $ref: "cache://k2", size: 1024, mime: "audio/wav" }); expect(isCacheRef(ref)).toBe(true); }); it("rejects values without a string $ref", () => { - expect(isCacheRef({})).toBe(false); - expect(isCacheRef({ ref: "cache://k" })).toBe(false); - expect(isCacheRef({ $ref: 42 })).toBe(false); - expect(isCacheRef({ $ref: null })).toBe(false); + expect(isCacheRef({ kind: CACHE_REF_KIND })).toBe(false); + expect(isCacheRef({ kind: CACHE_REF_KIND, ref: "cache://k" })).toBe(false); + expect(isCacheRef({ kind: CACHE_REF_KIND, $ref: 42 })).toBe(false); + expect(isCacheRef({ kind: CACHE_REF_KIND, $ref: null })).toBe(false); + }); + + it("rejects values that lack the kind brand even when $ref is a string", () => { + expect(isCacheRef({ $ref: "cache://k" })).toBe(false); + expect(isCacheRef({ $ref: "cache://k", size: 10 })).toBe(false); + }); + + it("rejects values whose kind is not the literal brand", () => { + expect(isCacheRef({ kind: "other", $ref: "cache://k" })).toBe(false); + expect(isCacheRef({ kind: 1, $ref: "cache://k" })).toBe(false); }); it("rejects primitives and null", () => { @@ -38,16 +50,47 @@ describe("isCacheRef", () => { expect(isCacheRef(true)).toBe(false); }); - it("accepts a ref where $ref is the empty string (still string-typed)", () => { - expect(isCacheRef({ $ref: "" })).toBe(true); + it("accepts a branded ref where $ref is the empty string (still string-typed)", () => { + expect(isCacheRef(makeCacheRef({ $ref: "" }))).toBe(true); + }); + + it("does NOT confuse JSON-Schema $ref strings with cache refs", () => { + // Before the kind brand, a shape-only check would have matched any + // {$ref: string} — including JSON-Schema $refs in metadata or attacker + // payloads pointing at other-run cache keys. With the brand, a JSON-Schema + // ref no longer passes isCacheRef and the cache resolver will not walk it. + const jsonSchemaRef = { $ref: "#/$defs/Foo" }; + expect(isCacheRef(jsonSchemaRef)).toBe(false); + }); +}); + +describe("makeCacheRef", () => { + it("brands the constructed object with CACHE_REF_KIND", () => { + const ref = makeCacheRef({ $ref: "cache://x" }); + expect(ref.kind).toBe(CACHE_REF_KIND); + }); + + it("omits size and mime when not supplied", () => { + const ref = makeCacheRef({ $ref: "cache://x" }); + expect("size" in ref).toBe(false); + expect("mime" in ref).toBe(false); + }); + + it("preserves size and mime when supplied", () => { + const ref = makeCacheRef({ $ref: "cache://x", size: 99, mime: "image/png" }); + expect(ref.size).toBe(99); + expect(ref.mime).toBe("image/png"); }); - it("does not confuse JSON-Schema $ref strings with cache refs by shape", () => { - // JSON Schema $ref also uses { $ref: string }. Shape is identical at this - // layer; discrimination by call site / port-context is the contract, not - // shape inspection. This test documents the limitation. - const jsonSchemaRef = { $ref: "#/definitions/Foo" }; - expect(isCacheRef(jsonSchemaRef)).toBe(true); + it("survives JSON round-trip and still passes isCacheRef", () => { + const original = makeCacheRef({ $ref: "cache://round-trip", size: 7, mime: "text/plain" }); + const wire = JSON.stringify(original); + const received = JSON.parse(wire); + expect(isCacheRef(received)).toBe(true); + expect(received.kind).toBe(CACHE_REF_KIND); + expect(received.$ref).toBe("cache://round-trip"); + expect(received.size).toBe(7); + expect(received.mime).toBe("text/plain"); }); }); diff --git a/packages/test/src/test/task-graph/Spec2QueueRowAndRehydrate.test.ts b/packages/test/src/test/task-graph/Spec2QueueRowAndRehydrate.test.ts index f453b30ec..e8a63640b 100644 --- a/packages/test/src/test/task-graph/Spec2QueueRowAndRehydrate.test.ts +++ b/packages/test/src/test/task-graph/Spec2QueueRowAndRehydrate.test.ts @@ -4,18 +4,24 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { CacheRef, StreamEvent } from "@workglow/task-graph"; +import type { + CacheRef, + JobHandleLike, + StreamEvent, + TaskInput, + TaskOutput, +} from "@workglow/task-graph"; import { CACHE_REGISTRY, DefaultCacheRegistry, IExecuteContext, isCacheRef, + makeCacheRef, resolveJobOutput, Task, TaskOutputRepository, TaskRegistry, } from "@workglow/task-graph"; -import type { JobHandleLike, TaskInput, TaskOutput } from "@workglow/task-graph"; import { Container, ServiceRegistry, sleep } from "@workglow/util"; import { DataPortSchema } from "@workglow/util/schema"; import { beforeAll, beforeEach, describe, expect, it } from "vitest"; @@ -68,7 +74,7 @@ class StreamingMemoryRepo extends TaskOutputRepository { } const key = `inmem://${taskType}::${JSON.stringify(inputs)}`; this.streamed.set(key, merged); - return { $ref: key, size, mime: "application/octet-stream" }; + return makeCacheRef({ $ref: key, size, mime: "application/octet-stream" }); } override async getOutputByRef(ref: CacheRef): Promise { const bytes = this.streamed.get(ref.$ref); @@ -125,8 +131,37 @@ class BigBlobStreamTask extends Task, BinOut> { } } +type ArrayBufOut = { bytes: ArrayBuffer }; + +class BigArrayBufferStreamTask extends Task, ArrayBufOut> { + public static override type = "Spec2QueueRowTest_BigArrayBufferStream"; + public static override category = "Test"; + public static override cacheable = true; + + public static override outputSchema(): DataPortSchema { + return { + type: "object", + properties: { bytes: { type: "object", format: "binary", "x-stream": "binary" } }, + additionalProperties: false, + } as const satisfies DataPortSchema; + } + + async *executeStream( + _input: Record, + _ctx: IExecuteContext + ): AsyncIterable> { + for (let i = 0; i < CHUNKS; i++) { + const chunk = new Uint8Array(CHUNK).fill(i & 0xff); + yield { type: "binary-delta", port: "bytes", binaryDelta: chunk }; + if (i % 4 === 3) await sleep(0); + } + yield { type: "finish", data: {} as ArrayBufOut }; + } +} + beforeAll(() => { TaskRegistry.registerTask(BigBlobStreamTask as any); + TaskRegistry.registerTask(BigArrayBufferStreamTask as any); }); let services: ServiceRegistry; @@ -222,6 +257,33 @@ describe("Spec 2 — saved-row size & cross-process rehydration", () => { expect((resolved.bytes as Blob).size).toBe(CHUNKS * CHUNK); }); + it("rehydration below threshold inlines a Blob for format:'blob' (canonical BinaryFormat)", async () => { + const task = new BigBlobStreamTask(); + // Threshold above the full payload → rehydrate inline. + const output = await task.run( + {}, + { + registry: services, + referenceThresholdBytes: CHUNKS * CHUNK + 1, + } + ); + expect(output.bytes).toBeInstanceOf(Blob); + expect((output.bytes as Blob).size).toBe(CHUNKS * CHUNK); + }); + + it("rehydration below threshold inlines an ArrayBuffer for format:'binary' (canonical BinaryFormat)", async () => { + const task = new BigArrayBufferStreamTask(); + const output = await task.run( + {}, + { + registry: services, + referenceThresholdBytes: CHUNKS * CHUNK + 1, + } + ); + expect(output.bytes).toBeInstanceOf(ArrayBuffer); + expect((output.bytes as ArrayBuffer).byteLength).toBe(CHUNKS * CHUNK); + }); + it("dangling refs (cache cleared between save and read) resolve to undefined — best-effort contract", async () => { const task = new BigBlobStreamTask(); await task.run({}, { registry: services, referenceThresholdBytes: 0 }); diff --git a/packages/test/src/test/task-graph/StreamBinaryPump.test.ts b/packages/test/src/test/task-graph/StreamBinaryPump.test.ts index 2968139b7..fb76b3c5f 100644 --- a/packages/test/src/test/task-graph/StreamBinaryPump.test.ts +++ b/packages/test/src/test/task-graph/StreamBinaryPump.test.ts @@ -38,11 +38,11 @@ */ import type { CacheRef, ITask, StreamEvent, TaskInput, TaskOutput } from "@workglow/task-graph"; -import type { DataPortSchema } from "@workglow/util/schema"; import { Dataflow, getBinaryPortId, IExecuteContext, + makeCacheRef, StreamPump, Task, TaskGraph, @@ -51,6 +51,7 @@ import { TaskStatus, } from "@workglow/task-graph"; import { setLogger, sleep } from "@workglow/util"; +import type { DataPortSchema } from "@workglow/util/schema"; import { beforeEach, describe, expect, it } from "vitest"; import { getTestingLogger } from "../../binding/TestingLogger"; @@ -212,7 +213,7 @@ class StreamingMemoryRepo extends TaskOutputRepository { size += c.byteLength; for (const b of c) this.streamedBytes.push(b); } - return { $ref: `inmem://${taskType}::${JSON.stringify(inputs)}`, size }; + return makeCacheRef({ $ref: `inmem://${taskType}::${JSON.stringify(inputs)}`, size }); } } diff --git a/packages/test/src/test/task-graph/StreamBinaryTypes.test.ts b/packages/test/src/test/task-graph/StreamBinaryTypes.test.ts index 5a48d3e85..df2ca1fe9 100644 --- a/packages/test/src/test/task-graph/StreamBinaryTypes.test.ts +++ b/packages/test/src/test/task-graph/StreamBinaryTypes.test.ts @@ -3,17 +3,18 @@ * Copyright 2026 Steven Roussey * SPDX-License-Identifier: Apache-2.0 */ -import { describe, expect, it } from "vitest"; import type { StreamBinaryDelta, StreamEvent, StreamMode } from "@workglow/task-graph"; import { + assertBinaryFormat, + edgeNeedsAccumulation, + getBinaryPortId, + getOutputStreamMode, getPortStreamMode, getStreamingPorts, - getOutputStreamMode, - getBinaryPortId, - edgeNeedsAccumulation, materializeBinary, } from "@workglow/task-graph"; import type { DataPortSchema } from "@workglow/util/schema"; +import { describe, expect, it } from "vitest"; const binarySchema = { type: "object", @@ -32,6 +33,30 @@ const mixedSchema = { additionalProperties: false, } as const satisfies DataPortSchema; +const unannotatedBinarySchema = { + type: "object", + properties: { + bytes: { type: "object", "x-stream": "binary" }, + }, + additionalProperties: false, +} as const satisfies DataPortSchema; + +const typoFormatSchema = { + type: "object", + properties: { + bytes: { type: "object", format: "Blob", "x-stream": "binary" }, + }, + additionalProperties: false, +} as const satisfies DataPortSchema; + +const unknownFormatSchema = { + type: "object", + properties: { + bytes: { type: "object", format: "wat", "x-stream": "binary" }, + }, + additionalProperties: false, +} as const satisfies DataPortSchema; + describe("StreamBinaryDelta type", () => { it("is assignable to StreamEvent and carries a Uint8Array delta", () => { const evt: StreamEvent = { @@ -95,6 +120,30 @@ describe("binary-aware port helpers", () => { }); }); +describe("assertBinaryFormat", () => { + it("returns 'blob' when format is 'blob'", () => { + expect(assertBinaryFormat(binarySchema, "bytes")).toBe("blob"); + }); + + it("returns 'binary' when format is 'binary'", () => { + expect(assertBinaryFormat(mixedSchema, "bytes")).toBe("binary"); + }); + + it("returns 'blob' for undefined / absent format (canonical default)", () => { + expect(assertBinaryFormat(unannotatedBinarySchema, "bytes")).toBe("blob"); + }); + + it("throws on a casing typo such as 'Blob'", () => { + expect(() => assertBinaryFormat(typoFormatSchema, "bytes")).toThrow( + /Allowed: "blob" \| "binary"/ + ); + }); + + it("throws on an unknown format value", () => { + expect(() => assertBinaryFormat(unknownFormatSchema, "bytes")).toThrow(/wat/); + }); +}); + describe("materializeBinary", () => { const chunks = [new Uint8Array([1, 2]), new Uint8Array([3, 4, 5])]; @@ -111,16 +160,8 @@ describe("materializeBinary", () => { expect(Array.from(new Uint8Array(buf))).toEqual([1, 2, 3, 4, 5]); }); - it("defaults to Blob when format is undefined", () => { - expect(materializeBinary(chunks, undefined)).toBeInstanceOf(Blob); - }); - it("handles an empty chunk list", () => { expect(materializeBinary([], "binary")).toBeInstanceOf(ArrayBuffer); expect((materializeBinary([], "binary") as ArrayBuffer).byteLength).toBe(0); }); - - it("treats an unknown format as binary (ArrayBuffer)", () => { - expect(materializeBinary(chunks, "wat")).toBeInstanceOf(ArrayBuffer); - }); }); diff --git a/packages/test/src/test/task-graph/StreamProcessorBinaryRefSink.test.ts b/packages/test/src/test/task-graph/StreamProcessorBinaryRefSink.test.ts index f14b9d42d..d69d005bb 100644 --- a/packages/test/src/test/task-graph/StreamProcessorBinaryRefSink.test.ts +++ b/packages/test/src/test/task-graph/StreamProcessorBinaryRefSink.test.ts @@ -5,7 +5,13 @@ */ import type { BinaryRefSink, CacheRef, StreamEvent } from "@workglow/task-graph"; -import { IExecuteContext, isCacheRef, Task, TaskRegistry } from "@workglow/task-graph"; +import { + IExecuteContext, + isCacheRef, + makeCacheRef, + Task, + TaskRegistry, +} from "@workglow/task-graph"; import { sleep } from "@workglow/util"; import { DataPortSchema } from "@workglow/util/schema"; import { beforeAll, describe, expect, it } from "vitest"; @@ -103,7 +109,7 @@ function makeSink(): { rejectCollected(err); throw err; } - const ref = { $ref, size: bytes.length, mime: "application/octet-stream" }; + const ref = makeCacheRef({ $ref, size: bytes.length, mime: "application/octet-stream" }); resolveCollected({ ref, bytes }); return ref; }; diff --git a/packages/test/src/test/task-graph/StreamingBackpressure.test.ts b/packages/test/src/test/task-graph/StreamingBackpressure.test.ts index 0728927d2..bf00b44a9 100644 --- a/packages/test/src/test/task-graph/StreamingBackpressure.test.ts +++ b/packages/test/src/test/task-graph/StreamingBackpressure.test.ts @@ -20,10 +20,12 @@ * downstream tasks from starting. */ -import type { CachePolicy, StreamEvent } from "@workglow/task-graph"; +import type { BinaryRefSink, CachePolicy, CacheRef, StreamEvent } from "@workglow/task-graph"; import { Dataflow, IExecuteContext, + isCacheRef, + makeCacheRef, Task, TaskGraph, TaskGraphRunner, @@ -314,6 +316,198 @@ describe("Streaming backpressure and stress", () => { }); }); + describe("binary backpressure", () => { + // Sized so a fast producer would overwhelm a slow sink: 100 chunks of 1 MiB + // each = 100 MiB total, sink consumes one chunk every 50 ms (~5 s end-to-end). + const CHUNK = 1024 * 1024; + const CHUNKS = 100; + const HIGH_WATER = 4 * 1024 * 1024; + + type BinOut = { bytes: Blob }; + + /** + * Streams `CHUNKS` × `CHUNK` bytes; awaits the producer-side `push` + * promise so the byte-bounded backpressure check actually parks when the + * router buffer reaches the high-water mark. + */ + class FastBinaryProducer extends Task, BinOut> { + public static override type = "StreamingBackpressure_FastBinaryProducer"; + public static override cachePolicy: CachePolicy = { kind: "none" }; + public static override cacheable = true; + + public static override inputSchema(): DataPortSchema { + return { type: "object", properties: {}, additionalProperties: false } as const; + } + public static override outputSchema(): DataPortSchema { + return { + type: "object", + properties: { bytes: { type: "object", format: "blob", "x-stream": "binary" } }, + additionalProperties: false, + } as const satisfies DataPortSchema; + } + + async *executeStream( + _input: Record, + _ctx: IExecuteContext + ): AsyncIterable> { + for (let i = 0; i < CHUNKS; i++) { + const chunk = new Uint8Array(CHUNK).fill(i & 0xff); + yield { type: "binary-delta", port: "bytes", binaryDelta: chunk }; + } + yield { type: "finish", data: {} as BinOut }; + } + } + + it("keeps router buffer at-or-below the high-water mark while a slow sink drains", async () => { + const router = await import("@workglow/task-graph"); + const { BinaryStreamRouter } = router as unknown as { + BinaryStreamRouter: new ( + sink: BinaryRefSink, + highWaterMarkBytes: number + ) => { + push(chunk: Uint8Array): Promise; + end(): void; + ref(): Promise; + readonly _bufferedBytes: number; + }; + }; + + let observedPeak = 0; + const consumed: number[] = []; + const sink: BinaryRefSink = async (chunks) => { + for await (const c of chunks) { + consumed.push(c.byteLength); + // Slow consumer: 50 ms per chunk. + await new Promise((res) => setTimeout(res, 50)); + } + return makeCacheRef({ $ref: "inmem://bp", size: consumed.reduce((a, b) => a + b, 0) }); + }; + + const r = new BinaryStreamRouter(sink, HIGH_WATER); + for (let i = 0; i < CHUNKS; i++) { + const chunk = new Uint8Array(CHUNK).fill(i & 0xff); + await r.push(chunk); + observedPeak = Math.max(observedPeak, r._bufferedBytes); + } + r.end(); + const ref = await r.ref(); + + // Peak buffer stayed at-or-below the high-water mark + at most one chunk + // (the chunk that pushed us over the mark is counted before we park). + expect(observedPeak).toBeLessThanOrEqual(HIGH_WATER + CHUNK); + + // Every byte was delivered. + const totalDelivered = consumed.reduce((a, b) => a + b, 0); + expect(totalDelivered).toBe(CHUNKS * CHUNK); + expect(ref.size).toBe(CHUNKS * CHUNK); + expect(isCacheRef(ref)).toBe(true); + }, 30_000); + + it("releases a parked push() promise within 100ms when the router is ended (abort path)", async () => { + const router = await import("@workglow/task-graph"); + const { BinaryStreamRouter } = router as unknown as { + BinaryStreamRouter: new ( + sink: BinaryRefSink, + highWaterMarkBytes: number + ) => { + push(chunk: Uint8Array): Promise; + end(): void; + fail(err: Error): void; + ref(): Promise; + readonly _bufferedBytes: number; + }; + }; + + // Sink starts consuming but the gate keeps it parked so the producer + // buffer fills to the high-water mark. The orphaned-Promise bug pre-fix: + // `end()` did not release the parked producer, so the test would + // wait until the test-level timeout. + let gateRelease: (() => void) | undefined; + const gate = new Promise((res) => { + gateRelease = res; + }); + const seen: Uint8Array[] = []; + const sink: BinaryRefSink = async (chunks) => { + await gate; // park the sink so the buffer stays full + for await (const c of chunks) { + seen.push(c); + } + return makeCacheRef({ $ref: "inmem://parked", size: 0 }); + }; + + // High-water mark of 1 byte so a single non-empty chunk parks the next push. + const r = new BinaryStreamRouter(sink, 1); + // The first push parks immediately (1 byte >= 1-byte mark, no consumer). + const parked = r.push(new Uint8Array([0xff])); + + let parkedResolved = false; + parked.then(() => { + parkedResolved = true; + }); + + // Park before measuring. + await new Promise((res) => setTimeout(res, 10)); + expect(parkedResolved).toBe(false); + + const t0 = Date.now(); + r.end(); + await parked; + const elapsed = Date.now() - t0; + + expect(parkedResolved).toBe(true); + expect(elapsed).toBeLessThan(100); + + // Release the sink so the test exits cleanly. + gateRelease?.(); + await r.ref(); + void seen; + }, 10_000); + + it("runs a 100MiB stream end-to-end through StreamProcessor with byte-bounded backpressure", async () => { + const producer = new FastBinaryProducer({ id: "producer" }); + + let received = 0; + const sink: BinaryRefSink = async (chunks) => { + for await (const c of chunks) { + received += c.byteLength; + // Slow consumer — gives the producer time to outrun it without + // bound if backpressure didn't apply. The byte-bounded invariant + // itself is exercised directly in the BinaryStreamRouter unit + // test above; here we verify the full StreamProcessor path + // delivers every byte through the throttle without hangs or drops. + await new Promise((res) => setTimeout(res, 2)); + } + return makeCacheRef({ $ref: "inmem://e2e", size: received }); + }; + + const processor = (producer as any).runner.streamProcessor as { + run(input: any, ctx: any, deps: any): Promise; + }; + const abortController = new AbortController(); + const ctx = { + abortController, + shouldAccumulate: false, + telemetrySpan: undefined, + dispose: () => {}, + } as any; + const sinks: ReadonlyMap = new Map([["bytes", sink]]); + + const output = (await processor.run({}, ctx, { + registry: undefined as any, + resourceScope: undefined, + inputStreams: undefined, + onProgress: async () => {}, + own: (t: T) => t, + binaryRefSinks: sinks, + binaryHighWaterBytes: HIGH_WATER, + })) as BinOut | undefined; + + expect(output).toBeDefined(); + // Every byte arrived; backpressure didn't drop chunks. + expect(received).toBe(CHUNKS * CHUNK); + }, 30_000); + }); + describe("StreamError propagation", () => { it("fails the source task on a StreamError event and keeps downstream from completing", async () => { const graph = new TaskGraph(); diff --git a/packages/test/src/test/task-graph/TaskOutputRepositoryStream.test.ts b/packages/test/src/test/task-graph/TaskOutputRepositoryStream.test.ts index aa852b90f..5aad59f2f 100644 --- a/packages/test/src/test/task-graph/TaskOutputRepositoryStream.test.ts +++ b/packages/test/src/test/task-graph/TaskOutputRepositoryStream.test.ts @@ -3,9 +3,9 @@ * Copyright 2026 Steven Roussey * SPDX-License-Identifier: Apache-2.0 */ -import { describe, expect, it } from "vitest"; -import { RunPrivateCacheRepo, TaskOutputRepository } from "@workglow/task-graph"; import type { CacheRef, TaskInput, TaskOutput } from "@workglow/task-graph"; +import { makeCacheRef, RunPrivateCacheRepo, TaskOutputRepository } from "@workglow/task-graph"; +import { describe, expect, it } from "vitest"; class StreamingMemoryRepo extends TaskOutputRepository { public readonly streamed = new Map(); @@ -47,7 +47,7 @@ class StreamingMemoryRepo extends TaskOutputRepository { const key = taskType + JSON.stringify(inputs); this.streamed.set(key, merged); this.streamedMetadata.set(key, metadata); - return { $ref: `inmem://${key}`, size: total }; + return makeCacheRef({ $ref: `inmem://${key}`, size: total }); } override async getOutputByRef(ref: CacheRef): Promise { const key = ref.$ref.replace(/^inmem:\/\//, ""); @@ -144,12 +144,7 @@ describe("TaskOutputRepository.saveOutputStream", () => { it("saveOutputStream returns a CacheRef the same backing can resolve to bytes", async () => { const repo = new StreamingMemoryRepo({}); - const ref = await repo.saveOutputStream( - "T", - { k: 1 }, - gen(new Uint8Array([7, 8, 9])), - {} - ); + const ref = await repo.saveOutputStream("T", { k: 1 }, gen(new Uint8Array([7, 8, 9])), {}); expect(typeof ref.$ref).toBe("string"); expect(ref.size).toBe(3); const hydrated = await repo.getOutputByRef(ref); diff --git a/packages/test/src/test/task-graph/TaskRegistry.test.ts b/packages/test/src/test/task-graph/TaskRegistry.test.ts index 1a65cd88f..b0bdf7267 100644 --- a/packages/test/src/test/task-graph/TaskRegistry.test.ts +++ b/packages/test/src/test/task-graph/TaskRegistry.test.ts @@ -30,10 +30,58 @@ class TaskB extends Task { } } +// Binary-stream port format checks happen at registration time so a typo +// surfaces near the task definition rather than during a streaming run. + +class BinaryPortTypoTask extends Task { + static override readonly type = "TaskRegistryTest_BinaryFormatTypo"; + static override inputSchema(): DataPortSchema { + return { type: "object", properties: {}, additionalProperties: false } as const; + } + static override outputSchema(): DataPortSchema { + return { + type: "object", + properties: { bytes: { type: "object", format: "Blob", "x-stream": "binary" } }, + additionalProperties: false, + } as const satisfies DataPortSchema; + } +} + +class BinaryPortValidBlobTask extends Task { + static override readonly type = "TaskRegistryTest_BinaryValidBlob"; + static override inputSchema(): DataPortSchema { + return { type: "object", properties: {}, additionalProperties: false } as const; + } + static override outputSchema(): DataPortSchema { + return { + type: "object", + properties: { bytes: { type: "object", format: "blob", "x-stream": "binary" } }, + additionalProperties: false, + } as const satisfies DataPortSchema; + } +} + +class BinaryPortValidBinaryTask extends Task { + static override readonly type = "TaskRegistryTest_BinaryValidBinary"; + static override inputSchema(): DataPortSchema { + return { type: "object", properties: {}, additionalProperties: false } as const; + } + static override outputSchema(): DataPortSchema { + return { + type: "object", + properties: { bytes: { type: "object", format: "binary", "x-stream": "binary" } }, + additionalProperties: false, + } as const satisfies DataPortSchema; + } +} + describe("TaskRegistry", () => { afterEach(() => { // Clean up any registrations made during the test TaskRegistry.unregisterTask(TaskA.type); + TaskRegistry.unregisterTask(BinaryPortTypoTask.type); + TaskRegistry.unregisterTask(BinaryPortValidBlobTask.type); + TaskRegistry.unregisterTask(BinaryPortValidBinaryTask.type); }); it("registers a task constructor", () => { @@ -67,4 +115,22 @@ describe("TaskRegistry", () => { it("unregisterTask returns false when the type was not registered", () => { expect(TaskRegistry.unregisterTask("NonExistentType")).toBe(false); }); + + it("throws at registration when a binary port uses a typo format like 'Blob'", () => { + expect(() => TaskRegistry.registerTask(BinaryPortTypoTask)).toThrow( + /invalid binary stream port/ + ); + // And the task is NOT in the registry afterwards. + expect(TaskRegistry.all.get(BinaryPortTypoTask.type)).toBeUndefined(); + }); + + it("accepts a binary port with format 'blob'", () => { + expect(() => TaskRegistry.registerTask(BinaryPortValidBlobTask)).not.toThrow(); + expect(TaskRegistry.all.get(BinaryPortValidBlobTask.type)).toBe(BinaryPortValidBlobTask); + }); + + it("accepts a binary port with format 'binary'", () => { + expect(() => TaskRegistry.registerTask(BinaryPortValidBinaryTask)).not.toThrow(); + expect(TaskRegistry.all.get(BinaryPortValidBinaryTask.type)).toBe(BinaryPortValidBinaryTask); + }); }); diff --git a/packages/test/src/test/task-graph/TaskRunnerRefPath.test.ts b/packages/test/src/test/task-graph/TaskRunnerRefPath.test.ts index b644da9b2..64e214260 100644 --- a/packages/test/src/test/task-graph/TaskRunnerRefPath.test.ts +++ b/packages/test/src/test/task-graph/TaskRunnerRefPath.test.ts @@ -4,17 +4,17 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { CacheRef, StreamEvent } from "@workglow/task-graph"; +import type { CacheRef, StreamEvent, TaskInput, TaskOutput } from "@workglow/task-graph"; import { CACHE_REGISTRY, DefaultCacheRegistry, IExecuteContext, isCacheRef, + makeCacheRef, Task, TaskOutputRepository, TaskRegistry, } from "@workglow/task-graph"; -import type { TaskInput, TaskOutput } from "@workglow/task-graph"; import { Container, ServiceRegistry, sleep } from "@workglow/util"; import { DataPortSchema } from "@workglow/util/schema"; import { beforeAll, beforeEach, describe, expect, it } from "vitest"; @@ -66,7 +66,7 @@ class StreamingMemoryRepo extends TaskOutputRepository { } const key = `inmem://${taskType}::${JSON.stringify(inputs)}`; this.streamed.set(key, merged); - return { $ref: key, size, mime: "application/octet-stream" }; + return makeCacheRef({ $ref: key, size, mime: "application/octet-stream" }); } override async getOutputByRef(ref: CacheRef): Promise { const bytes = this.streamed.get(ref.$ref); diff --git a/packages/test/src/test/task-graph/resolveJobOutput.test.ts b/packages/test/src/test/task-graph/resolveJobOutput.test.ts index 0eb6d58e2..0f321a8c5 100644 --- a/packages/test/src/test/task-graph/resolveJobOutput.test.ts +++ b/packages/test/src/test/task-graph/resolveJobOutput.test.ts @@ -4,21 +4,20 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { describe, expect, it } from "vitest"; -import { resolveJobOutput } from "@workglow/task-graph"; import type { CacheRef, CacheRefResolver, JobHandleLike } from "@workglow/task-graph"; +import { makeCacheRef, resolveJobOutput } from "@workglow/task-graph"; +import { describe, expect, it, vi } from "vitest"; const handleOf = (value: T): JobHandleLike => ({ waitFor: async () => value, }); -const ref = (key: string, size = 0): CacheRef => ({ $ref: key, size }); +const ref = (key: string, size = 0): CacheRef => makeCacheRef({ $ref: key, size }); describe("resolveJobOutput", () => { it("awaits the job and hydrates a top-level ref through a function resolver", async () => { const blob = new Blob([new Uint8Array([1, 2, 3])]); - const resolver: CacheRefResolver = async (r) => - r.$ref === "cache://A" ? blob : undefined; + const resolver: CacheRefResolver = async (r) => (r.$ref === "cache://A" ? blob : undefined); const handle = handleOf({ bytes: ref("cache://A", 3) as unknown as Blob }); const out = await resolveJobOutput(handle, resolver); expect(out.bytes).toBe(blob); @@ -67,6 +66,20 @@ describe("resolveJobOutput", () => { await expect(resolveJobOutput(handle, async () => undefined)).rejects.toThrow("job failed"); }); + it("never invokes getOutputByRef for attacker-supplied {$ref} shapes without the brand", async () => { + // Cross-tenant attack vector: a task output contains a metadata field whose + // shape collides with the legacy CacheRef discriminator + // (`{$ref: "cache://OTHER_RUN/secret"}`). With the literal `kind` brand, + // the resolver MUST NOT pass the unbranded value to `getOutputByRef`, so + // bytes from another run/tenant can't be read by shape collision alone. + const getOutputByRef = vi.fn(async (_r: CacheRef) => new Blob([new Uint8Array([1, 2, 3])])); + const backing = { getOutputByRef }; + const handle = handleOf({ note: { $ref: "cache://OTHER_RUN/secret" } }); + const out = await resolveJobOutput(handle, backing); + expect(out.note).toEqual({ $ref: "cache://OTHER_RUN/secret" }); + expect(getOutputByRef).not.toHaveBeenCalled(); + }); + it("forwards ResolveOutputOptions to the underlying walker", async () => { let inFlight = 0; let observedMax = 0; diff --git a/packages/test/src/test/task-graph/resolveOutput.test.ts b/packages/test/src/test/task-graph/resolveOutput.test.ts index 1a5a4ead7..df95a5972 100644 --- a/packages/test/src/test/task-graph/resolveOutput.test.ts +++ b/packages/test/src/test/task-graph/resolveOutput.test.ts @@ -4,15 +4,12 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { describe, expect, it, vi } from "vitest"; -import { resolveOutput } from "@workglow/task-graph"; import type { CacheRef, CacheRefResolver } from "@workglow/task-graph"; +import { makeCacheRef, resolveOutput } from "@workglow/task-graph"; +import { describe, expect, it, vi } from "vitest"; -const ref = (key: string, size?: number, mime?: string): CacheRef => ({ - $ref: key, - ...(size !== undefined ? { size } : {}), - ...(mime !== undefined ? { mime } : {}), -}); +const ref = (key: string, size?: number, mime?: string): CacheRef => + makeCacheRef({ $ref: key, size, mime }); const fakeResolver = (table: Record): CacheRefResolver => @@ -27,6 +24,19 @@ describe("resolveOutput", () => { expect(resolver).not.toHaveBeenCalled(); }); + it("does not walk JSON-Schema-shaped {$ref: string} objects (no brand)", async () => { + // Brand discrimination matters here: a JSON-Schema $ref embedded in + // metadata must NOT be passed to the cache resolver, since the cache + // backing would treat the JSON-Schema pointer as a cache key. Identity is + // preserved because the tree has no branded refs to resolve. + const resolver = vi.fn(); + const input = { schema: { $ref: "#/$defs/Foo" }, name: "ok" }; + const out = await resolveOutput(input, resolver); + expect(out).toBe(input); + expect(out.schema).toBe(input.schema); + expect(resolver).not.toHaveBeenCalled(); + }); + it("resolves a top-level ref to bytes", async () => { const blob = new Blob([new Uint8Array([1, 2, 3])]); const table = { "cache://x": blob };