@query-farm/vgi-rpc 0.6.4 → 0.7.0

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]);
   }
 }

package/src/errors.ts

@@ -20,3 +20,79 @@ export class VersionError extends Error {
     this.name = "VersionError";
   }
 }
+
+/** Well-known values for the `vgi_rpc.error_kind` batch metadata key. Mirrors
+ *  Python's `vgi_rpc.metadata.ERROR_KIND_*` constants. */
+export const ERROR_KIND_METHOD_NOT_IMPLEMENTED = "method_not_implemented";
+export const ERROR_KIND_SESSION_LOST = "session_lost";
+export const ERROR_KIND_SERVER_DRAINING = "server_draining";
+export const ERROR_KIND_PROTOCOL_VERSION_MISMATCH = "protocol_version_mismatch";
+
+/** Raised when the client's declared `vgi_rpc.protocol_version` is
+ *  incompatible with the server's. Subclass of `VersionError` so existing
+ *  catch sites continue to write a typed error stream and keep serving.
+ *  Carries a directional message that tells the reader which side to
+ *  upgrade. Mirrors Python's `vgi_rpc.rpc.ProtocolVersionError`. */
+export class ProtocolVersionError extends VersionError {
+  static readonly errorKind = ERROR_KIND_PROTOCOL_VERSION_MISMATCH;
+  readonly errorKind = ERROR_KIND_PROTOCOL_VERSION_MISMATCH;
+  constructor(message: string) {
+    super(message);
+    this.name = "ProtocolVersionError";
+  }
+}
+
+const SEMVER_REGEX = /^(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)$/;
+
+/** Parse a canonical semver string into `[major, minor, patch]`. Throws on
+ *  any input that isn't `MAJOR.MINOR.PATCH` with non-negative integers and
+ *  no leading zeros (except literal `0`). No prereleases, no build metadata.
+ *  Mirrors Python's `vgi_rpc.metadata.parse_version`. */
+export function parseProtocolVersion(value: string): [number, number, number] {
+  const m = SEMVER_REGEX.exec(value);
+  if (!m) {
+    throw new Error(
+      `Invalid protocol version '${value}': expected canonical semver ` +
+        "MAJOR.MINOR.PATCH with non-negative integers and no leading zeros " +
+        "(no prereleases or build metadata).",
+    );
+  }
+  return [Number(m[1]), Number(m[2]), Number(m[3])];
+}
+
+/** Raised when a client invokes a method the server does not implement.
+ *
+ *  Mirrors Python's `vgi_rpc.rpc.MethodNotImplementedError`. The static
+ *  `errorKind` is hoisted onto the error batch metadata as
+ *  `vgi_rpc.error_kind` so clients can branch on the typed marker without
+ *  string-matching the message.
+ */
+export class MethodNotImplementedError extends Error {
+  static readonly errorKind = ERROR_KIND_METHOD_NOT_IMPLEMENTED;
+  readonly errorKind = ERROR_KIND_METHOD_NOT_IMPLEMENTED;
+  constructor(message: string) {
+    super(message);
+    this.name = "MethodNotImplementedError";
+  }
+}
+
+/** Raised when a sticky session token is malformed, expired, evicted, or
+ *  bound to a different worker / principal. HTTP-only. */
+export class SessionLostError extends Error {
+  static readonly errorKind = ERROR_KIND_SESSION_LOST;
+  readonly errorKind = ERROR_KIND_SESSION_LOST;
+  constructor(message: string) {
+    super(message);
+    this.name = "SessionLostError";
+  }
+}
+
+/** Raised when `ctx.openSession` is called while the server is draining. */
+export class ServerDrainingError extends Error {
+  static readonly errorKind = ERROR_KIND_SERVER_DRAINING;
+  readonly errorKind = ERROR_KIND_SERVER_DRAINING;
+  constructor(message: string) {
+    super(message);
+    this.name = "ServerDrainingError";
+  }
+}