Add protocol state to all output batches and catalog results

- Update ProtocolOutput.metadata() in table_function.py to include
  vgi.protocol_state: output metadata on all output batches
- Update ProtocolOutput.metadata() in table_in_out_function.py similarly
- Update worker.py catalog result serialization to include protocol state
  on all catalog_result messages using IPC streams

This completes the protocol state requirement where all record batches
must include protocol state metadata for validation.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: rustyconover

vgi/table_function.py

@@ -151,25 +151,29 @@ class ProtocolOutput:
 
     def metadata(
         self, invocation: vgi.invocation.Invocation
-    ) -> pa.KeyValueMetadata | None:
+    ) -> pa.KeyValueMetadata:
         """Create metadata for this output based on the status.
 
         Args:
             invocation: The Invocation for this function invocation, passed through
                 to Message.add_to_metadata() for correlation information.
 
         Returns:
-            KeyValueMetadata containing status and optional log message fields.
+            KeyValueMetadata containing protocol state, status and optional log
+            message fields.
 
         """
-        metadata_dict: dict[str, str] = {}
+        # Start with protocol state (required for VGI protocol)
+        metadata_dict: dict[bytes, bytes] = {
+            vgi.ipc_utils.PROTOCOL_STATE_KEY: vgi.ipc_utils.ProtocolState.OUTPUT.encode()
+        }
 
+        # Add log message metadata if present
         if self.log_message is not None:
-            metadata_dict = self.log_message.add_to_metadata(invocation, metadata_dict)
+            for k, v in self.log_message.add_to_metadata(invocation, {}).items():
+                metadata_dict[k.encode()] = v.encode()
 
-        return pa.KeyValueMetadata(
-            {k.encode(): v.encode() for k, v in metadata_dict.items()}
-        )
+        return pa.KeyValueMetadata(metadata_dict)
 
     @classmethod
     def from_process_result(cls, process_result: OutputComplete) -> ProtocolOutput:

vgi/table_in_out_function.py

@@ -161,25 +161,30 @@ class ProtocolOutput:
 
     def metadata(
         self, invocation: vgi.invocation.Invocation
-    ) -> pa.KeyValueMetadata | None:
+    ) -> pa.KeyValueMetadata:
         """Create metadata for this output based on the status.
 
         Args:
             invocation: The Invocation for this function invocation, passed through
                 to Message.add_to_metadata() for correlation information.
 
         Returns:
-            KeyValueMetadata containing status and optional log message fields.
+            KeyValueMetadata containing protocol state, status and optional log
+            message fields.
 
         """
-        metadata_dict: dict[str, str] = {"vgi.status": self.status.value}
+        # Start with protocol state (required for VGI protocol) and status
+        metadata_dict: dict[bytes, bytes] = {
+            vgi.ipc_utils.PROTOCOL_STATE_KEY: vgi.ipc_utils.ProtocolState.OUTPUT.encode(),
+            b"vgi.status": self.status.value.encode(),
+        }
 
+        # Add log message metadata if present
         if self.log_message is not None:
-            metadata_dict = self.log_message.add_to_metadata(invocation, metadata_dict)
+            for k, v in self.log_message.add_to_metadata(invocation, {}).items():
+                metadata_dict[k.encode()] = v.encode()
 
-        return pa.KeyValueMetadata(
-            {k.encode(): v.encode() for k, v in metadata_dict.items()}
-        )
+        return pa.KeyValueMetadata(metadata_dict)
 
     @classmethod
     def from_process_result(

vgi/worker.py

@@ -974,28 +974,47 @@ def _handle_catalog_invocation(
 
         fn_log.debug("catalog_method_result", result=result)
 
-        # Serialize and stream result
+        # Serialize and stream result with protocol state metadata
         # Result types:
         # - None → empty batch (0 rows, 0 columns)
         # - list of primitives → convert to single-column batch (e.g., catalogs())
         # - Dataclass with serialize() → serialize to bytes, write
         # - Iterable of dataclasses → stream multiple serialized items
+        from vgi.ipc_utils import ProtocolState, protocol_state_metadata
+
+        catalog_result_metadata = protocol_state_metadata(ProtocolState.CATALOG_RESULT)
+
         if result is None:
             # Write empty batch to signal no result
             batch = pa.RecordBatch.from_pydict({})
-            stdout.write(batch.schema.serialize().to_pybytes())
-            stdout.write(batch.serialize().to_pybytes())
+            with pa.ipc.new_stream(cast(IOBase, stdout), batch.schema) as writer:
+                writer.write_batch(batch, custom_metadata=catalog_result_metadata)
         elif isinstance(result, list) and (
             not result or not hasattr(result[0], "serialize")
         ):
             # List of primitives (e.g., strings from catalogs())
             batch = pa.RecordBatch.from_pydict({"value": result})
-            stdout.write(batch.schema.serialize().to_pybytes())
-            stdout.write(batch.serialize().to_pybytes())
+            with pa.ipc.new_stream(cast(IOBase, stdout), batch.schema) as writer:
+                writer.write_batch(batch, custom_metadata=catalog_result_metadata)
         elif hasattr(result, "serialize"):
-            # Single dataclass result - write serialized bytes directly
-            result_bytes = result.serialize()
-            stdout.write(result_bytes)
+            # Single dataclass result - serialize with protocol state if possible
+            # Check if serialize() accepts custom_metadata parameter
+            import inspect
+
+            sig = inspect.signature(result.serialize)
+            if "custom_metadata" in sig.parameters:
+                result_bytes = result.serialize(custom_metadata=catalog_result_metadata)
+                stdout.write(result_bytes)
+            else:
+                # Fallback: use serialize() then we need to re-wrap with protocol state
+                # This is less efficient but ensures protocol state is included
+                result_bytes = result.serialize()
+                # Deserialize and re-serialize with protocol state
+                from vgi.ipc_utils import deserialize_record_batch
+
+                batch, _ = deserialize_record_batch(result_bytes)
+                with pa.ipc.new_stream(cast(IOBase, stdout), batch.schema) as writer:
+                    writer.write_batch(batch, custom_metadata=catalog_result_metadata)
         elif isinstance(result, list) and result and hasattr(result[0], "to_row_dict"):
             # List of catalog objects with to_row_dict() - use efficient batch writing
             # Determine the schema based on method and type parameter
@@ -1045,12 +1064,12 @@ def _handle_catalog_invocation(
                 # Fallback: use the ARROW_SCHEMA from the first item's class
                 schema = type(result[0]).ARROW_SCHEMA
 
-            # Collect all rows and write as single batch
+            # Collect all rows and write as single batch with protocol state
             rows = [item.to_row_dict() for item in result]
             batch = pa.RecordBatch.from_pylist(rows, schema=schema)
 
             with pa.ipc.new_stream(cast(IOBase, stdout), schema) as writer:
-                writer.write_batch(batch)
+                writer.write_batch(batch, custom_metadata=catalog_result_metadata)
         elif isinstance(result, list) and not result:
             # Empty list - need to determine schema for empty batch
             from vgi.catalog import (
@@ -1086,7 +1105,7 @@ def _handle_catalog_invocation(
 
             batch = pa.RecordBatch.from_pylist([], schema=schema)
             with pa.ipc.new_stream(cast(IOBase, stdout), schema) as writer:
-                writer.write_batch(batch)
+                writer.write_batch(batch, custom_metadata=catalog_result_metadata)
         else:
             raise TypeError(
                 f"Catalog method returned unsupported type: "