@query-farm/vgi-rpc 0.1.0

package/src/http/token.ts

@@ -0,0 +1,149 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+import { createHmac, timingSafeEqual } from "node:crypto";
+
+const TOKEN_VERSION = 2;
+const HMAC_LEN = 32;
+// 1 (version) + 8 (created_at) + 4*3 (three length prefixes) + 32 (hmac)
+const MIN_TOKEN_LEN = 1 + 8 + 12 + HMAC_LEN;
+
+/**
+ * Pack a state token matching the Python v2 wire format.
+ *
+ * Layout:
+ *   [1B version=2]
+ *   [8B created_at uint64 LE (seconds since epoch)]
+ *   [4B state_len uint32 LE] [state_len bytes]
+ *   [4B schema_len uint32 LE] [schema_len bytes]
+ *   [4B input_schema_len uint32 LE] [input_schema_len bytes]
+ *   [32B HMAC-SHA256(signing_key, all above bytes)]
+ */
+export function packStateToken(
+  stateBytes: Uint8Array,
+  schemaBytes: Uint8Array,
+  inputSchemaBytes: Uint8Array,
+  signingKey: Uint8Array,
+  createdAt?: number,
+): string {
+  const now = createdAt ?? Math.floor(Date.now() / 1000);
+
+  const payloadLen =
+    1 + 8 + 4 + stateBytes.length + 4 + schemaBytes.length + 4 + inputSchemaBytes.length;
+  const buf = Buffer.alloc(payloadLen);
+  let offset = 0;
+
+  // version
+  buf.writeUInt8(TOKEN_VERSION, offset);
+  offset += 1;
+
+  // created_at as uint64 LE
+  buf.writeBigUInt64LE(BigInt(now), offset);
+  offset += 8;
+
+  // state
+  buf.writeUInt32LE(stateBytes.length, offset);
+  offset += 4;
+  buf.set(stateBytes, offset);
+  offset += stateBytes.length;
+
+  // output schema
+  buf.writeUInt32LE(schemaBytes.length, offset);
+  offset += 4;
+  buf.set(schemaBytes, offset);
+  offset += schemaBytes.length;
+
+  // input schema
+  buf.writeUInt32LE(inputSchemaBytes.length, offset);
+  offset += 4;
+  buf.set(inputSchemaBytes, offset);
+  offset += inputSchemaBytes.length;
+
+  // HMAC
+  const mac = createHmac("sha256", signingKey).update(buf).digest();
+  const token = Buffer.concat([buf, mac]);
+
+  return token.toString("base64");
+}
+
+export interface UnpackedToken {
+  stateBytes: Uint8Array;
+  schemaBytes: Uint8Array;
+  inputSchemaBytes: Uint8Array;
+  createdAt: number;
+}
+
+/**
+ * Unpack and verify a state token.
+ * Throws on tampered, expired, or malformed tokens.
+ */
+export function unpackStateToken(
+  tokenBase64: string,
+  signingKey: Uint8Array,
+  tokenTtl: number,
+): UnpackedToken {
+  const token = Buffer.from(tokenBase64, "base64");
+
+  if (token.length < MIN_TOKEN_LEN) {
+    throw new Error("State token too short");
+  }
+
+  // Split payload and mac
+  const payload = token.subarray(0, token.length - HMAC_LEN);
+  const receivedMac = token.subarray(token.length - HMAC_LEN);
+
+  // Verify HMAC first (before inspecting any fields)
+  const expectedMac = createHmac("sha256", signingKey).update(payload).digest();
+  if (!timingSafeEqual(receivedMac, expectedMac)) {
+    throw new Error("State token HMAC verification failed");
+  }
+
+  let offset = 0;
+
+  // Version
+  const version = payload.readUInt8(offset);
+  offset += 1;
+  if (version !== TOKEN_VERSION) {
+    throw new Error(`Unsupported state token version: ${version}`);
+  }
+
+  // created_at
+  const createdAt = Number(payload.readBigUInt64LE(offset));
+  offset += 8;
+
+  // TTL check
+  if (tokenTtl > 0) {
+    const now = Math.floor(Date.now() / 1000);
+    if (now - createdAt > tokenTtl) {
+      throw new Error("State token expired");
+    }
+  }
+
+  // state bytes
+  const stateLen = payload.readUInt32LE(offset);
+  offset += 4;
+  if (offset + stateLen > payload.length) {
+    throw new Error("State token truncated (state)");
+  }
+  const stateBytes = payload.slice(offset, offset + stateLen);
+  offset += stateLen;
+
+  // output schema bytes
+  const schemaLen = payload.readUInt32LE(offset);
+  offset += 4;
+  if (offset + schemaLen > payload.length) {
+    throw new Error("State token truncated (schema)");
+  }
+  const schemaBytes = payload.slice(offset, offset + schemaLen);
+  offset += schemaLen;
+
+  // input schema bytes
+  const inputSchemaLen = payload.readUInt32LE(offset);
+  offset += 4;
+  if (offset + inputSchemaLen > payload.length) {
+    throw new Error("State token truncated (input schema)");
+  }
+  const inputSchemaBytes = payload.slice(offset, offset + inputSchemaLen);
+
+  return { stateBytes, schemaBytes, inputSchemaBytes, createdAt };
+}

