Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 40 additions & 2 deletions packages/protocol/src/bridge-frames.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
// ---------------------------------------------------------------------------
Expand Down Expand Up @@ -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
Expand All @@ -388,6 +423,9 @@ export const bridgeFrameSchema = z.discriminatedUnion("type", [
resultFrameSchema,
errorFrameSchema,
chunkFrameSchema,
blobStartFrameSchema,
blobChunkFrameSchema,
blobEndFrameSchema,
progressFrameSchema,
comfyEventFrameSchema,
blenderEventFrameSchema
Expand Down
6 changes: 5 additions & 1 deletion packages/protocol/src/bridge-protocol.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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", () => {
Expand Down
147 changes: 139 additions & 8 deletions packages/runtime/src/python-bridge-base.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
Expand Down Expand Up @@ -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);
Expand All @@ -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) {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
Expand Down
3 changes: 2 additions & 1 deletion packages/runtime/tests/python-bridge-base-coverage.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
4 changes: 2 additions & 2 deletions packages/runtime/tests/python-bridge-models.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
});
});

Expand Down
43 changes: 41 additions & 2 deletions packages/runtime/tests/python-websocket-bridge.test-helpers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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;
Expand Down
Loading
Loading