flo-python 0.1.0.dev4__py3-none-any.whl → 0.1.2__py3-none-any.whl

flo/wire.py

@@ -4,33 +4,39 @@ Binary serialization/deserialization for the Flo protocol.
 Header: 32 bytes, little-endian, CRC32 validated.
 """
 
+import json
 import struct
 from binascii import crc32
 from dataclasses import dataclass
+from typing import Any
 
 from . import types as types
 from .exceptions import (
+    BlockTooLongError,
     IncompleteResponseError,
     InvalidChecksumError,
     InvalidMagicError,
     KeyTooLargeError,
     NamespaceTooLargeError,
-    UnsupportedVersionError,
+    TableMismatchError,
+    ValidationError,
     ValueTooLargeError,
 )
+from .kv_body import is_table_op
 from .types import (
     HEADER_SIZE,
     MAGIC,
+    MAX_BLOCK_MS,
     MAX_KEY_SIZE,
     MAX_NAMESPACE_SIZE,
     MAX_VALUE_SIZE,
+    TABLE_HASH,
     VERSION,
     DequeueResult,
-    KVEntry,
     Message,
     OpCode,
     OptionTag,
-    ScanResult,
+    PendingEntry,
     StatusCode,
     StorageTier,
     StreamAppendResult,
@@ -38,18 +44,17 @@ from .types import (
     StreamInfo,
     StreamReadResult,
     StreamRecord,
-    VersionEntry,
 )
 
 # Request header (32 bytes):
 #   magic(u32) + payload_len(u32) + request_id(u64)
-#   + crc32(u32) + op_code(u16) + version(u8) + flags(u8) + reserved(8s)
-REQUEST_HEADER_FORMAT = "<IIQIHBB8s"
+#   + crc32(u32) + op_code(u16) + version(u8) + flags(u8) + table_hash(u64)
+REQUEST_HEADER_FORMAT = "<IIQIHBBQ"
 
 # Response header (32 bytes):
 #   magic(u32) + data_len(u32) + request_id(u64)
-#   + crc32(u32) + version(u8) + status(u8) + flags(u8) + _pad(u8) + reserved(8s)
-RESPONSE_HEADER_FORMAT = "<IIQIBBBB8s"
+#   + crc32(u32) + version(u8) + status(u8) + flags(u8) + _pad(u8) + table_hash(u64)
+RESPONSE_HEADER_FORMAT = "<IIQIBBBBQ"
 
 
 # =============================================================================
@@ -57,6 +62,18 @@ RESPONSE_HEADER_FORMAT = "<IIQIBBBB8s"
 # =============================================================================
 
 
+def _validate_block_ms(block_ms: int) -> None:
+    """Refuse a block_ms the server would refuse or u32 cannot encode, before the round trip."""
+    if not isinstance(block_ms, int) or isinstance(block_ms, bool):
+        raise ValidationError(f"block_ms must be an int, got {block_ms!r}")
+    if block_ms < 0:
+        raise ValidationError(f"block_ms must not be negative, got {block_ms}")
+    if block_ms > MAX_BLOCK_MS:
+        raise BlockTooLongError(
+            f"a blocking wait (block_ms) is at most {MAX_BLOCK_MS} ms (5 minutes), got {block_ms}"
+        )
+
+
 class OptionsBuilder:
     """Helper for building TLV-encoded options."""
 
@@ -70,6 +87,8 @@ class OptionsBuilder:
 
     def add_u32(self, tag: OptionTag, value: int) -> "OptionsBuilder":
         """Add a u32 option."""
+        if tag == OptionTag.BLOCK_MS:
+            _validate_block_ms(value)
         self._buffer.extend(bytes([tag, 4]))
         self._buffer.extend(struct.pack("<I", value))
         return self
@@ -178,7 +197,7 @@ def compute_crc32(header_bytes: bytes, payload: bytes) -> int:
 
     The CRC32 is computed over:
     - Header bytes 0-15 (magic, payload_length, request_id)
-    - Header bytes 20-31 (op_code/version/flags/reserved or version/status/flags/pad/reserved)
+    - Header bytes 20-31 (op_code/version/flags/table_hash or version/status/flags/pad/table_hash)
     - Payload bytes
     """
     # Hash bytes 0-15 and 20-31 of header (skip crc32 field at 16-19)
