diff --git a/client.py b/client.py index 35efb8e..06c6950 100644 --- a/client.py +++ b/client.py @@ -103,6 +103,20 @@ CLIENT_CAPABILITIES = [ # this token; legacy clients surface NEGOTIATION_FAILED at # descriptor-load time rather than crashing at execute. Capability.MODEL_3D_BUNDLE_INLINE, + # Cross-node task handle envelope (vendor-native task / video / + # job IDs with lineage metadata for chained partner provider + # nodes — Tripo Texture/Refine/Rig/Retarget/Conversion, Kling + # VideoExtend, future partner chains). The decoder in + # ``serialization.decode_task_handle_envelope`` returns a + # ``TaskHandle`` dataclass rather than a plain ``str`` so the + # ``parent_chain`` lineage round-trips when the handle is passed + # through to a downstream provider node (re-encoded by + # ``serialization.encode_task_handle`` on the submit path). + # Servers gate handle-emitting AND handle-accepting descriptors + # on this token (whole-descriptor gating — partial-socket + # synthesis is rejected because a chain that breaks meaningfully + # if either end is missing is worse than no chain at all). + Capability.IO_TASK_HANDLE, ] # Reported in the X-RNP-Client-Version header. diff --git a/notes/run_task_handle_decoder.py b/notes/run_task_handle_decoder.py new file mode 100644 index 0000000..6569d6e --- /dev/null +++ b/notes/run_task_handle_decoder.py @@ -0,0 +1,392 @@ +"""Smoke test for the client-side cross-node task_handle decoder. + +Verifies: + +1. ``Capability.IO_TASK_HANDLE`` is intact at the vendored protocol + layer and exposed as the well-formed ``io:task_handle`` token. +2. ``"task_handle"`` is in ``HEAVY_TYPES`` and ``is_envelope`` returns + True for a well-formed task_handle envelope. +3. ``decode_task_handle_envelope`` returns a ``TaskHandle`` dataclass + with every field round-tripped verbatim (including a non-empty + ``parent_chain``). +4. ``str(handle)`` surfaces the native_id (so UI previews / log lines + stay human-meaningful) without leaking the lineage chain. +5. ``encode_task_handle(decode_task_handle_envelope(env)) == env`` — + the round-trip is lossless so a downstream provider node can re- + emit the handle unchanged. +6. The top-level ``decode_envelope`` dispatcher routes task_handle + envelopes to the new decoder (so proxy_node's generic + ``_deserialize_output`` picks it up without provider-specific + plumbing). +7. Malformed task_handle envelopes (missing fields, wrong types, + comma-union ``kind``, non-list ``parent_chain``, bad parent_chain + entries, wrong encoding) raise ``RnpProtocolError`` with + ``code=INTERNAL`` — no opaque fallthrough. +8. End-to-end against the server's ``make_task_handle_envelope`` + helper (imported from the comfy-rnp-server worktree) — the + envelope the server packs is exactly what the client decoder + unpacks, and ``encode_task_handle`` of the decoded result + re-produces the same wire shape. +9. ``CLIENT_CAPABILITIES`` advertises ``io:task_handle`` so server- + side capability-gated descriptors negotiate cleanly. + +Run with the ComfyUI venv (needs torch / comfy_api on PYTHONPATH): + + python notes/run_task_handle_decoder.py +""" +from __future__ import annotations + +import importlib.util +import os +import sys +import types + +# --------------------------------------------------------------------------- +# Locate the worktree under test. +# --------------------------------------------------------------------------- +_HERE = os.path.dirname(os.path.abspath(__file__)) +_CLIENT_DIR = os.path.dirname(_HERE) +_COMFYUI_DIR = os.path.dirname(os.path.dirname(_CLIENT_DIR)) +_SERVER_DIR = os.path.normpath( + os.path.join(_COMFYUI_DIR, "..", "comfy-rnp-server") +) + + +def _load_module(name: str, path: str) -> types.ModuleType: + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + sys.modules[name] = mod + spec.loader.exec_module(mod) + return mod + + +# Emulate the client package layout so the relative ``from . import +# client as rnp_client`` at the top of serialization.py resolves +# without pulling aiohttp / comfy_api_nodes. +pkg = types.ModuleType("comfy_remote_nodes_test") +pkg.__path__ = [_CLIENT_DIR] +sys.modules["comfy_remote_nodes_test"] = pkg + +protocol = _load_module( + "comfy_remote_nodes_test.protocol", + os.path.join(_CLIENT_DIR, "protocol.py"), +) +# serialization imports ``client`` for ``rnp_client``; the task_handle +# decoder doesn't use it, so a tiny stub suffices. +sys.modules["comfy_remote_nodes_test.client"] = types.ModuleType( + "comfy_remote_nodes_test.client", +) +serialization = _load_module( + "comfy_remote_nodes_test.serialization", + os.path.join(_CLIENT_DIR, "serialization.py"), +) + + +# --------------------------------------------------------------------------- +# 1. Capability constant intact + well-formed token. +# --------------------------------------------------------------------------- +assert protocol.Capability.IO_TASK_HANDLE == "io:task_handle", ( + f"Capability.IO_TASK_HANDLE wrong value: " + f"{protocol.Capability.IO_TASK_HANDLE!r}" +) +print("PASS: Capability.IO_TASK_HANDLE = 'io:task_handle'.") + + +# --------------------------------------------------------------------------- +# 2. HEAVY_TYPES + is_envelope. +# --------------------------------------------------------------------------- +assert "task_handle" in protocol.HEAVY_TYPES, ( + f"'task_handle' missing from HEAVY_TYPES: {sorted(protocol.HEAVY_TYPES)}" +) + +hand_envelope = { + "type": "task_handle", + "encoding": "vendor_inline", + "vendor": "tripo", + "kind": "MODEL_TASK_ID", + "native_id": "abc123-tripo-uuid", + "origin_node_id": "42", + "parent_chain": [ + {"vendor": "tripo", "kind": "MODEL_TASK_ID", + "native_id": "root-uuid", "origin_node_id": "9"}, + {"vendor": "tripo", "kind": "MODEL_TASK_ID", + "native_id": "prev-uuid", "origin_node_id": "17"}, + ], +} + +assert protocol.is_envelope(hand_envelope), ( + "is_envelope returned False for a well-formed task_handle envelope; " + "HEAVY_TYPES update somehow regressed." +) +print("PASS: 'task_handle' in HEAVY_TYPES and is_envelope recognises it.") + + +# --------------------------------------------------------------------------- +# 3. decode_task_handle_envelope returns a populated TaskHandle. +# --------------------------------------------------------------------------- +handle = serialization.decode_task_handle_envelope(hand_envelope) +assert isinstance(handle, serialization.TaskHandle), ( + f"decode_task_handle_envelope should return TaskHandle, got " + f"{type(handle).__name__}" +) +assert handle.vendor == "tripo" +assert handle.kind == "MODEL_TASK_ID" +assert handle.native_id == "abc123-tripo-uuid" +assert handle.origin_node_id == "42" +assert len(handle.parent_chain) == 2 +assert handle.parent_chain[0] == { + "vendor": "tripo", "kind": "MODEL_TASK_ID", + "native_id": "root-uuid", "origin_node_id": "9", +} +assert handle.parent_chain[1] == { + "vendor": "tripo", "kind": "MODEL_TASK_ID", + "native_id": "prev-uuid", "origin_node_id": "17", +} +print("PASS: decode_task_handle_envelope returns TaskHandle with full lineage.") + + +# --------------------------------------------------------------------------- +# 4. str(handle) surfaces the native_id. +# --------------------------------------------------------------------------- +assert str(handle) == "abc123-tripo-uuid", ( + f"str(TaskHandle) should be the native_id; got {str(handle)!r}" +) +print("PASS: str(TaskHandle) surfaces the native_id.") + + +# --------------------------------------------------------------------------- +# 5. encode(decode(env)) == env — lossless round-trip. +# --------------------------------------------------------------------------- +re_encoded = serialization.encode_task_handle(handle) +assert re_encoded == hand_envelope, ( + f"round-trip mismatch:\n original: {hand_envelope}\n " + f"re-encoded: {re_encoded}" +) +# Empty chain default is preserved. +solo = serialization.decode_task_handle_envelope({ + "type": "task_handle", + "encoding": "vendor_inline", + "vendor": "kling", + "kind": "KLING_VIDEO_ID", + "native_id": "video-xyz", + "origin_node_id": "5", + "parent_chain": [], +}) +assert solo.parent_chain == [], "empty parent_chain should round-trip empty" +assert serialization.encode_task_handle(solo)["parent_chain"] == [] +print("PASS: encode_task_handle round-trips lossless (including empty chain).") + + +# --------------------------------------------------------------------------- +# 6. decode_envelope dispatcher routes to the new decoder. +# --------------------------------------------------------------------------- +dispatched = serialization.decode_envelope(hand_envelope) +assert isinstance(dispatched, serialization.TaskHandle), ( + f"decode_envelope didn't dispatch task_handle to the right decoder; " + f"got {type(dispatched).__name__}" +) +assert dispatched.native_id == "abc123-tripo-uuid" +print("PASS: decode_envelope dispatches 'task_handle' -> decode_task_handle_envelope.") + + +# --------------------------------------------------------------------------- +# 7. Malformed envelopes raise RnpProtocolError(INTERNAL). +# --------------------------------------------------------------------------- +def _expect_protocol_error(bad_env: dict, needle: str) -> None: + try: + serialization.decode_task_handle_envelope(bad_env) + except protocol.RnpProtocolError as e: + assert e.code == protocol.ErrorCode.INTERNAL, f"wrong error code: {e.code}" + assert needle in str(e), f"error msg doesn't mention {needle!r}: {e}" + else: + raise AssertionError(f"expected RnpProtocolError for {needle!r}") + + +# Wrong encoding. +_expect_protocol_error( + {"type": "task_handle", "encoding": "replay_ticket_inline", + "vendor": "tripo", "kind": "MODEL_TASK_ID", + "native_id": "x", "origin_node_id": "1"}, + "Unsupported task_handle encoding", +) +# Missing top-level fields. +for missing in ("vendor", "kind", "native_id", "origin_node_id"): + bad = { + "type": "task_handle", "encoding": "vendor_inline", + "vendor": "tripo", "kind": "MODEL_TASK_ID", + "native_id": "x", "origin_node_id": "1", + } + del bad[missing] + _expect_protocol_error(bad, repr(missing)) +# Empty string fields. +_expect_protocol_error( + {"type": "task_handle", "encoding": "vendor_inline", + "vendor": "", "kind": "MODEL_TASK_ID", + "native_id": "x", "origin_node_id": "1"}, + "'vendor'", +) +# Comma-union kind rejected (input-acceptance string, not emit kind). +_expect_protocol_error( + {"type": "task_handle", "encoding": "vendor_inline", + "vendor": "tripo", + "kind": "MODEL_TASK_ID,RIG_TASK_ID,RETARGET_TASK_ID", + "native_id": "x", "origin_node_id": "1"}, + "singular socket kind", +) +# parent_chain not a list. +_expect_protocol_error( + {"type": "task_handle", "encoding": "vendor_inline", + "vendor": "tripo", "kind": "MODEL_TASK_ID", + "native_id": "x", "origin_node_id": "1", + "parent_chain": "not-a-list"}, + "'parent_chain'", +) +# Bad parent_chain entry — missing field. +_expect_protocol_error( + {"type": "task_handle", "encoding": "vendor_inline", + "vendor": "tripo", "kind": "MODEL_TASK_ID", + "native_id": "x", "origin_node_id": "1", + "parent_chain": [{"vendor": "tripo", "kind": "MODEL_TASK_ID", + "native_id": "y"}]}, # missing origin_node_id + "parent_chain[0]", +) +# Bad parent_chain entry — comma-union kind in lineage. +_expect_protocol_error( + {"type": "task_handle", "encoding": "vendor_inline", + "vendor": "tripo", "kind": "MODEL_TASK_ID", + "native_id": "x", "origin_node_id": "1", + "parent_chain": [{"vendor": "tripo", + "kind": "MODEL_TASK_ID,RIG_TASK_ID", + "native_id": "y", "origin_node_id": "3"}]}, + "parent_chain[0]", +) +# Non-dict parent_chain entry. +_expect_protocol_error( + {"type": "task_handle", "encoding": "vendor_inline", + "vendor": "tripo", "kind": "MODEL_TASK_ID", + "native_id": "x", "origin_node_id": "1", + "parent_chain": ["not-a-dict"]}, + "parent_chain[0]", +) +print("PASS: malformed task_handle envelopes raise RnpProtocolError(INTERNAL).") + + +# --------------------------------------------------------------------------- +# 8. End-to-end: server make_task_handle_envelope -> client decode. +# Skips when the server worktree isn't on disk. +# --------------------------------------------------------------------------- +server_envelopes_path = os.path.join( + _SERVER_DIR, "comfy_rnp_protocol", "envelopes.py", +) +if os.path.exists(server_envelopes_path): + server_pkg = types.ModuleType("comfy_rnp_protocol_e2e_handle") + server_pkg.__path__ = [os.path.join(_SERVER_DIR, "comfy_rnp_protocol")] + sys.modules["comfy_rnp_protocol_e2e_handle"] = server_pkg + _load_module( + "comfy_rnp_protocol_e2e_handle.constants", + os.path.join(_SERVER_DIR, "comfy_rnp_protocol", "constants.py"), + ) + server_envelopes = _load_module( + "comfy_rnp_protocol_e2e_handle.envelopes", + os.path.join(_SERVER_DIR, "comfy_rnp_protocol", "envelopes.py"), + ) + assert hasattr(server_envelopes, "make_task_handle_envelope"), ( + "Server envelopes module missing make_task_handle_envelope — " + "server-side PR #39 should have added it." + ) + server_env = server_envelopes.make_task_handle_envelope( + vendor="tripo", + kind="MODEL_TASK_ID", + native_id="server-built-uuid", + origin_node_id="100", + parent_chain=[ + {"vendor": "tripo", "kind": "MODEL_TASK_ID", + "native_id": "ancestor-uuid", "origin_node_id": "50"}, + ], + ) + assert server_env["type"] == "task_handle", server_env + assert server_env["encoding"] == "vendor_inline", server_env + # Client decode of the server-packed bytes. + server_decoded = serialization.decode_task_handle_envelope(server_env) + assert isinstance(server_decoded, serialization.TaskHandle) + assert server_decoded.native_id == "server-built-uuid" + assert server_decoded.origin_node_id == "100" + assert len(server_decoded.parent_chain) == 1 + assert server_decoded.parent_chain[0]["native_id"] == "ancestor-uuid" + # Round-trip back through encode_task_handle reproduces the server's wire shape. + re_encoded_server = serialization.encode_task_handle(server_decoded) + assert re_encoded_server == server_env, ( + f"client encode -> server wire mismatch:\n" + f" server: {server_env}\n client: {re_encoded_server}" + ) + # Server-side rejection paths fire as expected (defence in depth — + # the client decoder mirrors the same rejections, but a malicious / + # buggy producer can't sneak a bad envelope past the server-side + # builder either). + try: + server_envelopes.make_task_handle_envelope( + vendor="tripo", + kind="MODEL_TASK_ID,RIG_TASK_ID", # comma-union rejected + native_id="x", + origin_node_id="1", + ) + except ValueError as e: + assert "singular socket kind" in str(e), ( + f"server-side comma-union rejection wrong msg: {e!r}" + ) + else: + raise AssertionError( + "server should reject comma-union 'kind' at envelope build time" + ) + # Server enforces parent_chain depth cap (32). + try: + server_envelopes.make_task_handle_envelope( + vendor="tripo", + kind="MODEL_TASK_ID", + native_id="x", + origin_node_id="1", + parent_chain=[ + {"vendor": "tripo", "kind": "MODEL_TASK_ID", + "native_id": f"u{i}", "origin_node_id": str(i)} + for i in range(33) + ], + ) + except ValueError as e: + assert "depth" in str(e), ( + f"server-side depth-cap rejection wrong msg: {e!r}" + ) + else: + raise AssertionError( + "server should reject parent_chain over MAX_CHAIN_DEPTH" + ) + print( + "PASS: end-to-end server make_task_handle_envelope -> client " + "decode_task_handle_envelope -> encode_task_handle round-trips " + "losslessly, and server-side rejections fire." + ) +else: + print( + f"SKIP end-to-end: server worktree not found at {_SERVER_DIR!r}" + ) + + +# --------------------------------------------------------------------------- +# 9. CLIENT_CAPABILITIES advertises io:task_handle. +# client.py imports aiohttp / comfy_api_nodes — same source-level +# contract check we used for the GLB cap in PR #1. +# --------------------------------------------------------------------------- +client_src = open( + os.path.join(_CLIENT_DIR, "client.py"), "r", encoding="utf-8", +).read() +assert "CLIENT_CAPABILITIES = [" in client_src, ( + "client.py missing CLIENT_CAPABILITIES list" +) +cc_start = client_src.index("CLIENT_CAPABILITIES = [") +cc_end = client_src.index("]", cc_start) +cc_block = client_src[cc_start:cc_end] +assert "Capability.IO_TASK_HANDLE" in cc_block, ( + "Capability.IO_TASK_HANDLE not advertised in CLIENT_CAPABILITIES" +) +print("PASS: CLIENT_CAPABILITIES advertises Capability.IO_TASK_HANDLE.") + +print("ALL CHECKS PASSED.") diff --git a/protocol.py b/protocol.py index 086ebe7..616c02f 100644 --- a/protocol.py +++ b/protocol.py @@ -225,9 +225,61 @@ class Capability: # ``uri`` per file for URL-fetch-out-of-band semantics; not in # this PR. MODEL_3D_BUNDLE_INLINE = "model_3d:bundle_inline" + # Cross-node task handle envelope: wraps a vendor-native task / + # video / job ID (Tripo MODEL_TASK_ID, Kling video_id, future + # partner handles) with enough lineage metadata to chain provider + # nodes (Tripo Texture/Refine/Rig/Retarget/Conversion, Kling + # VideoExtend, ...) across a single workflow graph. Wire shape: + # + # { + # "type": "task_handle", + # "encoding": "vendor_inline", + # "vendor": "tripo", + # "kind": "MODEL_TASK_ID", + # "native_id": "abc123-tripo-uuid", + # "origin_node_id": "42", + # "parent_chain": [ + # {"vendor": "tripo", "kind": "MODEL_TASK_ID", + # "native_id": "prev-uuid", "origin_node_id": "17"} + # ] + # } + # + # ``vendor`` is a lowercase provider slug. ``kind`` is the + # singular emitted socket kind string (``MODEL_TASK_ID`` / + # ``RIG_TASK_ID`` / ``RETARGET_TASK_ID`` / ``KLING_VIDEO_ID``) — + # NEVER the comma-union string that provider nodes use for input + # type-acceptance (e.g. ``TripoConversionNode`` accepts + # ``"MODEL_TASK_ID,RIG_TASK_ID,RETARGET_TASK_ID"``; the wire kind + # is whatever the upstream actually emitted). ``native_id`` is + # the vendor-native opaque string. ``origin_node_id`` is the + # upstream workflow node's unique id (log correlation + future + # replay). + # + # ``parent_chain`` is a FLAT list (not recursive nested + # envelopes) of ``{vendor, kind, native_id, origin_node_id}`` + # refs in root-to-immediate-parent order. Each downstream node + # appends its parent's ref before the parent's chain is copied + # through; the chain captures the full lineage. Hard-capped at + # depth 32 server-side to prevent runaway growth. + # + # Negotiation: server emits ``task_handle`` envelopes only when + # the inbound request advertises this capability; legacy clients + # surface NEGOTIATION_FAILED at descriptor-load time on any + # provider that publishes a handle-typed output OR accepts one + # as input. Whole-descriptor gating — partial-socket synthesis + # is rejected because the chain breaks meaningfully if either + # end is missing. + # + # The client decoder returns a ``TaskHandle`` dataclass (NOT a + # plain string) so lineage metadata round-trips when the value + # is passed through to the next provider node — re-encoding the + # same dataclass on a downstream submit preserves ``parent_chain`` + # without the proxy_node needing to know anything provider- + # specific. + IO_TASK_HANDLE = "io:task_handle" -HEAVY_TYPES = frozenset({"image", "video", "audio", "mask", "model_3d"}) +HEAVY_TYPES = frozenset({"image", "video", "audio", "mask", "model_3d", "task_handle"}) def is_envelope(value: Any) -> bool: diff --git a/serialization.py b/serialization.py index c1420cb..059933f 100644 --- a/serialization.py +++ b/serialization.py @@ -795,16 +795,165 @@ def decode_model3d_envelope(envelope: dict[str, Any]) -> Any: ) +# --------------------------------------------------------------------------- +# Cross-node task handle +# +# A ``task_handle`` envelope wraps a vendor-native task / video / job +# ID (Tripo MODEL_TASK_ID, Kling video_id, ...) with lineage metadata +# so chained provider nodes can pass the handle through the workflow +# without losing the upstream chain. The decoder returns a +# :class:`TaskHandle` dataclass — re-encoding the same dataclass on a +# downstream submit (via :func:`encode_task_handle`) preserves the +# full ``parent_chain`` without the proxy_node having to know +# anything provider-specific. Plain ``str`` round-trip would silently +# drop lineage, so the decoder deliberately surfaces the rich type +# even at the cost of slightly more work on the consumer side. +# --------------------------------------------------------------------------- + +from dataclasses import dataclass, field + + +@dataclass +class TaskHandle: + """Cross-node handle to a vendor-native async task / video / job. + + See :data:`comfy_remote_nodes.protocol.Capability.IO_TASK_HANDLE` + for the wire-shape contract. The dataclass mirrors the envelope + fields verbatim so :func:`encode_task_handle` is a trivial + round-trip. ``parent_chain`` is a flat list of + ``{vendor, kind, native_id, origin_node_id}`` dicts in + root-to-immediate-parent order; treat it as opaque on the client + — the server is the only component that ever consumes it. + """ + + vendor: str + kind: str + native_id: str + origin_node_id: str + parent_chain: list[dict[str, str]] = field(default_factory=list) + + def __str__(self) -> str: + # Workflow runtime sometimes coerces socket values to ``str`` + # (UI previews, log lines). Surfacing the native id keeps that + # output human-meaningful without leaking the lineage chain + # into UI strings. + return self.native_id + + +def _validate_task_handle_ref(ref: Any, *, context: str) -> dict[str, str]: + """Validate one ``{vendor, kind, native_id, origin_node_id}`` ref. + + Mirror of the server-side ``_validate_task_handle_ref`` — + rejects malformed parent_chain entries with the same field-by- + field error messages so a bad envelope surfaces the same way on + both sides. + """ + if not isinstance(ref, dict): + raise RnpProtocolError( + f"{context}: expected dict, got {type(ref).__name__}", + code=ErrorCode.INTERNAL, + ) + out: dict[str, str] = {} + for fld in ("vendor", "kind", "native_id", "origin_node_id"): + val = ref.get(fld) + if not isinstance(val, str) or not val: + raise RnpProtocolError( + f"{context}: field {fld!r} must be a non-empty string " + f"(got {val!r})", + code=ErrorCode.INTERNAL, + ) + out[fld] = val + if "," in out["kind"]: + raise RnpProtocolError( + f"{context}: 'kind' must be a singular socket kind, not a " + f"comma-union string (got {out['kind']!r})", + code=ErrorCode.INTERNAL, + ) + return out + + +def decode_task_handle_envelope(envelope: dict[str, Any]) -> TaskHandle: + """Decode a ``task_handle`` envelope into a :class:`TaskHandle`. + + Only the ``vendor_inline`` encoding is recognised today (a future + sibling ``replay_ticket_inline`` would carry a signed server- + local replay token for cross-server portability — not in this + PR). All four top-level fields and every ``parent_chain`` entry + are validated; rejections raise ``RnpProtocolError(INTERNAL)`` + with the same field-by-field messages the server emits at + envelope-build time. + """ + encoding = envelope.get("encoding") + if encoding != "vendor_inline": + raise RnpProtocolError( + f"Unsupported task_handle encoding: {encoding!r}", + code=ErrorCode.INTERNAL, + ) + top = _validate_task_handle_ref( + { + "vendor": envelope.get("vendor"), + "kind": envelope.get("kind"), + "native_id": envelope.get("native_id"), + "origin_node_id": envelope.get("origin_node_id"), + }, + context="task_handle", + ) + chain_in = envelope.get("parent_chain", []) + if not isinstance(chain_in, list): + raise RnpProtocolError( + f"task_handle 'parent_chain' must be a list (got " + f"{type(chain_in).__name__})", + code=ErrorCode.INTERNAL, + ) + chain_out = [ + _validate_task_handle_ref(entry, context=f"parent_chain[{i}]") + for i, entry in enumerate(chain_in) + ] + return TaskHandle( + vendor=top["vendor"], + kind=top["kind"], + native_id=top["native_id"], + origin_node_id=top["origin_node_id"], + parent_chain=chain_out, + ) + + +def encode_task_handle(handle: TaskHandle) -> dict[str, Any]: + """Encode a :class:`TaskHandle` back into a wire-format envelope. + + Symmetric counterpart of :func:`decode_task_handle_envelope` so a + handle decoded from one provider's output and fed into another + provider's input round-trips losslessly — the proxy_node calls + this whenever it sees a ``TaskHandle`` in the input values for a + descriptor that accepts the ``task_handle`` envelope type. + + The returned dict is the same shape the server's + ``make_task_handle_envelope`` produces, so the server-side + decoder accepts it without any branching on "is this a fresh + handle or a re-encoded one" — the wire is the source of truth. + """ + return { + "type": "task_handle", + "encoding": "vendor_inline", + "vendor": handle.vendor, + "kind": handle.kind, + "native_id": handle.native_id, + "origin_node_id": handle.origin_node_id, + "parent_chain": [dict(entry) for entry in handle.parent_chain], + } + + # --------------------------------------------------------------------------- # Dispatch by envelope type # --------------------------------------------------------------------------- _DECODERS = { - "image": decode_image_envelope, - "mask": decode_mask_envelope, - "audio": decode_audio_envelope, - "video": decode_video_envelope, - "model_3d": decode_model3d_envelope, + "image": decode_image_envelope, + "mask": decode_mask_envelope, + "audio": decode_audio_envelope, + "video": decode_video_envelope, + "model_3d": decode_model3d_envelope, + "task_handle": decode_task_handle_envelope, } @@ -856,6 +1005,9 @@ __all__ = [ "decode_video_envelope", "decode_model3d_envelope", "_bundled_file3d_class", + "TaskHandle", + "decode_task_handle_envelope", + "encode_task_handle", "decode_envelope", "is_envelope", "is_image_tensor",