@query-farm/vgi-rpc 0.6.4 → 0.7.1

package/src/constants.ts

@@ -1,25 +1,57 @@
 // © Copyright 2025-2026, Query.Farm LLC - https://query.farm
 // SPDX-License-Identifier: Apache-2.0
 
-/** Well-known metadata keys matching Python's metadata.py */
+// Well-known metadata keys matching Python's metadata.py.
 
+/** Batch-metadata key carrying the invoked RPC method name. */
 export const RPC_METHOD_KEY = "vgi_rpc.method";
+/** Batch-metadata key carrying a log batch's severity level. */
 export const LOG_LEVEL_KEY = "vgi_rpc.log_level";
+/** Batch-metadata key carrying a log batch's message text. */
 export const LOG_MESSAGE_KEY = "vgi_rpc.log_message";
+/** Batch-metadata key carrying a log batch's structured extra fields. */
 export const LOG_EXTRA_KEY = "vgi_rpc.log_extra";
+/** Batch-metadata key carrying the wire request-framing version. */
 export const REQUEST_VERSION_KEY = "vgi_rpc.request_version";
+/** Current wire request-framing version. Distinct from the application-level
+ *  {@link PROTOCOL_VERSION_KEY protocol version}. */
 export const REQUEST_VERSION = "1";
 
+/** Batch-metadata key identifying the server instance that produced a batch. */
 export const SERVER_ID_KEY = "vgi_rpc.server_id";
+/** Batch-metadata key carrying the client-supplied request id. */
 export const REQUEST_ID_KEY = "vgi_rpc.request_id";
 
+/** Batch-metadata key carrying the service / protocol name. */
 export const PROTOCOL_NAME_KEY = "vgi_rpc.protocol_name";
+/** Batch-metadata key carrying the `__describe__` response schema version. */
 export const DESCRIBE_VERSION_KEY = "vgi_rpc.describe_version";
-export const DESCRIBE_VERSION = "3";
+export const PROTOCOL_HASH_KEY = "vgi_rpc.protocol_hash";
+/** Current `__describe__` response schema version (the slim 8-column schema). */
+export const DESCRIBE_VERSION = "4";
 
+/** Application protocol surface version. Carried on every request batch from
+ *  a client bound to a Protocol that declares `protocolVersion`; also emitted
+ *  in the __describe__ response metadata. Format: canonical semver
+ *  MAJOR.MINOR.PATCH. Enforced at the dispatch boundary on the server: exact
+ *  major+minor match required, patch ignored. Distinct from `REQUEST_VERSION`
+ *  (wire framing). Mirrors Python's `PROTOCOL_VERSION_KEY`. */
+export const PROTOCOL_VERSION_KEY = "vgi_rpc.protocol_version";
+
+/** Reserved method name for the introspection (`__describe__`) call. */
 export const DESCRIBE_METHOD_NAME = "__describe__";
 
+/** Batch-metadata key carrying the base64-encoded stream continuation/state token. */
 export const STATE_KEY = "vgi_rpc.stream_state#b64";
+export const CANCEL_KEY = "vgi_rpc.cancel";
 
 export const LOCATION_KEY = "vgi_rpc.location";
 export const LOCATION_SHA256_KEY = "vgi_rpc.location.sha256";
+
+/** HTTP response header set when an RPC error is returned over the HTTP transport. */
+export const RPC_ERROR_HEADER = "X-VGI-RPC-Error";
+
+/** Top-level metadata key on an EXCEPTION batch identifying the error category.
+ *  Hoisted by `buildErrorBatch` when the thrown error has a static or instance
+ *  `errorKind` property. Mirrors Python's `vgi_rpc.metadata.ERROR_KIND_KEY`. */
+export const ERROR_KIND_KEY = "vgi_rpc.error_kind";

package/src/crypto.ts