package/src/http/types.ts

@@ -0,0 +1,49 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+/** Configuration options for createHttpHandler(). */
+export interface HttpHandlerOptions {
+  /** URL path prefix for all endpoints. Default: "/vgi" */
+  prefix?: string;
+  /** HMAC-SHA256 signing key for state tokens. Random 32 bytes if omitted. */
+  signingKey?: Uint8Array;
+  /** State token time-to-live in seconds. Default: 3600 (1 hour). 0 disables TTL checks. */
+  tokenTtl?: number;
+  /** CORS allowed origins. If set, CORS headers are added to all responses. */
+  corsOrigins?: string;
+  /** Maximum request body size in bytes. Advertised via VGI-Max-Request-Bytes header. */
+  maxRequestBytes?: number;
+  /** Maximum bytes before a producer stream emits a continuation token. */
+  maxStreamResponseBytes?: number;
+  /** Server ID included in response metadata. Random if omitted. */
+  serverId?: string;
+  /** Custom state serializer for stream state objects. Default: JSON with BigInt support. */
+  stateSerializer?: StateSerializer;
+  /** zstd compression level for responses (1-22). If set, responses are
+   *  compressed when the client sends Accept-Encoding: zstd. */
+  compressionLevel?: number;
+}
+
+/** Serializer for stream state objects stored in state tokens. */
+export interface StateSerializer {
+  serialize(state: any): Uint8Array;
+  deserialize(bytes: Uint8Array): any;
+}
+
+/** Default state serializer using JSON (with BigInt support). */
+export const jsonStateSerializer: StateSerializer = {
+  serialize(state: any): Uint8Array {
+    return new TextEncoder().encode(
+      JSON.stringify(state, (_key, value) =>
+        typeof value === "bigint" ? `__bigint__:${value}` : value,
+      ),
+    );
+  },
+  deserialize(bytes: Uint8Array): any {
+    return JSON.parse(new TextDecoder().decode(bytes), (_key, value) =>
+      typeof value === "string" && value.startsWith("__bigint__:")
+        ? BigInt(value.slice(11))
+        : value,
+    );
+  },
+};

package/src/index.ts

@@ -0,0 +1,52 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+export { VgiRpcServer } from "./server.js";
+export { Protocol } from "./protocol.js";
+export {
+  MethodType,
+  OutputCollector,
+  type LogContext,
+  type MethodDefinition,
+  type UnaryHandler,
+  type HeaderInit,
+  type ProducerInit,
+  type ProducerFn,
+  type ExchangeInit,
+  type ExchangeFn,
+} from "./types.js";
+export {
+  type SchemaLike,
+  toSchema,
+  inferParamTypes,
+  str,
+  bytes,
+  int,
+  int32,
+  float,
+  float32,
+  bool,
+} from "./schema.js";
+export { RpcError, VersionError } from "./errors.js";
+export {
+  createHttpHandler,
+  ARROW_CONTENT_TYPE,
+  type HttpHandlerOptions,
+  type StateSerializer,
+} from "./http/index.js";
+export {
+  RPC_METHOD_KEY,
+  REQUEST_VERSION_KEY,
+  REQUEST_VERSION,
+  LOG_LEVEL_KEY,
+  LOG_MESSAGE_KEY,
+  LOG_EXTRA_KEY,
+  SERVER_ID_KEY,
+  REQUEST_ID_KEY,
+  PROTOCOL_NAME_KEY,
+  DESCRIBE_VERSION_KEY,
+  DESCRIBE_VERSION,
+  DESCRIBE_METHOD_NAME,
+  STATE_KEY,
+} from "./constants.js";
+export * from "./client/index.js";

package/src/protocol.ts

