feat(runtime): receive chunked worker blobs (#5688)

* feat(runtime): receive chunked worker blobs

* fix(runtime): build blob transfer payload explicitly

* test(runtime): update protocol v5 expectations

Author: georgi

packages/protocol/src/bridge-frames.ts

@@ -21,8 +21,9 @@
  *    Python worker repo's test suite validates against, so the two sides of
  *    the bridge can never silently drift apart.
  *
- * The dispatcher only ever switches on seven frame `type`s — `discover`,
- * `result`, `error`, `chunk`, `progress`, `comfy.event`, and `blender.event`
+ * The dispatcher switches on the response frame `type`s modeled below,
+ * including execution results, streamed output, blob transfer, progress, and
+ * integration events.
  * — everything else is either a request the JS side sends (`execute`,
  * `worker.status`, `provider.*`, `models.*`, `comfy.execute`,
  * `blender.execute`, …) or silently ignored. Only the seven response types
@@ -229,6 +230,37 @@ export const chunkFrameSchema = z.object({
   data: resultOrChunkDataSchema
 });
 
+const blobNameSchema = z.string().min(1);
+
+export const blobStartFrameSchema = z.object({
+  type: z.literal("blob.start"),
+  request_id: requestIdSchema,
+  data: z.object({
+    name: blobNameSchema,
+    size: z.number().int().nonnegative()
+  })
+});
+
+export const blobChunkFrameSchema = z.object({
+  type: z.literal("blob.chunk"),
+  request_id: requestIdSchema,
+  data: z.object({
+    name: blobNameSchema,
+    offset: z.number().int().nonnegative(),
+    bytes: z.instanceof(Uint8Array)
+  })
+});
+
+export const blobEndFrameSchema = z.object({
+  type: z.literal("blob.end"),
+  request_id: requestIdSchema,
+  data: z.object({
+    name: blobNameSchema,
+    size: z.number().int().nonnegative(),
+    sha256: z.string().regex(/^[a-f0-9]{64}$/)
+  })
+});
+
 // ---------------------------------------------------------------------------
 // progress
 // ---------------------------------------------------------------------------
@@ -375,6 +407,9 @@ export const bridgeFrameSchemas = {
   result: resultFrameSchema,
   error: errorFrameSchema,
   chunk: chunkFrameSchema,
+  "blob.start": blobStartFrameSchema,
+  "blob.chunk": blobChunkFrameSchema,
+  "blob.end": blobEndFrameSchema,
   progress: progressFrameSchema,
   "comfy.event": comfyEventFrameSchema,
   "blender.event": blenderEventFrameSchema
@@ -388,6 +423,9 @@ export const bridgeFrameSchema = z.discriminatedUnion("type", [
   resultFrameSchema,
   errorFrameSchema,
   chunkFrameSchema,
+  blobStartFrameSchema,
+  blobChunkFrameSchema,
+  blobEndFrameSchema,
   progressFrameSchema,
   comfyEventFrameSchema,
   blenderEventFrameSchema

packages/protocol/src/bridge-protocol.ts

@@ -46,7 +46,7 @@
  *   3. Update `MIN_NODETOOL_CORE_VERSION` to that new release.
  */
 
-export const BRIDGE_PROTOCOL_VERSION = 4;
+export const BRIDGE_PROTOCOL_VERSION = 5;
 
 /**
  * Hard floor: the JS runtime rejects (at `discover`) any worker reporting a
@@ -63,6 +63,10 @@ export const BRIDGE_PROTOCOL_VERSION = 4;
  * JS side sends them unconditionally; only the new `job.start` / `job.end` /
  * `models.evict` message types are gated, because a pre-v4 worker answers
  * those with `Unknown message type`.
+ *
+ * v5 adds an optional `chunked-v1` execute-result blob transfer. The JS side
+ * requests it only from workers reporting v5 or newer; older workers keep
+ * returning the legacy inline `blobs` map.
  */
 export const MIN_BRIDGE_PROTOCOL_VERSION = 1;
 

packages/protocol/tests/error-and-protocol-coverage.test.ts

@@ -79,7 +79,7 @@ describe("error-helpers.isTRPCErrorWithCode", () => {
 
 describe("bridge-protocol constants", () => {
   it("BRIDGE_PROTOCOL_VERSION is the current speaking version", () => {
-    expect(BRIDGE_PROTOCOL_VERSION).toBe(4);
+    expect(BRIDGE_PROTOCOL_VERSION).toBe(5);
   });
 
   it("MIN_BRIDGE_PROTOCOL_VERSION is the hard floor at 1", () => {

packages/runtime/src/python-bridge-base.ts

@@ -121,6 +121,14 @@ interface PendingRequest {
   resolve: (value: ExecuteResult) => void;
   reject: (error: Error) => void;
   onProgress?: (event: ProgressEvent) => void;
+  blobTransfers?: Map<string, BlobTransfer>;
+  completedBlobs?: Record<string, Uint8Array>;
+}
+
+interface BlobTransfer {
+  size: number;
+  received: number;
+  chunks: Uint8Array[];
 }
 
 interface PendingStreamRequest {
@@ -138,6 +146,9 @@ const DEFAULT_STATUS_TIMEOUT_MS = Number(
 const DEFAULT_DOWNLOAD_IDLE_TIMEOUT_MS = Number(
   safeProcessEnv()["NODETOOL_PYTHON_DOWNLOAD_IDLE_TIMEOUT_MS"] ?? 5 * 60 * 1000
 );
+const MAX_RESULT_BLOB_BYTES = Number(
+  safeProcessEnv()["NODETOOL_PYTHON_MAX_RESULT_BLOB_BYTES"] ?? 2 * 1024 * 1024 * 1024
+);
 
 /**
  * Transport-agnostic Python bridge. Subclasses provide the transport via
@@ -312,11 +323,20 @@ export abstract class PythonBridgeBase
       const pending = this._pending.get(requestId);
       if (pending) {
         this._pending.delete(requestId);
+        if (pending.blobTransfers?.size) {
+          pending.reject(
+            new Error("Python worker ended the result before all blob transfers completed")
+          );
+          return;
+        }
         const data = msg.data as {
           outputs: Record<string, unknown>;
           blobs: Record<string, Uint8Array>;
         };
-        pending.resolve({ outputs: data.outputs, blobs: data.blobs ?? {} });
+        pending.resolve({
+          outputs: data.outputs,
+          blobs: { ...(data.blobs ?? {}), ...(pending.completedBlobs ?? {}) }
+        });
       }
     } else if (type === "error" && requestId) {
       const streamReq = this._pendingStream.get(requestId);
@@ -341,6 +361,12 @@ export abstract class PythonBridgeBase
       if (streamReq) {
         streamReq.onChunk(msg.data as Record<string, unknown>);
       }
+    } else if (type === "blob.start" && requestId) {
+      this._startBlobTransfer(requestId, msg.data as Record<string, unknown>);
+    } else if (type === "blob.chunk" && requestId) {
+      this._appendBlobChunk(requestId, msg.data as Record<string, unknown>);
+    } else if (type === "blob.end" && requestId) {
+      this._finishBlobTransfer(requestId, msg.data as Record<string, unknown>);
     } else if (type === "progress" && requestId) {
       const pending = this._pending.get(requestId);
       if (pending?.onProgress) {
@@ -368,6 +394,103 @@ export abstract class PythonBridgeBase
     }
   }
 
+  private _rejectBlobTransfer(requestId: string, message: string): void {
+    const pending = this._pending.get(requestId);
+    if (!pending) return;
+    this._pending.delete(requestId);
+    pending.reject(new Error(message));
+    try {
+      this.cancel(requestId);
+    } catch {
+      // The worker may already have completed; cancellation is best-effort.
+    }
+  }
+
+  private _startBlobTransfer(
+    requestId: string,
+    data: Record<string, unknown>
+  ): void {
+    const pending = this._pending.get(requestId);
+    const name = data["name"];
+    const size = data["size"];
+    if (!pending || typeof name !== "string" || typeof size !== "number") return;
+    if (!Number.isSafeInteger(size) || size < 0 || size > MAX_RESULT_BLOB_BYTES) {
+      this._rejectBlobTransfer(
+        requestId,
+        `Python worker declared invalid blob size ${String(size)} for "${name}"`
+      );
+      return;
+    }
+    pending.blobTransfers ??= new Map();
+    pending.completedBlobs ??= Object.create(null) as Record<string, Uint8Array>;
+    if (pending.blobTransfers.has(name) || Object.hasOwn(pending.completedBlobs, name)) {
+      this._rejectBlobTransfer(requestId, `Python worker started duplicate blob "${name}"`);
+      return;
+    }
+    pending.blobTransfers.set(name, { size, received: 0, chunks: [] });
+  }
+
+  private _appendBlobChunk(
+    requestId: string,
+    data: Record<string, unknown>
+  ): void {
+    const pending = this._pending.get(requestId);
+    const name = data["name"];
+    const offset = data["offset"];
+    const bytes = data["bytes"];
+    if (
+      !pending ||
+      typeof name !== "string" ||
+      typeof offset !== "number" ||
+      !(bytes instanceof Uint8Array)
+    ) return;
+    const transfer = pending.blobTransfers?.get(name);
+    if (!transfer || offset !== transfer.received || offset + bytes.length > transfer.size) {
+      this._rejectBlobTransfer(
+        requestId,
+        `Python worker sent an out-of-order or oversized chunk for blob "${name}"`
+      );
+      return;
+    }
+    transfer.chunks.push(bytes);
+    transfer.received += bytes.length;
+  }
+
+  private _finishBlobTransfer(
+    requestId: string,
+    data: Record<string, unknown>
+  ): void {
+    const pending = this._pending.get(requestId);
+    const name = data["name"];
+    const size = data["size"];
+    const expectedDigest = data["sha256"];
+    if (
+      !pending ||
+      typeof name !== "string" ||
+      typeof size !== "number" ||
+      typeof expectedDigest !== "string"
+    ) return;
+    const transfer = pending.blobTransfers?.get(name);
+    if (!transfer || size !== transfer.size || transfer.received !== transfer.size) {
+      this._rejectBlobTransfer(requestId, `Python worker truncated blob "${name}"`);
+      return;
+    }
+    const blob = new Uint8Array(transfer.size);
+    let offset = 0;
+    for (const chunk of transfer.chunks) {
+      blob.set(chunk, offset);
+      offset += chunk.length;
+    }
+    const digest = nodeCrypto?.createHash("sha256").update(blob).digest("hex");
+    if (!digest || digest !== expectedDigest) {
+      this._rejectBlobTransfer(requestId, `Python worker blob "${name}" failed SHA-256 verification`);
+      return;
+    }
+    pending.blobTransfers?.delete(name);
+    pending.completedBlobs ??= Object.create(null) as Record<string, Uint8Array>;
+    pending.completedBlobs[name] = blob;
+  }
+
   /**
    * A frame decoded fine off the wire (valid msgpack) but failed its
    * `@nodetool-ai/protocol` schema — a worker-side protocol bug, not a
@@ -514,19 +637,27 @@ export abstract class PythonBridgeBase
 
     log.debug("Python bridge execute dispatched", { nodeType, requestId });
 
+    const executeData: Record<string, unknown> = {
+      node_type: nodeType,
+      fields,
+      secrets,
+      blobs,
+      ...this._identityPayload(identity)
+    };
+    if (
+      this._workerStatus?.protocol_version != null &&
+      this._workerStatus.protocol_version >= 5
+    ) {
+      executeData["blob_transfer"] = "chunked-v1";
+    }
+
     const executePromise = new Promise<ExecuteResult>((resolve, reject) => {
       this._pending.set(requestId, { resolve, reject, onProgress });
       try {
         this._send({
           type: "execute",
           request_id: requestId,
-          data: {
-            node_type: nodeType,
-            fields,
-            secrets,
-            blobs,
-            ...this._identityPayload(identity)
-          }
+          data: executeData
         });
       } catch (err) {
         this._pending.delete(requestId);

packages/runtime/tests/python-bridge-base-coverage.test.ts

@@ -330,7 +330,8 @@ describe("PythonBridgeBase — execute", () => {
       node_type: "n.T",
       fields: { a: 1 },
       secrets: { KEY: "v" },
-      blobs: {}
+      blobs: {},
+      blob_transfer: "chunked-v1"
     });
     bridge.handle({
       type: "result",

packages/runtime/tests/python-bridge-models.test.ts

@@ -19,8 +19,8 @@ import {
 import type { ModelDownloadUpdate } from "../src/python-bridge-types.js";
 
 describe("bridge protocol version", () => {
-  it("is 4 (models.* + comfy.* + run identity/job.* support)", () => {
-    expect(BRIDGE_PROTOCOL_VERSION).toBe(4);
+  it("is 5 (chunked result blobs)", () => {
+    expect(BRIDGE_PROTOCOL_VERSION).toBe(5);
   });
 });
 

packages/runtime/tests/python-websocket-bridge.test-helpers.ts

@@ -8,6 +8,7 @@
 
 import { WebSocketServer, type WebSocket as WsServerSocket } from "ws";
 import { AddressInfo } from "node:net";
+import { createHash } from "node:crypto";
 import { pack, unpack } from "msgpackr";
 
 import { BRIDGE_PROTOCOL_VERSION } from "@nodetool-ai/protocol/bridge-protocol";
@@ -35,6 +36,10 @@ interface FakeWorkerOptions {
   answerStatus?: boolean;
   /** When false, execute requests are silently ignored (no reply). */
   answerExecute?: boolean;
+  /** Return this execute output through protocol-v5 chunked blob frames. */
+  executeBlob?: Uint8Array;
+  /** Override the chunked blob digest to exercise integrity failures. */
+  executeBlobSha256?: string;
   /**
    * execute.stream behavior:
    *  - "chunks": emit two chunks then the empty result terminator (default)
@@ -288,6 +293,8 @@ export function startFakeWorker(
     answerDiscover: initialOptions.answerDiscover ?? true,
     answerStatus: initialOptions.answerStatus ?? true,
     answerExecute: initialOptions.answerExecute ?? true,
+    executeBlob: initialOptions.executeBlob ?? new Uint8Array(),
+    executeBlobSha256: initialOptions.executeBlobSha256 ?? "",
     streamMode: initialOptions.streamMode ?? "chunks",
     protocolVersion: initialOptions.protocolVersion ?? BRIDGE_PROTOCOL_VERSION,
     downloadMode: initialOptions.downloadMode ?? "progress",
@@ -372,13 +379,45 @@ export function startFakeWorker(
           break;
         case "execute": {
           if (!opts.answerExecute) break;
-          const data = msg.data as { fields?: Record<string, unknown> };
+          const data = msg.data as {
+            fields?: Record<string, unknown>;
+            blob_transfer?: string;
+          };
+          if (opts.executeBlob.length > 0 && data.blob_transfer === "chunked-v1") {
+            const blob = opts.executeBlob;
+            send({
+              type: "blob.start",
+              request_id: requestId,
+              data: { name: "out", size: blob.length }
+            });
+            for (let offset = 0; offset < blob.length; offset += 3) {
+              send({
+                type: "blob.chunk",
+                request_id: requestId,
+                data: { name: "out", offset, bytes: blob.subarray(offset, offset + 3) }
+              });
+            }
+            send({
+              type: "blob.end",
+              request_id: requestId,
+              data: {
+                name: "out",
+                size: blob.length,
+                sha256:
+                  opts.executeBlobSha256 ||
+                  createHash("sha256").update(blob).digest("hex")
+              }
+            });
+          }
           send({
             type: "result",
             request_id: requestId,
             data: {
               outputs: { out: (data.fields?.value as string) ?? "executed" },
-              blobs: {}
+              blobs:
+                opts.executeBlob.length > 0 && data.blob_transfer !== "chunked-v1"
+                  ? { out: opts.executeBlob }
+                  : {}
             }
           });
           break;