@@ -0,0 +1,95 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Generic AEAD seal/open primitives shared by stream-state and sticky-session
+ * tokens. Mirrors Python's `vgi_rpc.crypto` module: a tiny envelope around
+ * XChaCha20-Poly1305 with a leading version byte so future format bumps stay
+ * self-describing.
+ *
+ * Wire format (returned by {@link sealBytes}, accepted by {@link openBytes}):
+ *
+ * ```
+ *   [1B  version  (1..255)]
+ *   [24B nonce    (XChaCha20-Poly1305, random per envelope)]
+ *   [..  ciphertext + 16B Poly1305 tag]
+ * ```
+ *
+ * The plaintext frame is fully up to the caller — only the version + nonce +
+ * tag overhead is fixed. AAD (`aad`) is bound at the crypto layer so an
+ * envelope sealed for one identity fails decryption when presented by another.
+ */
+
+import { xchacha20poly1305 } from "@noble/ciphers/chacha.js";
+import { randomBytes } from "./util/web-crypto.js";
+
+const NONCE_LEN = 24;
+const TAG_LEN = 16;
+const VERSION_LEN = 1;
+const MIN_ENVELOPE_LEN = VERSION_LEN + NONCE_LEN + TAG_LEN;
+
+/** Thrown by {@link openBytes} for any envelope it cannot open — malformed,
+ *  tampered, wrong key, wrong AAD, wrong version, truncated, all surface the
+ *  same way so callers cannot distinguish failure modes via error content. */
+export class SealError extends Error {
+  constructor(message: string) {
+    super(message);
+    this.name = "SealError";
+  }
+}
+
+/** Normalise a key to 32 bytes by SHA-256 hashing when it isn't already 32B.
+ *  Mirrors Python's `normalize_key` so any callers can pass operator-provided
+ *  keys of arbitrary length without a separate stretching step. */
+export async function normalizeKey(key: Uint8Array): Promise<Uint8Array> {
+  if (key.length === 32) return key;
+  const digest = await crypto.subtle.digest("SHA-256", key as BufferSource);
+  return new Uint8Array(digest);
+}
+
+export interface SealOptions {
+  /** Associated data bound at the crypto layer — typically a principal or
+   *  request-scoped identifier. Must match between seal and open. */
+  aad: Uint8Array;
+  /** Envelope version byte. Defaults to 1; carry through to {@link openBytes}. */
+  version?: number;
+}
+
+/** Seal `plaintext` under `key` with AEAD, returning the wire envelope. */
+export function sealBytes(plaintext: Uint8Array, key: Uint8Array, opts: SealOptions): Uint8Array {
+  if (key.length !== 32) {
+    throw new Error("AEAD key must be 32 bytes — call normalizeKey() first");
+  }
+  const version = opts.version ?? 1;
+  if (version < 1 || version > 255) {
+    throw new Error(`AEAD envelope version must fit in one byte; got ${version}`);
+  }
+  const nonce = randomBytes(NONCE_LEN);
+  const ciphertext = xchacha20poly1305(key, nonce, opts.aad as Uint8Array).encrypt(plaintext);
+  const wire = new Uint8Array(VERSION_LEN + NONCE_LEN + ciphertext.length);
+  wire[0] = version;
+  wire.set(nonce, VERSION_LEN);
+  wire.set(ciphertext, VERSION_LEN + NONCE_LEN);
+  return wire;
+}
+
+/** Open and verify an envelope produced by {@link sealBytes}. */
+export function openBytes(envelope: Uint8Array, key: Uint8Array, opts: SealOptions): Uint8Array {
+  if (key.length !== 32) {
+    throw new Error("AEAD key must be 32 bytes — call normalizeKey() first");
+  }
+  if (envelope.length < MIN_ENVELOPE_LEN) {
+    throw new SealError("envelope truncated");
+  }
+  const expectedVersion = opts.version ?? 1;
+  if (envelope[0] !== expectedVersion) {
+    throw new SealError(`unsupported envelope version: ${envelope[0]}`);
+  }
+  const nonce = envelope.subarray(VERSION_LEN, VERSION_LEN + NONCE_LEN);
+  const ciphertext = envelope.subarray(VERSION_LEN + NONCE_LEN);
+  try {
+    return xchacha20poly1305(key, nonce, opts.aad as Uint8Array).decrypt(ciphertext);
+  } catch {
+    throw new SealError("envelope verification failed");
+  }
+}

package/src/dispatch/describe.ts