@@ -233,13 +252,19 @@ def serialize_request(
     payload.extend(struct.pack("<H", len(key)))
     payload.extend(key)
 
-    # Value: [len:u32][data]
-    payload.extend(struct.pack("<I", len(value)))
-    payload.extend(value)
+    if is_table_op(op_code):
+        # A table op's body, coded by its schema, is `value`; it takes no options.
+        if options:
+            raise ValidationError(f"op {op_code} takes a body, not options")
+        payload.extend(value)
+    else:
+        # Value: [len:u32][data]
+        payload.extend(struct.pack("<I", len(value)))
+        payload.extend(value)
 
-    # Options: [len:u16][data]
-    payload.extend(struct.pack("<H", len(options)))
-    payload.extend(options)
+        # Options: [len:u16][data]
+        payload.extend(struct.pack("<H", len(options)))
+        payload.extend(options)
 
     payload_bytes = bytes(payload)
 
@@ -253,7 +278,7 @@ def serialize_request(
         op_code,
         VERSION,
         0,  # flags
-        b"\x00" * 8,  # reserved
+        TABLE_HASH,
     )
 
     # Compute CRC32
@@ -269,7 +294,7 @@ def serialize_request(
         op_code,
         VERSION,
         0,  # flags
-        b"\x00" * 8,  # reserved
+        TABLE_HASH,
     )
 
     return header + payload_bytes
@@ -309,21 +334,21 @@ def parse_response_header(data: bytes) -> tuple[StatusCode, int, int, int]:
     Raises:
         IncompleteResponseError: If data is too short.
         InvalidMagicError: If magic number is invalid.
-        UnsupportedVersionError: If protocol version is unsupported.
-        InvalidChecksumError: If CRC32 validation fails.
+        TableMismatchError: If the server was built from another protocol
+            version or op table; its body can't be read.
     """
     if len(data) < HEADER_SIZE:
         raise IncompleteResponseError(f"Response too short: {len(data)} < {HEADER_SIZE}")
 
-    magic, data_len, request_id, crc, version, status, flags, _pad, reserved = struct.unpack(
+    magic, data_len, request_id, crc, version, status, flags, _pad, table_hash = struct.unpack(
         RESPONSE_HEADER_FORMAT, data[:HEADER_SIZE]
     )
 
     if magic != MAGIC:
         raise InvalidMagicError(f"Invalid magic: 0x{magic:08X} != 0x{MAGIC:08X}")
 
-    if version != VERSION:
-        raise UnsupportedVersionError(f"Unsupported version: {version} != {VERSION}")
+    if version != VERSION or table_hash != TABLE_HASH:
+        raise TableMismatchError(version, table_hash)
 
     return StatusCode(status), data_len, request_id, crc
 
@@ -362,117 +387,12 @@ def parse_response(data: bytes) -> RawResponse:
 # =============================================================================
 
 
-def parse_scan_response(data: bytes) -> ScanResult:
-    """Parse scan response data.
-
-    Format: [has_more:u8][cursor_len:u32][cursor:bytes][count:u32][entries...]
-    Entry format: [key_len:u16][key][value_len:u32][value]
-    """
-    if len(data) < 9:
-        raise IncompleteResponseError("Scan response too short")
-
-    offset = 0
-
-    # has_more
-    has_more = data[offset] != 0
-    offset += 1
-
-    # cursor
-    cursor_len = struct.unpack("<I", data[offset : offset + 4])[0]
-    offset += 4
-
-    if len(data) < offset + cursor_len:
-        raise IncompleteResponseError("Scan response cursor incomplete")
-
-    cursor = data[offset : offset + cursor_len] if cursor_len > 0 else None
-    offset += cursor_len
-
-    # count
-    if len(data) < offset + 4:
-        raise IncompleteResponseError("Scan response count incomplete")
-
-    count = struct.unpack("<I", data[offset : offset + 4])[0]
-    offset += 4
-
-    # entries
-    entries: list[KVEntry] = []
-    for _ in range(count):
-        # key
-        if len(data) < offset + 2:
-            raise IncompleteResponseError("Scan response key length incomplete")
-
-        key_len = struct.unpack("<H", data[offset : offset + 2])[0]
-        offset += 2
-
-        if len(data) < offset + key_len:
-            raise IncompleteResponseError("Scan response key incomplete")
-
-        key = data[offset : offset + key_len]
-        offset += key_len
-
-        # value
-        if len(data) < offset + 4:
-            raise IncompleteResponseError("Scan response value length incomplete")
-
-        value_len = struct.unpack("<I", data[offset : offset + 4])[0]
-        offset += 4
-
-        value: bytes | None = None
-        if value_len > 0:
-            if len(data) < offset + value_len:
-                raise IncompleteResponseError("Scan response value incomplete")
-            value = data[offset : offset + value_len]
-            offset += value_len
-
-        entries.append(KVEntry(key=key, value=value))
-
-    return ScanResult(entries=entries, cursor=cursor, has_more=has_more)
-
-
-def parse_history_response(data: bytes) -> list[VersionEntry]:
-    """Parse history response data.
-
-    Format: [count:u32][entries...]
-    Entry format: [version:u64][timestamp:i64][value_len:u32][value]
-    """
-    if len(data) < 4:
-        raise IncompleteResponseError("History response too short")
-
-    offset = 0
-
-    count = struct.unpack("<I", data[offset : offset + 4])[0]
-    offset += 4
-
-    entries: list[VersionEntry] = []
-    for _ in range(count):
-        if len(data) < offset + 20:
-            raise IncompleteResponseError("History response entry incomplete")
-
-        version = struct.unpack("<Q", data[offset : offset + 8])[0]
-        offset += 8
-
-        timestamp = struct.unpack("<q", data[offset : offset + 8])[0]
-        offset += 8
-
-        value_len = struct.unpack("<I", data[offset : offset + 4])[0]
-        offset += 4
-
-        if len(data) < offset + value_len:
-            raise IncompleteResponseError("History response value incomplete")
-
-        value = data[offset : offset + value_len]
-        offset += value_len
-
-        entries.append(VersionEntry(version=version, timestamp=timestamp, value=value))
-
-    return entries
-
-
 def parse_dequeue_response(data: bytes) -> DequeueResult:
     """Parse dequeue response data.
 
     Format: [count:u32][messages...]
     Message format: [seq:u64][payload_len:u32][payload]
+                    [enqueued_at_ms:i64][delivery_count:u32][priority:u8]
     """
     if len(data) < 4:
         raise IncompleteResponseError("Dequeue response too short")
@@ -499,7 +419,20 @@ def parse_dequeue_response(data: bytes) -> DequeueResult:
         payload = data[offset : offset + payload_len]
         offset += payload_len
 
-        messages.append(Message(seq=seq, payload=payload))
+        if len(data) < offset + 13:
+            raise IncompleteResponseError("Dequeue response message incomplete")
+        enqueued_at_ms, delivery_count, priority = struct.unpack("<qIB", data[offset : offset + 13])
+        offset += 13
+
+        messages.append(
+            Message(
+                seq=seq,
+                payload=payload,
+                enqueued_at_ms=enqueued_at_ms,
+                delivery_count=delivery_count,
+                priority=priority,
+            )
+        )
 
     return DequeueResult(messages=messages)
 
@@ -563,11 +496,10 @@ def parse_stream_read_response(data: bytes) -> StreamReadResult:
     count = struct.unpack("<I", data[pos : pos + 4])[0]
     pos += 4
 
+    # An answer that ends before its count of records is refused, not
+    # returned short: each field below checks it's there.
     records: list[StreamRecord] = []
     for _ in range(count):
-        if pos >= len(data):
-            break
-
         # Read sequence
         if pos + 8 > len(data):
             raise IncompleteResponseError("Stream record: missing sequence")
@@ -583,7 +515,10 @@ def parse_stream_read_response(data: bytes) -> StreamReadResult:
         # Read tier
         if pos + 1 > len(data):
             raise IncompleteResponseError("Stream record: missing tier")
-        tier = StorageTier(data[pos])
+        try:
+            tier = StorageTier(data[pos])
+        except ValueError:
+            raise IncompleteResponseError(f"Stream record: unknown tier {data[pos]}") from None
         pos += 1
 
         # Skip partition
@@ -597,7 +532,8 @@ def parse_stream_read_response(data: bytes) -> StreamReadResult:
         key_present = data[pos]
         pos += 1
 
-        # Skip key if present
+        # Read key (stream name) if present
+        stream_name = ""
         if key_present != 0:
             if pos + 4 > len(data):
                 raise IncompleteResponseError("Stream record: missing key length")
@@ -605,6 +541,7 @@ def parse_stream_read_response(data: bytes) -> StreamReadResult:
             pos += 4
             if pos + key_len > len(data):
                 raise IncompleteResponseError("Stream record: missing key data")
+            stream_name = data[pos : pos + key_len].decode("utf-8")
             pos += key_len
 
         # Read payload
@@ -618,17 +555,41 @@ def parse_stream_read_response(data: bytes) -> StreamReadResult:
         payload = data[pos : pos + payload_len]
         pos += payload_len
 
-        # Skip header_count (TODO: parse headers)
+        # Read headers: [header_count:u32]([key_len:u32][key][val_len:u32][val])*
         if pos + 4 > len(data):
             raise IncompleteResponseError("Stream record: missing header count")
+        header_count = struct.unpack("<I", data[pos : pos + 4])[0]
         pos += 4
 
+        headers = None
+        if header_count > 0:
+            headers = {}
+            for _ in range(header_count):
+                if pos + 4 > len(data):
+                    raise IncompleteResponseError("Stream record: missing header key length")
+                hkey_len = struct.unpack("<I", data[pos : pos + 4])[0]
+                pos += 4
+                if pos + hkey_len > len(data):
+                    raise IncompleteResponseError("Stream record: missing header key")
+                hkey = data[pos : pos + hkey_len].decode("utf-8")
+                pos += hkey_len
+                if pos + 4 > len(data):
+                    raise IncompleteResponseError("Stream record: missing header value length")
+                hval_len = struct.unpack("<I", data[pos : pos + 4])[0]
+                pos += 4
+                if pos + hval_len > len(data):
+                    raise IncompleteResponseError("Stream record: missing header value")
+                hval = data[pos : pos + hval_len].decode("utf-8")
+                pos += hval_len
+                headers[hkey] = hval
+
         records.append(
             StreamRecord(
                 id=StreamID(timestamp_ms=timestamp_ms, sequence=sequence),
                 tier=tier,
+                stream=stream_name,
                 payload=payload,
-                headers=None,
+                headers=headers,
             )
         )
 
@@ -661,6 +622,33 @@ def parse_stream_info_response(data: bytes) -> StreamInfo:
     )
 
 
+def build_stream_batch_value(
+    payload: bytes,
+    headers: dict[str, str] | None = None,
+) -> bytes:
+    """Frame a single record for stream_append (batch of one).
+
+    Wire format:
+    [record_count:u32=1][payload_len:u32][payload]
+    [header_count:u16]([key_len:u16][key][val_len:u16][val])*
+    """
+    hdrs = headers or {}
+    parts: list[bytes] = [
+        struct.pack("<I", 1),
+        struct.pack("<I", len(payload)),
+        payload,
+        struct.pack("<H", len(hdrs)),
+    ]
+    for key, val in hdrs.items():
+        key_bytes = key.encode("utf-8")
+        val_bytes = val.encode("utf-8")
+        parts.append(struct.pack("<H", len(key_bytes)))
+        parts.append(key_bytes)
+        parts.append(struct.pack("<H", len(val_bytes)))
+        parts.append(val_bytes)
+    return b"".join(parts)
+
+
 def serialize_group_value(group: str, consumer: str) -> bytes:
     """Serialize group and consumer names for consumer group operations.
 
@@ -697,95 +685,140 @@ def serialize_group_ack_value(group: str, ids: list[StreamID], consumer: str = "
     return bytes(result)
 
 
-# =============================================================================
-# Action/Worker Value Serialization
-# =============================================================================
+def serialize_group_pending_value(group: str, consumer: str = "") -> bytes:
+    """Serialize group (and optional consumer filter) for STREAM_GROUP_PENDING.
 
+    Format: [group_len:u16][group]([consumer_len:u16][consumer])?
 
-def serialize_action_register_value(
-    action_type: int,
-    timeout_ms: int,
-    max_retries: int,
-    description: str | None = None,
+    An empty ``consumer`` returns the whole group's PEL; a non-empty one filters
+    to that consumer's entries.
+    """
+    group_bytes = group.encode("utf-8")
+
+    result = bytearray()
+    result.extend(struct.pack("<H", len(group_bytes)))
+    result.extend(group_bytes)
+    if consumer:
+        consumer_bytes = consumer.encode("utf-8")
+        result.extend(struct.pack("<H", len(consumer_bytes)))
+        result.extend(consumer_bytes)
+    return bytes(result)
+
+
+def serialize_group_claim_value(
+    group: str,
+    consumer: str,
+    min_idle_ms: int,
+    start_id: StreamID,
+    count: int,
 ) -> bytes:
-    """Serialize action register value.
+    """Serialize the request value for STREAM_GROUP_CLAIM.
 
-    Format: [action_type:u8][timeout_ms:u32][max_retries:u32]
-            [has_desc:u8][desc_len:u16]?[desc]?
-            [has_wasm_module:u8=0]...(reserved zero fields for wire compat)
+    Format: [group_len:u16][group][consumer_len:u16][consumer]
+            [min_idle_ms:u32][start_ts:u64][start_seq:u64][count:u32]
     """
+    group_bytes = group.encode("utf-8")
+    consumer_bytes = consumer.encode("utf-8")
+
     result = bytearray()
+    result.extend(struct.pack("<H", len(group_bytes)))
+    result.extend(group_bytes)
+    result.extend(struct.pack("<H", len(consumer_bytes)))
+    result.extend(consumer_bytes)
+    result.extend(struct.pack("<I", min_idle_ms))
+    result.extend(struct.pack("<Q", start_id.timestamp_ms))
+    result.extend(struct.pack("<Q", start_id.sequence))
+    result.extend(struct.pack("<I", count))
+    return bytes(result)
 
-    # action_type
-    result.append(action_type & 0xFF)
 
-    # timeout_ms
-    result.extend(struct.pack("<I", timeout_ms))
+def parse_pending_entries(data: bytes) -> list[PendingEntry]:
+    """Parse the STREAM_GROUP_PENDING response.
 
-    # max_retries
-    result.extend(struct.pack("<I", max_retries))
+    Format: [count:u32]([ts:u64][seq:u64][delivery_count:u32]
+            [consumer_len:u16][consumer])*
+    """
+    if len(data) < 4:
+        return []
 
-    # description (optional)
-    if description:
-        desc_bytes = description.encode("utf-8")
-        result.append(1)  # has_desc
-        result.extend(struct.pack("<H", len(desc_bytes)))
-        result.extend(desc_bytes)
-    else:
-        result.append(0)
+    pos = 0
+    count = struct.unpack("<I", data[pos : pos + 4])[0]
+    pos += 4
 
-    # wasm_module (reserved, always zero)
-    result.append(0)
+    entries: list[PendingEntry] = []
+    for _ in range(count):
+        if pos + 8 + 8 + 4 + 2 > len(data):
+            raise IncompleteResponseError("Pending entry: truncated header")
+        ts = struct.unpack("<Q", data[pos : pos + 8])[0]
+        pos += 8
+        seq = struct.unpack("<Q", data[pos : pos + 8])[0]
+        pos += 8
+        delivery_count = struct.unpack("<I", data[pos : pos + 4])[0]
+        pos += 4
+        clen = struct.unpack("<H", data[pos : pos + 2])[0]
+        pos += 2
+        if pos + clen > len(data):
+            raise IncompleteResponseError("Pending entry: truncated consumer")
+        consumer = data[pos : pos + clen].decode("utf-8")
+        pos += clen
+        entries.append(
+            PendingEntry(
+                id=StreamID(timestamp_ms=ts, sequence=seq),
+                consumer=consumer,
+                delivery_count=delivery_count,
+            )
+        )
+    return entries
 
-    # wasm_entrypoint (reserved, always zero)
-    result.append(0)
 
-    # wasm_memory_limit (reserved, always zero)
-    result.append(0)
+# =============================================================================
+# Action/Worker Value Serialization
+# =============================================================================
 
-    # trigger_stream (optional, not used)
-    result.append(0)
 
-    # trigger_group (optional, not used)
-    result.append(0)
+def serialize_action_register_value(
+    action_type: int,
+    timeout_ms: int,
+    max_retries: int,
+    owner: str | None = None,
+) -> bytes:
+    """Serialize action register value.
 
-    return bytes(result)
+    Format: [action_type:u8][timeout_ms:u32][max_retries:u32][owner_len:u16][owner]
+    """
+    for name, value in (("timeout_ms", timeout_ms), ("max_retries", max_retries)):
+        if not 0 <= value < 2**32:
+            raise ValueError(f"{name} must be in [0, 2**32), not {value}")
+    owner_bytes = (owner or "").encode("utf-8")
+    if len(owner_bytes) > 0xFFFF:
+        raise ValueError("owner is longer than 65535 bytes")
+    return (
+        struct.pack("<BIIH", action_type & 0xFF, timeout_ms, max_retries, len(owner_bytes))
+        + owner_bytes
+    )
 
 
 def serialize_action_invoke_value(
     input_data: bytes,
-    priority: int = 10,
-    idempotency_key: str | None = None,
+    labels: dict[str, Any] | None = None,
 ) -> bytes:
     """Serialize action invoke value.
 
-    Format: [priority:u8][delay_ms:i64][has_caller:u8]
-            [has_idempotency_key:u8][key_len:u16]?[key]?[input...]
-    """
-    result = bytearray()
-
-    # priority
-    result.append(priority & 0xFF)
+    Format: [has_labels:u8]([labels_len:u16][labels...])?[input...]
 
-    # delay_ms (default 0)
-    result.extend(struct.pack("<q", 0))
+    has_labels is 0 or 1. labels is a compact JSON object; only workers whose
+    registered labels contain every key/value receive the run.
 
-    # caller_id (optional, none)
-    result.append(0)
-
-    # idempotency_key (optional)
-    if idempotency_key:
-        key_bytes = idempotency_key.encode("utf-8")
-        result.append(1)
-        result.extend(struct.pack("<H", len(key_bytes)))
-        result.extend(key_bytes)
-    else:
-        result.append(0)
-
-    # input
-    result.extend(input_data)
+    Raises:
+        ValueError: If the encoded labels exceed 65535 bytes.
+    """
+    if labels is None:
+        return b"\x00" + input_data
 
-    return bytes(result)
+    labels_bytes = json.dumps(labels, separators=(",", ":")).encode("utf-8")
+    if len(labels_bytes) > 0xFFFF:
+        raise ValueError(f"invoke labels are {len(labels_bytes)} bytes; the limit is 65535")
+    return b"\x01" + struct.pack("<H", len(labels_bytes)) + labels_bytes + input_data
 
 
 def serialize_action_list_value(limit: int = 100) -> bytes:
@@ -933,19 +966,13 @@ def serialize_worker_fail_value(
     return bytes(result)
 
 
-def serialize_worker_list_value(limit: int = 100) -> bytes:
-    """Serialize worker list value.
-
-    Format: [limit:u32]
-    """
-    return struct.pack("<I", limit)
-
-
 def parse_task_assignment(data: bytes) -> "types.TaskAssignment":
     """Parse task assignment from server response.
 
     Format: [task_id_len:u16][task_id][task_type_len:u16][task_type]
-            [created_at:i64][attempt:u32][payload...]
+            [created_at:i64][attempt:u32]
+            [has_caller:u8]([run_id_len:u16][run_id][wf_name_len:u16][wf_name])?
+            [payload...]
     """
     if len(data) < 10:
         raise ValueError("Incomplete task assignment response")
@@ -982,6 +1009,32 @@ def parse_task_assignment(data: bytes) -> "types.TaskAssignment":
     attempt = struct.unpack_from("<I", data, pos)[0]
     pos += 4
 
+    # caller block: [has_caller:u8][run_id_len:u16][run_id][wf_name_len:u16][wf_name]
+    if pos + 1 > len(data):
+        raise ValueError("Incomplete task assignment: missing has_caller")
+    has_caller = data[pos]
+    pos += 1
+    caller_run_id = ""
+    caller_workflow_name = ""
+    if has_caller == 1:
+        if pos + 2 > len(data):
+            raise ValueError("Incomplete task assignment: missing caller_run_id length")
+        crid_len = struct.unpack_from("<H", data, pos)[0]
+        pos += 2
+        if pos + crid_len > len(data):
+            raise ValueError("Incomplete task assignment: missing caller_run_id")
+        caller_run_id = data[pos : pos + crid_len].decode("utf-8")
+        pos += crid_len
+
+        if pos + 2 > len(data):
+            raise ValueError("Incomplete task assignment: missing caller_workflow_name length")
+        cwf_len = struct.unpack_from("<H", data, pos)[0]
+        pos += 2
+        if pos + cwf_len > len(data):
+            raise ValueError("Incomplete task assignment: missing caller_workflow_name")
+        caller_workflow_name = data[pos : pos + cwf_len].decode("utf-8")
+        pos += cwf_len
+
     # payload (rest of data)
     payload = data[pos:]
 
@@ -991,4 +1044,6 @@ def parse_task_assignment(data: bytes) -> "types.TaskAssignment":
         payload=payload,
         created_at=created_at,
         attempt=attempt,
+        caller_run_id=caller_run_id,
+        caller_workflow_name=caller_workflow_name,
     )