@@ -0,0 +1,144 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+import { Schema } from "apache-arrow";
+import { type SchemaLike, toSchema, inferParamTypes } from "./schema.js";
+import {
+  MethodType,
+  type MethodDefinition,
+  type UnaryHandler,
+  type HeaderInit,
+  type ProducerInit,
+  type ProducerFn,
+  type ExchangeInit,
+  type ExchangeFn,
+} from "./types.js";
+
+const EMPTY_SCHEMA = new Schema([]);
+
+/**
+ * Fluent builder for defining RPC methods.
+ * Register unary, producer, and exchange methods, then pass to `VgiRpcServer`.
+ */
+export class Protocol {
+  readonly name: string;
+  private _methods: Map<string, MethodDefinition> = new Map();
+
+  constructor(name: string) {
+    this.name = name;
+  }
+
+  /**
+   * Register a unary (request-response) method.
+   * @param name - Method name exposed to clients
+   * @param config.params - Parameter schema (SchemaLike)
+   * @param config.result - Result schema (SchemaLike)
+   * @param config.handler - Async function receiving params and returning result values
+   * @param config.doc - Optional documentation string
+   * @param config.defaults - Optional default parameter values
+   * @param config.paramTypes - Optional parameter type hints (inferred from params if omitted)
+   */
+  unary(
+    name: string,
+    config: {
+      params: SchemaLike;
+      result: SchemaLike;
+      handler: UnaryHandler;
+      doc?: string;
+      defaults?: Record<string, any>;
+      paramTypes?: Record<string, string>;
+    },
+  ): this {
+    const params = toSchema(config.params);
+    this._methods.set(name, {
+      name,
+      type: MethodType.UNARY,
+      paramsSchema: params,
+      resultSchema: toSchema(config.result),
+      handler: config.handler,
+      doc: config.doc,
+      defaults: config.defaults,
+      paramTypes: config.paramTypes ?? inferParamTypes(params),
+    });
+    return this;
+  }
+
+  /**
+   * Register a producer (server-streaming) method.
+   * The generic `S` is inferred from the `init` return type and threaded to `produce`.
+   */
+  producer<S>(
+    name: string,
+    config: {
+      params: SchemaLike;
+      outputSchema: SchemaLike;
+      init: ProducerInit<S>;
+      produce: ProducerFn<S>;
+      headerSchema?: SchemaLike;
+      headerInit?: HeaderInit;
+      doc?: string;
+      defaults?: Record<string, any>;
+      paramTypes?: Record<string, string>;
+    },
+  ): this {
+    const params = toSchema(config.params);
+    this._methods.set(name, {
+      name,
+      type: MethodType.STREAM,
+      paramsSchema: params,
+      resultSchema: EMPTY_SCHEMA,
+      outputSchema: toSchema(config.outputSchema),
+      inputSchema: EMPTY_SCHEMA,
+      producerInit: config.init as ProducerInit,
+      producerFn: config.produce as ProducerFn,
+      headerSchema: config.headerSchema ? toSchema(config.headerSchema) : undefined,
+      headerInit: config.headerInit,
+      doc: config.doc,
+      defaults: config.defaults,
+      paramTypes: config.paramTypes ?? inferParamTypes(params),
+    });
+    return this;
+  }
+
+  /**
+   * Register an exchange (bidirectional-streaming) method.
+   * The generic `S` is inferred from the `init` return type and threaded to `exchange`.
+   */
+  exchange<S>(
+    name: string,
+    config: {
+      params: SchemaLike;
+      inputSchema: SchemaLike;
+      outputSchema: SchemaLike;
+      init: ExchangeInit<S>;
+      exchange: ExchangeFn<S>;
+      headerSchema?: SchemaLike;
+      headerInit?: HeaderInit;
+      doc?: string;
+      defaults?: Record<string, any>;
+      paramTypes?: Record<string, string>;
+    },
+  ): this {
+    const params = toSchema(config.params);
+    this._methods.set(name, {
+      name,
+      type: MethodType.STREAM,
+      paramsSchema: params,
+      resultSchema: EMPTY_SCHEMA,
+      inputSchema: toSchema(config.inputSchema),
+      outputSchema: toSchema(config.outputSchema),
+      exchangeInit: config.init as ExchangeInit,
+      exchangeFn: config.exchange as ExchangeFn,
+      headerSchema: config.headerSchema ? toSchema(config.headerSchema) : undefined,
+      headerInit: config.headerInit,
+      doc: config.doc,
+      defaults: config.defaults,
+      paramTypes: config.paramTypes ?? inferParamTypes(params),
+    });
+    return this;
+  }
+
+  getMethods(): Map<string, MethodDefinition> {
+    return new Map(this._methods);
+  }
+}

package/src/schema.ts