@@ -2,162 +2,201 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import {
-  Binary,
-  Bool,
-  Field,
-  makeData,
-  RecordBatch,
-  Schema,
-  Struct,
-  Utf8,
-  vectorFromArray,
-} from "@query-farm/apache-arrow";
+  batchFromColumns,
+  binary,
+  bool,
+  field,
+  schema as makeSchema,
+  utf8,
+  type VgiBatch,
+  withBatchMetadata,
+} from "../arrow/index.js";
 import {
   DESCRIBE_VERSION,
   DESCRIBE_VERSION_KEY,
+  PROTOCOL_HASH_KEY,
   PROTOCOL_NAME_KEY,
+  PROTOCOL_VERSION_KEY,
   REQUEST_VERSION,
   REQUEST_VERSION_KEY,
   SERVER_ID_KEY,
 } from "../constants.js";
 import type { MethodDefinition } from "../types.js";
 import { serializeSchema } from "../util/schema.js";
+import { sha256Hex } from "../util/web-crypto.js";
 
 /**
- * The schema for the __describe__ response, matching Python's _DESCRIBE_SCHEMA.
+ * Slim DESCRIBE_VERSION 4 schema. Python-flavoured fields (doc,
+ * param_types_json, param_defaults_json, param_docs_json) are not on the
+ * wire — Arrow IPC schema bytes are the authoritative type information;
+ * everything else is source-level metadata that callers should consult the
+ * Protocol class for.
  */