@@ -0,0 +1,114 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+import {
+  Schema,
+  Field,
+  DataType,
+  Utf8,
+  Binary,
+  Int64,
+  Int32,
+  Int16,
+  Float64,
+  Float32,
+  Bool,
+} from "apache-arrow";
+
+// ---------------------------------------------------------------------------
+// Convenient DataType singletons — re-export so users avoid arrow imports
+// ---------------------------------------------------------------------------
+
+/** Apache Arrow Utf8 type. Use as schema shorthand: `{ name: str }` */
+export const str = new Utf8();
+/** Apache Arrow Binary type. Use as schema shorthand: `{ data: bytes }` */
+export const bytes = new Binary();
+/** Apache Arrow Int64 type. Use as schema shorthand: `{ count: int }` */
+export const int = new Int64();
+/** Apache Arrow Int32 type. Use as schema shorthand: `{ count: int32 }` */
+export const int32 = new Int32();
+/** Apache Arrow Float64 type. Use as schema shorthand: `{ value: float }` */
+export const float = new Float64();
+/** Apache Arrow Float32 type. Use as schema shorthand: `{ value: float32 }` */
+export const float32 = new Float32();
+/** Apache Arrow Bool type. Use as schema shorthand: `{ flag: bool }` */
+export const bool = new Bool();
+
+// ---------------------------------------------------------------------------
+// SchemaLike — shorthand for declaring schemas
+// ---------------------------------------------------------------------------
+
+/**
+ * A schema specification that accepts:
+ * - A real `Schema` (passed through)
+ * - A record mapping field names to `DataType` instances or `Field` instances
+ * - An empty `{}` for an empty schema
+ */
+export type SchemaLike = Schema | Record<string, DataType | Field>;
+
+/**
+ * Convert a SchemaLike spec into a real `Schema`.
+ *
+ * - `Schema` → returned as-is
+ * - `Record<string, DataType>` → each entry becomes `new Field(name, type, false)`
+ * - `Record<string, Field>` → each entry is passed through
+ * - `{}` → `new Schema([])`
+ */
+export function toSchema(spec: SchemaLike): Schema {
+  if (spec instanceof Schema) return spec;
+
+  const fields: Field[] = [];
+  for (const [name, value] of Object.entries(spec)) {
+    if (value instanceof Field) {
+      fields.push(value);
+    } else if (value instanceof DataType) {
+      fields.push(new Field(name, value, false));
+    } else {
+      throw new TypeError(
+        `Invalid schema value for "${name}": expected DataType or Field, got ${typeof value}`,
+      );
+    }
+  }
+  return new Schema(fields);
+}
+
+// ---------------------------------------------------------------------------
+// inferParamTypes — derive paramTypes from a schema spec
+// ---------------------------------------------------------------------------
+
+const TYPE_MAP: [new (...args: any[]) => DataType, string][] = [
+  [Utf8, "str"],
+  [Binary, "bytes"],
+  [Bool, "bool"],
+  [Float64, "float"],
+  [Float32, "float"],
+  [Int64, "int"],
+  [Int32, "int"],
+  [Int16, "int"],
+];
+
+/**
+ * Derive a `paramTypes` record from a SchemaLike spec.
+ * Maps common Arrow scalar types to Python-style type strings.
+ * Returns `undefined` if any field has a complex type (List, Map_, Dictionary, etc.).
+ */
+export function inferParamTypes(
+  spec: SchemaLike,
+): Record<string, string> | undefined {
+  const schema = toSchema(spec);
+  if (schema.fields.length === 0) return undefined;
+
+  const result: Record<string, string> = {};
+  for (const field of schema.fields) {
+    let mapped: string | undefined;
+    for (const [ctor, name] of TYPE_MAP) {
+      if (field.type instanceof ctor) {
+        mapped = name;
+        break;
+      }
+    }
+    if (!mapped) return undefined;
+    result[field.name] = mapped;
+  }
+  return result;
+}

package/src/server.ts