-export const DESCRIBE_SCHEMA = new Schema([
-  new Field("name", new Utf8(), false),
-  new Field("method_type", new Utf8(), false),
-  new Field("doc", new Utf8(), true),
-  new Field("has_return", new Bool(), false),
-  new Field("params_schema_ipc", new Binary(), false),
-  new Field("result_schema_ipc", new Binary(), false),
-  new Field("param_types_json", new Utf8(), true),
-  new Field("param_defaults_json", new Utf8(), true),
-  new Field("has_header", new Bool(), false),
-  new Field("header_schema_ipc", new Binary(), true),
-  new Field("is_exchange", new Bool(), true),
-  new Field("param_docs_json", new Utf8(), true),
+export const DESCRIBE_SCHEMA = makeSchema([
+  field("name", utf8(), false),
+  field("method_type", utf8(), false),
+  field("has_return", bool(), false),
+  field("params_schema_ipc", binary(), false),
+  field("result_schema_ipc", binary(), false),
+  field("has_header", bool(), false),
+  field("header_schema_ipc", binary(), true),
+  field("is_exchange", bool(), true),
 ]);
 
+/**
+ * Compute the SHA-256 hex digest of the canonical describe payload. Mirrors
+ * Python's vgi_rpc.introspect.compute_protocol_hash byte-for-byte, so two
+ * implementations of the same Protocol that produce identical Arrow IPC
+ * schema bytes for params/result/header will hash to the same value.
+ */
+async function computeProtocolHash(
+  protocolName: string,
+  rows: ReadonlyArray<{
+    name: string;
+    methodType: string;
+    hasReturn: boolean;
+    hasHeader: boolean;
+    isExchange: boolean | null;
+    paramsIpc: Uint8Array;
+    resultIpc: Uint8Array;
+    headerIpc: Uint8Array | null;
+  }>,
+): Promise<string> {
+  // Web Crypto's `subtle.digest` takes a single buffer, so we accumulate the
+  // canonical byte stream into a single Uint8Array and hash once. The byte
+  // sequence below must remain byte-for-byte identical to the Python
+  // implementation in vgi_rpc.introspect.compute_protocol_hash.
+  const enc = new TextEncoder();
+  const parts: Uint8Array[] = [];
+  const push = (v: string | Uint8Array) => parts.push(typeof v === "string" ? enc.encode(v) : v);
+
+  push("vgi_rpc.describe.v");
+  push(DESCRIBE_VERSION);
+  push("|");
+  push(REQUEST_VERSION);
+  push("|");
+  push(protocolName);
+  push("|");
+  for (const r of rows) {
+    push(Uint8Array.of(0x1f));
+    push(r.name);
+    push(Uint8Array.of(0x1e));
+    push(r.methodType);
+    push(Uint8Array.of(0x1e));
+    push(r.hasReturn ? "1" : "0");
+    push(Uint8Array.of(0x1e));
+    push(r.hasHeader ? "1" : "0");
+    push(Uint8Array.of(0x1e));
+    push(r.isExchange === null ? "-" : r.isExchange ? "1" : "0");
+    push(Uint8Array.of(0x1e));
+    push(r.paramsIpc);
+    push(Uint8Array.of(0x1e));
+    push(r.resultIpc);
+    push(Uint8Array.of(0x1e));
+    if (r.headerIpc) push(r.headerIpc);
+  }
+
+  let total = 0;
+  for (const p of parts) total += p.length;
+  const buf = new Uint8Array(total);
+  let off = 0;
+  for (const p of parts) {
+    buf.set(p, off);
+    off += p.length;
+  }
+  return sha256Hex(buf);
+}
+
 /**
  * Build the __describe__ response batch and metadata.
  */
-export function buildDescribeBatch(
+export async function buildDescribeBatch(
   protocolName: string,
   methods: Map<string, MethodDefinition>,
   serverId: string,
-): { batch: RecordBatch; metadata: Map<string, string> } {
+  protocolVersion?: string,
+): Promise<{ batch: VgiBatch; metadata: Map<string, string> }> {
   // Sort methods by name for consistent ordering
   const sortedEntries = [...methods.entries()].sort(([a], [b]) => a.localeCompare(b));
 
   const names: (string | null)[] = [];
   const methodTypes: (string | null)[] = [];
-  const docs: (string | null)[] = [];
   const hasReturns: boolean[] = [];
   const paramsSchemas: (Uint8Array | null)[] = [];
   const resultSchemas: (Uint8Array | null)[] = [];
-  const paramTypesJsons: (string | null)[] = [];
-  const paramDefaultsJsons: (string | null)[] = [];
   const hasHeaders: boolean[] = [];
   const headerSchemas: (Uint8Array | null)[] = [];
   const isExchanges: (boolean | null)[] = [];
-  const paramDocsJsons: (string | null)[] = [];
+
+  const hashRows: Array<{
+    name: string;
+    methodType: string;
+    hasReturn: boolean;
+    hasHeader: boolean;
+    isExchange: boolean | null;
+    paramsIpc: Uint8Array;
+    resultIpc: Uint8Array;
+    headerIpc: Uint8Array | null;
+  }> = [];
 
   for (const [name, method] of sortedEntries) {
     names.push(name);
     methodTypes.push(method.type);
-    docs.push(method.doc ?? null);
 
-    // Unary methods with non-empty result schema have a return value
     const hasReturn = method.type === "unary" && method.resultSchema.fields.length > 0;
     hasReturns.push(hasReturn);
 
-    paramsSchemas.push(serializeSchema(method.paramsSchema));
-    resultSchemas.push(serializeSchema(method.resultSchema));
-
-    // Build param_types_json
-    if (method.paramTypes && Object.keys(method.paramTypes).length > 0) {
-      paramTypesJsons.push(JSON.stringify(method.paramTypes));
-    } else {
-      paramTypesJsons.push(null);
-    }
-
-    // Build param_defaults_json
-    if (method.defaults && Object.keys(method.defaults).length > 0) {
-      const safe: Record<string, any> = {};
-      for (const [k, v] of Object.entries(method.defaults)) {
-        if (v === null || typeof v === "string" || typeof v === "number" || typeof v === "boolean") {
-          safe[k] = v;
-        }
-      }
-      paramDefaultsJsons.push(Object.keys(safe).length > 0 ? JSON.stringify(safe) : null);
-    } else {
-      paramDefaultsJsons.push(null);
-    }
-
-    hasHeaders.push(!!method.headerSchema);
-    headerSchemas.push(method.headerSchema ? serializeSchema(method.headerSchema) : null);
-
-    // is_exchange: true for exchange, false for producer, null for unary
-    if (method.exchangeFn) {
-      isExchanges.push(true);
-    } else if (method.producerFn) {
-      isExchanges.push(false);
-    } else {
-      isExchanges.push(null);
-    }
-
-    // param_docs_json: no docstring source in TypeScript, always null
-    paramDocsJsons.push(null);
+    const paramsIpc = serializeSchema(method.paramsSchema);
+    const resultIpc = serializeSchema(method.resultSchema);
+    paramsSchemas.push(paramsIpc);
+    resultSchemas.push(resultIpc);
+
+    const hasHeader = !!method.headerSchema;
+    hasHeaders.push(hasHeader);
+    const headerIpc = method.headerSchema ? serializeSchema(method.headerSchema) : null;
+    headerSchemas.push(headerIpc);
+
+    let isExchange: boolean | null;
+    if (method.exchangeFn) isExchange = true;
+    else if (method.producerFn) isExchange = false;
+    else isExchange = null;
+    isExchanges.push(isExchange);
+
+    hashRows.push({
+      name,
+      methodType: method.type,
+      hasReturn,
+      hasHeader,
+      isExchange,
+      paramsIpc,
+      resultIpc,
+      headerIpc,
+    });
   }
 
-  // Build the batch using vectorFromArray for each column
-  const nameArr = vectorFromArray(names, new Utf8());
-  const methodTypeArr = vectorFromArray(methodTypes, new Utf8());
-  const docArr = vectorFromArray(docs, new Utf8());
-  const hasReturnArr = vectorFromArray(hasReturns, new Bool());
-  const paramsSchemaArr = vectorFromArray(paramsSchemas, new Binary());
-  const resultSchemaArr = vectorFromArray(resultSchemas, new Binary());
-  const paramTypesArr = vectorFromArray(paramTypesJsons, new Utf8());
-  const paramDefaultsArr = vectorFromArray(paramDefaultsJsons, new Utf8());
-  const hasHeaderArr = vectorFromArray(hasHeaders, new Bool());
-  const headerSchemaArr = vectorFromArray(headerSchemas, new Binary());
-  const isExchangeArr = vectorFromArray(isExchanges, new Bool());
-  const paramDocsArr = vectorFromArray(paramDocsJsons, new Utf8());
-
-  const children = [
-    nameArr.data[0],
-    methodTypeArr.data[0],
-    docArr.data[0],
-    hasReturnArr.data[0],
-    paramsSchemaArr.data[0],
-    resultSchemaArr.data[0],
-    paramTypesArr.data[0],
-    paramDefaultsArr.data[0],
-    hasHeaderArr.data[0],
-    headerSchemaArr.data[0],
-    isExchangeArr.data[0],
-    paramDocsArr.data[0],
-  ];
-
-  const structType = new Struct(DESCRIBE_SCHEMA.fields);
-  const data = makeData({
-    type: structType,
-    length: sortedEntries.length,
-    children,
-    nullCount: 0,
+  const baseBatch = batchFromColumns(DESCRIBE_SCHEMA, {
+    name: names,
+    method_type: methodTypes,
+    has_return: hasReturns,
+    params_schema_ipc: paramsSchemas,
+    result_schema_ipc: resultSchemas,
+    has_header: hasHeaders,
+    header_schema_ipc: headerSchemas,
+    is_exchange: isExchanges,
   });
 
-  // Build metadata for the batch
+  const protocolHash = await computeProtocolHash(protocolName, hashRows);
+
   const metadata = new Map<string, string>();
   metadata.set(PROTOCOL_NAME_KEY, protocolName);
   metadata.set(REQUEST_VERSION_KEY, REQUEST_VERSION);
   metadata.set(DESCRIBE_VERSION_KEY, DESCRIBE_VERSION);
+  metadata.set(PROTOCOL_HASH_KEY, protocolHash);
   metadata.set(SERVER_ID_KEY, serverId);
+  if (protocolVersion) {
+    metadata.set(PROTOCOL_VERSION_KEY, protocolVersion);
+  }
 
-  const batch = new RecordBatch(DESCRIBE_SCHEMA, data, metadata);
-
+  const batch = withBatchMetadata(baseBatch, metadata);
   return { batch, metadata };
 }

package/src/dispatch/stream.ts

@@ -1,16 +1,16 @@
 // © Copyright 2025-2026, Query.Farm LLC - https://query.farm
 // SPDX-License-Identifier: Apache-2.0
 
-import { Schema } from "@query-farm/apache-arrow";
+import { conformBatchToSchema, schema as makeSchema, withBatchMetadata } from "../arrow/index.js";
+import { CANCEL_KEY } from "../constants.js";
 import { type ExternalLocationConfig, maybeExternalizeBatch } from "../external.js";
-import type { MethodDefinition } from "../types.js";
+import type { MethodDefinition, TransportKind } from "../types.js";
 import { OutputCollector } from "../types.js";
-import { conformBatchToSchema } from "../util/conform.js";
 import type { IpcStreamReader } from "../wire/reader.js";
 import { buildErrorBatch, buildResultBatch } from "../wire/response.js";
 import type { IpcStreamWriter } from "../wire/writer.js";
 
-const EMPTY_SCHEMA = new Schema([]);
+const EMPTY_SCHEMA = makeSchema([]);
 
 /**
  * Dispatch a stream RPC call (producer or exchange).
@@ -35,6 +35,7 @@ export async function dispatchStream(
   serverId: string,
   requestId: string | null,
   externalConfig?: ExternalLocationConfig,
+  kind?: TransportKind,
 ): Promise<void> {
   const isProducer = !!method.producerFn;
 
@@ -48,7 +49,7 @@ export async function dispatchStream(
   } catch (error: any) {
     const errSchema = method.headerSchema ?? EMPTY_SCHEMA;
     const errBatch = buildErrorBatch(errSchema, error, serverId, requestId);
-    writer.writeStream(errSchema, [errBatch]);
+    await writer.writeStream(errSchema, [errBatch]);
     // Still need to consume the input stream from the client
     const inputSchema = await reader.openNextStream();
     if (inputSchema) {
@@ -71,14 +72,14 @@ export async function dispatchStream(
   // Write header IPC stream if method has a header schema
   if (method.headerSchema && method.headerInit) {
     try {
-      const headerOut = new OutputCollector(method.headerSchema, true, serverId, requestId);
+      const headerOut = new OutputCollector(method.headerSchema, true, serverId, requestId, undefined, undefined, kind);
       const headerValues = method.headerInit(params, state, headerOut);
       const headerBatch = buildResultBatch(method.headerSchema, headerValues, serverId, requestId);
       const headerBatches = [...headerOut.batches.map((b) => b.batch), headerBatch];
-      writer.writeStream(method.headerSchema, headerBatches);
+      await writer.writeStream(method.headerSchema, headerBatches);
     } catch (error: any) {
       const errBatch = buildErrorBatch(method.headerSchema, error, serverId, requestId);
-      writer.writeStream(method.headerSchema, [errBatch]);
+      await writer.writeStream(method.headerSchema, [errBatch]);
       // Drain input stream so client doesn't hang
       const inputSchema = await reader.openNextStream();
       if (inputSchema) {
@@ -92,7 +93,7 @@ export async function dispatchStream(
   const inputSchema = await reader.openNextStream();
   if (!inputSchema) {
     const errBatch = buildErrorBatch(outputSchema, new Error("Expected input stream but got EOF"), serverId, requestId);
-    writer.writeStream(outputSchema, [errBatch]);
+    await writer.writeStream(outputSchema, [errBatch]);
     return;
   }
 
@@ -101,32 +102,54 @@ export async function dispatchStream(
   // same stream. We use IncrementalStream which writes bytes synchronously.
   const stream = writer.openStream(outputSchema);
 
-  // Expected input schema for casting compatible types (e.g., decimal→double)
-  const expectedInputSchema = method.inputSchema;
+  // Expected input schema for casting compatible types (e.g., decimal→double).
+  // State.__inputSchema overrides the method-registration schema per call,
+  // mirroring the __outputSchema pattern. Used by dynamic-input exchange
+  // methods (e.g., VGI's init, which binds to a user-supplied input shape).
+  const expectedInputSchema = state?.__inputSchema ?? method.inputSchema;
 
   try {
     while (true) {
       let inputBatch = await reader.readNextBatch();
       if (!inputBatch) break;
 
+      // Client cancellation: if the input batch carries vgi_rpc.cancel metadata,
+      // end the stream cleanly without calling the producer/exchange handler.
+      // The onCancel hook (if registered) runs once so state objects can
+      // release resources.
+      if (inputBatch.metadata?.get(CANCEL_KEY)) {
+        if (method.onCancel) {
+          try {
+            await method.onCancel(state);
+          } catch (err) {
+            console.debug?.(`onCancel hook failed: ${err instanceof Error ? err.message : err}`);
+          }
+        }
+        break;
+      }
+
       // Cast compatible input types when schema doesn't match exactly.
-      // If conformance fails (e.g., completely different schemas like a dummy
-      // registration schema vs actual data), pass the original batch through —
-      // the exchange handler may handle dynamic schemas internally.
-      if (expectedInputSchema && !isProducer && inputBatch.schema !== expectedInputSchema) {
+      // Gated on effectiveProducer (not isProducer) so methods that flip to
+      // producer mode via state.__isProducer skip the conform entirely — the
+      // tick batches they receive have a dummy shape that shouldn't be checked.
+      // Any conformance failure falls through with the original batch; the
+      // exchange handler owns input-shape validation if it cares.
+      if (expectedInputSchema && !effectiveProducer && inputBatch.schema !== expectedInputSchema) {
         try {
           inputBatch = conformBatchToSchema(inputBatch, expectedInputSchema);
         } catch (e) {
-          if (e instanceof TypeError) {
-            // Field name/count mismatch — propagate as error (matches Python behavior).
-            throw e;
-          }
-          // Other conformance failures: pass through for dynamic schema handlers.
+          // Field name/count mismatch (TypeError) is a hard contract violation —
+          // propagate as an RpcError so callers see a structured failure instead
+          // of a downstream hang or silent garbage. Other conform failures (e.g.
+          // type-cast issues for dynamic-input handlers) fall through with the
+          // original batch — handlers that bind their input shape per-call
+          // should set state.__inputSchema so the conform doesn't run at all.
+          if (e instanceof TypeError) throw e;
           console.debug?.(`Schema conformance skipped: ${e instanceof Error ? e.message : e}`);
         }
       }
 
-      const out = new OutputCollector(outputSchema, effectiveProducer, serverId, requestId);
+      const out = new OutputCollector(outputSchema, effectiveProducer, serverId, requestId, undefined, undefined, kind);
 
       if (isProducer) {
         await method.producerFn!(state, out);
@@ -139,7 +162,13 @@ export async function dispatchStream(
         if (externalConfig) {
           batch = await maybeExternalizeBatch(batch, externalConfig);
         }
-        stream.write(batch);
+        // Attach per-emit metadata (e.g. vgi_batch_index,
+        // vgi_partition_values#b64) as the RecordBatch message's
+        // custom_metadata so the C++ extension can read it off the wire.
+        if (emitted.metadata && emitted.metadata.size > 0) {
+          batch = withBatchMetadata(batch, emitted.metadata);
+        }
+        await stream.write(batch);
       }
 
       if (out.finished) {
@@ -147,10 +176,10 @@ export async function dispatchStream(
       }
     }
   } catch (error: any) {
-    stream.write(buildErrorBatch(outputSchema, error, serverId, requestId));
+    await stream.write(buildErrorBatch(outputSchema, error, serverId, requestId));
   }
 
-  stream.close();
+  await stream.close();
 
   // Drain remaining input so transport stays synchronized for next request.
   // Matches Python's _drain_stream() called after every streaming method.

package/src/dispatch/unary.ts

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import { type ExternalLocationConfig, maybeExternalizeBatch } from "../external.js";
-import type { MethodDefinition } from "../types.js";
+import type { MethodDefinition, TransportKind } from "../types.js";
 import { OutputCollector } from "../types.js";
 import { buildErrorBatch, buildResultBatch } from "../wire/response.js";
 import type { IpcStreamWriter } from "../wire/writer.js";
@@ -19,9 +19,10 @@ export async function dispatchUnary(
   serverId: string,
   requestId: string | null,
   externalConfig?: ExternalLocationConfig,
+  kind?: TransportKind,
 ): Promise<void> {
   const schema = method.resultSchema;
-  const out = new OutputCollector(schema, true, serverId, requestId);
+  const out = new OutputCollector(schema, true, serverId, requestId, undefined, undefined, kind);
 
   try {
     const result = await method.handler!(params, out);
@@ -31,9 +32,9 @@ export async function dispatchUnary(
     }
     // Collect log batches (from clientLog) + result batch
     const batches = [...out.batches.map((b) => b.batch), resultBatch];
-    writer.writeStream(schema, batches);
+    await writer.writeStream(schema, batches);
   } catch (error: any) {
     const batch = buildErrorBatch(schema, error, serverId, requestId);
-    writer.writeStream(schema, [batch]);
+    await writer.writeStream(schema, [batch]);
   }
 }