@@ -0,0 +1,159 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+import { Schema } from "apache-arrow";
+import { Protocol } from "./protocol.js";
+import { IpcStreamReader } from "./wire/reader.js";
+import { IpcStreamWriter } from "./wire/writer.js";
+import { parseRequest } from "./wire/request.js";
+import { buildErrorBatch } from "./wire/response.js";
+import { buildDescribeBatch } from "./dispatch/describe.js";
+import { dispatchUnary } from "./dispatch/unary.js";
+import { dispatchStream } from "./dispatch/stream.js";
+import { DESCRIBE_METHOD_NAME } from "./constants.js";
+import { MethodType } from "./types.js";
+import { RpcError, VersionError } from "./errors.js";
+
+const EMPTY_SCHEMA = new Schema([]);
+
+/**
+ * RPC server that reads Arrow IPC requests from stdin and writes responses to stdout.
+ * Supports unary and streaming (producer/exchange) methods.
+ */
+export class VgiRpcServer {
+  private protocol: Protocol;
+  private enableDescribe: boolean;
+  private serverId: string;
+  private describeBatch: import("apache-arrow").RecordBatch | null = null;
+
+  constructor(
+    protocol: Protocol,
+    options?: { enableDescribe?: boolean; serverId?: string },
+  ) {
+    this.protocol = protocol;
+    this.enableDescribe = options?.enableDescribe ?? true;
+    this.serverId =
+      options?.serverId ?? crypto.randomUUID().replace(/-/g, "").slice(0, 12);
+
+    if (this.enableDescribe) {
+      const { batch } = buildDescribeBatch(
+        protocol.name,
+        protocol.getMethods(),
+        this.serverId,
+      );
+      this.describeBatch = batch;
+    }
+  }
+
+  /** Start the server loop. Reads requests until stdin closes. */
+  async run(): Promise<void> {
+    const stdin = process.stdin as unknown as ReadableStream<Uint8Array>;
+
+    // Warn if running interactively
+    if (process.stdin.isTTY || process.stdout.isTTY) {
+      process.stderr.write(
+        "WARNING: This process communicates via Arrow IPC on stdin/stdout " +
+          "and is not intended to be run interactively.\n" +
+          "It should be launched as a subprocess by an RPC client " +
+          "(e.g. vgi_rpc.connect()).\n",
+      );
+    }
+
+    const reader = await IpcStreamReader.create(stdin);
+    const writer = new IpcStreamWriter();
+
+    try {
+      while (true) {
+        await this.serveOne(reader, writer);
+      }
+    } catch (e: any) {
+      // EOF or broken pipe → clean exit
+      if (
+        e.message?.includes("closed") ||
+        e.message?.includes("Expected Schema Message") ||
+        e.message?.includes("null or length 0") ||
+        e.code === "EPIPE" ||
+        e.code === "ERR_STREAM_PREMATURE_CLOSE" ||
+        e.code === "ERR_STREAM_DESTROYED" ||
+        (e instanceof Error && e.message.includes("EOF"))
+      ) {
+        return;
+      }
+      // ArrowInvalid or unexpected error
+      throw e;
+    } finally {
+      await reader.cancel();
+    }
+  }
+
+  private async serveOne(
+    reader: IpcStreamReader,
+    writer: IpcStreamWriter,
+  ): Promise<void> {
+    const stream = await reader.readStream();
+    if (!stream) {
+      throw new Error("EOF");
+    }
+
+    const { schema, batches } = stream;
+    if (batches.length === 0) {
+      const err = new RpcError("ProtocolError", "Request stream contains no batches", "");
+      const errBatch = buildErrorBatch(EMPTY_SCHEMA, err, this.serverId, null);
+      writer.writeStream(EMPTY_SCHEMA, [errBatch]);
+      return;
+    }
+
+    const batch = batches[0];
+    let methodName: string;
+    let params: Record<string, any>;
+    let requestId: string | null;
+
+    try {
+      const parsed = parseRequest(schema, batch);
+      methodName = parsed.methodName;
+      params = parsed.params;
+      requestId = parsed.requestId;
+    } catch (e: any) {
+      // Write error response for protocol/version errors
+      const errBatch = buildErrorBatch(EMPTY_SCHEMA, e, this.serverId, null);
+      writer.writeStream(EMPTY_SCHEMA, [errBatch]);
+      if (e instanceof VersionError || e instanceof RpcError) {
+        return; // Continue serving
+      }
+      throw e;
+    }
+
+    // Handle __describe__
+    if (methodName === DESCRIBE_METHOD_NAME && this.describeBatch) {
+      writer.writeStream(this.describeBatch.schema, [this.describeBatch]);
+      return;
+    }
+
+    // Look up method
+    const methods = this.protocol.getMethods();
+    const method = methods.get(methodName);
+    if (!method) {
+      const available = [...methods.keys()].sort();
+      const err = new Error(
+        `Unknown method: '${methodName}'. Available methods: [${available.join(", ")}]`,
+      );
+      const errBatch = buildErrorBatch(EMPTY_SCHEMA, err, this.serverId, requestId);
+      writer.writeStream(EMPTY_SCHEMA, [errBatch]);
+      return;
+    }
+
+    // Dispatch based on method type
+    if (method.type === MethodType.UNARY) {
+      await dispatchUnary(method, params, writer, this.serverId, requestId);
+    } else {
+      await dispatchStream(
+        method,
+        params,
+        writer,
+        reader,
+        this.serverId,
+        requestId,
+      );
+    }
+  }
+}