@query-farm/vgi-rpc 0.1.0

package/src/client/stream.ts

@@ -0,0 +1,297 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+import {
+  RecordBatch,
+  Schema,
+  Field,
+  makeData,
+  Struct,
+  vectorFromArray,
+} from "apache-arrow";
+import { STATE_KEY } from "../constants.js";
+import { RpcError } from "../errors.js";
+import { ARROW_CONTENT_TYPE, serializeIpcStream } from "../http/common.js";
+import {
+  inferArrowType,
+  dispatchLogOrError,
+  extractBatchRows,
+  readResponseBatches,
+} from "./ipc.js";
+import type { LogMessage, StreamSession } from "./types.js";
+
+type CompressFn = (data: Uint8Array, level: number) => Uint8Array;
+type DecompressFn = (data: Uint8Array) => Uint8Array;
+
+export class HttpStreamSession implements StreamSession {
+  private _baseUrl: string;
+  private _prefix: string;
+  private _method: string;
+  private _stateToken: string | null;
+  private _outputSchema: Schema;
+  private _inputSchema?: Schema;
+  private _onLog?: (msg: LogMessage) => void;
+  private _pendingBatches: RecordBatch[];
+  private _finished: boolean;
+  private _header: Record<string, any> | null;
+  private _compressionLevel?: number;
+  private _compressFn?: CompressFn;
+  private _decompressFn?: DecompressFn;
+
+  constructor(opts: {
+    baseUrl: string;
+    prefix: string;
+    method: string;
+    stateToken: string | null;
+    outputSchema: Schema;
+    inputSchema?: Schema;
+    onLog?: (msg: LogMessage) => void;
+    pendingBatches: RecordBatch[];
+    finished: boolean;
+    header: Record<string, any> | null;
+    compressionLevel?: number;
+    compressFn?: CompressFn;
+    decompressFn?: DecompressFn;
+  }) {
+    this._baseUrl = opts.baseUrl;
+    this._prefix = opts.prefix;
+    this._method = opts.method;
+    this._stateToken = opts.stateToken;
+    this._outputSchema = opts.outputSchema;
+    this._inputSchema = opts.inputSchema;
+    this._onLog = opts.onLog;
+    this._pendingBatches = opts.pendingBatches;
+    this._finished = opts.finished;
+    this._header = opts.header;
+    this._compressionLevel = opts.compressionLevel;
+    this._compressFn = opts.compressFn;
+    this._decompressFn = opts.decompressFn;
+  }
+
+  get header(): Record<string, any> | null {
+    return this._header;
+  }
+
+  private _buildHeaders(): Record<string, string> {
+    const headers: Record<string, string> = {
+      "Content-Type": ARROW_CONTENT_TYPE,
+    };
+    if (this._compressionLevel != null) {
+      headers["Content-Encoding"] = "zstd";
+      headers["Accept-Encoding"] = "zstd";
+    }
+    return headers;
+  }
+
+  private _prepareBody(content: Uint8Array): Uint8Array {
+    if (this._compressionLevel != null && this._compressFn) {
+      return this._compressFn(content, this._compressionLevel);
+    }
+    return content;
+  }
+
+  private async _readResponse(resp: Response): Promise<Uint8Array<ArrayBuffer>> {
+    let body = new Uint8Array(await resp.arrayBuffer());
+    if (resp.headers.get("Content-Encoding") === "zstd" && this._decompressFn) {
+      body = new Uint8Array(this._decompressFn(body));
+    }
+    return body;
+  }
+
+  /**
+   * Send an exchange request and return the data rows.
+   */
+  async exchange(input: Record<string, any>[]): Promise<Record<string, any>[]> {
+    if (this._stateToken === null) {
+      throw new RpcError(
+        "ProtocolError",
+        "Stream has finished \u2014 no state token available",
+        "",
+      );
+    }
+
+    // We need to determine the input schema from the data.
+    // Build a batch from the input rows using the output schema's field types.
+    // For exchange, the input schema matches what the server expects.
+    // We'll use the keys from input[0] to figure out columns.
+    if (input.length === 0) {
+      // Zero-row exchange: build an empty batch with state token.
+      // Use inputSchema from __describe__ if available; fall back to
+      // outputSchema so the server sees the correct column names.
+      const zeroSchema = this._inputSchema ?? this._outputSchema;
+      const emptyBatch = this._buildEmptyBatch(zeroSchema);
+      const metadata = new Map<string, string>();
+      metadata.set(STATE_KEY, this._stateToken);
+      const batchWithMeta = new RecordBatch(
+        zeroSchema,
+        emptyBatch.data,
+        metadata,
+      );
+      return this._doExchange(zeroSchema, [batchWithMeta]);
+    }
+
+    // Infer schema from first row values (input schema may differ from output).
+    const keys = Object.keys(input[0]);
+    const fields = keys.map((key) => {
+      // Find first non-null value to infer type
+      let sample: any = undefined;
+      for (const row of input) {
+        if (row[key] != null) { sample = row[key]; break; }
+      }
+      const arrowType = inferArrowType(sample);
+      const nullable = input.some((row) => row[key] == null);
+      return new Field(key, arrowType, nullable);
+    });
+
+    const inputSchema = new Schema(fields);
+    const children = inputSchema.fields.map((f) => {
+      const values = input.map((row) => row[f.name]);
+      return vectorFromArray(values, f.type).data[0];
+    });
+
+    const structType = new Struct(inputSchema.fields);
+    const data = makeData({
+      type: structType,
+      length: input.length,
+      children,
+      nullCount: 0,
+    });
+
+    const metadata = new Map<string, string>();
+    metadata.set(STATE_KEY, this._stateToken);
+    const batch = new RecordBatch(inputSchema, data, metadata);
+
+    return this._doExchange(inputSchema, [batch]);
+  }
+
+  private async _doExchange(
+    schema: Schema,
+    batches: RecordBatch[],
+  ): Promise<Record<string, any>[]> {
+    const body = serializeIpcStream(schema, batches);
+    const resp = await fetch(
+      `${this._baseUrl}${this._prefix}/${this._method}/exchange`,
+      {
+        method: "POST",
+        headers: this._buildHeaders(),
+        body: this._prepareBody(body) as unknown as BodyInit,
+      },
+    );
+
+    const responseBody = await this._readResponse(resp);
+    const { batches: responseBatches } = await readResponseBatches(responseBody);
+
+    let resultRows: Record<string, any>[] = [];
+    for (const batch of responseBatches) {
+      if (batch.numRows === 0) {
+        // Could be log/error or state token
+        dispatchLogOrError(batch, this._onLog);
+        // Check for state token
+        const token = batch.metadata?.get(STATE_KEY);
+        if (token) {
+          this._stateToken = token;
+        }
+        continue;
+      }
+
+      // Data batch — extract state token from metadata
+      const token = batch.metadata?.get(STATE_KEY);
+      if (token) {
+        this._stateToken = token;
+      }
+
+      resultRows = extractBatchRows(batch);
+    }
+
+    return resultRows;
+  }
+
+  private _buildEmptyBatch(schema: Schema): RecordBatch {
+    const children = schema.fields.map((f) => {
+      return makeData({ type: f.type, length: 0, nullCount: 0 });
+    });
+    const structType = new Struct(schema.fields);
+    const data = makeData({
+      type: structType,
+      length: 0,
+      children,
+      nullCount: 0,
+    });
+    return new RecordBatch(schema, data);
+  }
+
+  /**
+   * Iterate over producer stream batches.
+   */
+  async *[Symbol.asyncIterator](): AsyncIterableIterator<Record<string, any>[]> {
+    // Yield pre-loaded batches from init
+    for (const batch of this._pendingBatches) {
+      if (batch.numRows === 0) {
+        dispatchLogOrError(batch, this._onLog);
+        continue;
+      }
+      yield extractBatchRows(batch);
+    }
+    this._pendingBatches = [];
+
+    if (this._finished) return;
+    if (this._stateToken === null) return;
+
+    // Follow continuation tokens
+    while (true) {
+      const responseBody = await this._sendContinuation(this._stateToken);
+      const { batches } = await readResponseBatches(responseBody);
+
+      let gotContinuation = false;
+      for (const batch of batches) {
+        if (batch.numRows === 0) {
+          // Check for continuation token
+          const token = batch.metadata?.get(STATE_KEY);
+          if (token) {
+            this._stateToken = token;
+            gotContinuation = true;
+            continue;
+          }
+          // Log/error batch
+          dispatchLogOrError(batch, this._onLog);
+          continue;
+        }
+
+        yield extractBatchRows(batch);
+      }
+
+      if (!gotContinuation) break;
+    }
+  }
+
+  private async _sendContinuation(token: string): Promise<Uint8Array> {
+    const emptySchema = new Schema([]);
+    const metadata = new Map<string, string>();
+    metadata.set(STATE_KEY, token);
+
+    const structType = new Struct(emptySchema.fields);
+    const data = makeData({
+      type: structType,
+      length: 1,
+      children: [],
+      nullCount: 0,
+    });
+    const batch = new RecordBatch(emptySchema, data, metadata);
+    const body = serializeIpcStream(emptySchema, [batch]);
+
+    const resp = await fetch(
+      `${this._baseUrl}${this._prefix}/${this._method}/exchange`,
+      {
+        method: "POST",
+        headers: this._buildHeaders(),
+        body: this._prepareBody(body) as unknown as BodyInit,
+      },
+    );
+
+    return this._readResponse(resp);
+  }
+
+  close(): void {
+    // No-op for HTTP (stateless)
+  }
+}

package/src/client/types.ts

@@ -0,0 +1,31 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+export interface HttpConnectOptions {
+  prefix?: string;
+  onLog?: (msg: LogMessage) => void;
+  compressionLevel?: number;
+}
+
+export interface LogMessage {
+  level: string;
+  message: string;
+  extra?: Record<string, any>;
+}
+
+export interface StreamSession {
+  readonly header: Record<string, any> | null;
+  exchange(input: Record<string, any>[]): Promise<Record<string, any>[]>;
+  [Symbol.asyncIterator](): AsyncIterableIterator<Record<string, any>[]>;
+  close(): void;
+}
+
+export interface PipeConnectOptions {
+  onLog?: (msg: LogMessage) => void;
+}
+
+export interface SubprocessConnectOptions extends PipeConnectOptions {
+  cwd?: string;
+  env?: Record<string, string>;
+  stderr?: "inherit" | "pipe" | "ignore";
+}

package/src/constants.ts

@@ -0,0 +1,22 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+/** Well-known metadata keys matching Python's metadata.py */
+
+export const RPC_METHOD_KEY = "vgi_rpc.method";
+export const LOG_LEVEL_KEY = "vgi_rpc.log_level";
+export const LOG_MESSAGE_KEY = "vgi_rpc.log_message";
+export const LOG_EXTRA_KEY = "vgi_rpc.log_extra";
+export const REQUEST_VERSION_KEY = "vgi_rpc.request_version";
+export const REQUEST_VERSION = "1";
+
+export const SERVER_ID_KEY = "vgi_rpc.server_id";
+export const REQUEST_ID_KEY = "vgi_rpc.request_id";
+
+export const PROTOCOL_NAME_KEY = "vgi_rpc.protocol_name";
+export const DESCRIBE_VERSION_KEY = "vgi_rpc.describe_version";
+export const DESCRIBE_VERSION = "2";
+
+export const DESCRIBE_METHOD_NAME = "__describe__";
+
+export const STATE_KEY = "vgi_rpc.stream_state#b64";

package/src/dispatch/describe.ts

@@ -0,0 +1,155 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+import {
+  Schema,
+  Field,
+  RecordBatch,
+  Utf8,
+  Bool,
+  Binary,
+  vectorFromArray,
+  makeData,
+  Struct,
+} from "apache-arrow";
+import type { MethodDefinition } from "../types.js";
+import {
+  PROTOCOL_NAME_KEY,
+  REQUEST_VERSION_KEY,
+  REQUEST_VERSION,
+  DESCRIBE_VERSION_KEY,
+  DESCRIBE_VERSION,
+  SERVER_ID_KEY,
+} from "../constants.js";
+import { serializeSchema } from "../util/schema.js";
+
+/**
+ * The schema for the __describe__ response, matching Python's _DESCRIBE_SCHEMA.
+ */
+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),
+]);
+
+/**
+ * Build the __describe__ response batch and metadata.
+ */
+export function buildDescribeBatch(
+  protocolName: string,
+  methods: Map<string, MethodDefinition>,
+  serverId: string,
+): { batch: RecordBatch; 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)[] = [];
+
+  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,
+    );
+  }
+
+  // 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 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],
+  ];
+
+  const structType = new Struct(DESCRIBE_SCHEMA.fields);
+  const data = makeData({
+    type: structType,
+    length: sortedEntries.length,
+    children,
+    nullCount: 0,
+  });
+
+  // Build metadata for the batch
+  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(SERVER_ID_KEY, serverId);
+
+  const batch = new RecordBatch(DESCRIBE_SCHEMA, data, metadata);
+
+  return { batch, metadata };
+}

package/src/dispatch/stream.ts

@@ -0,0 +1,151 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+import { Schema } from "apache-arrow";
+import type { MethodDefinition } from "../types.js";
+import { OutputCollector } from "../types.js";
+import type { IpcStreamWriter } from "../wire/writer.js";
+import type { IpcStreamReader } from "../wire/reader.js";
+import { buildResultBatch, buildErrorBatch } from "../wire/response.js";
+
+const EMPTY_SCHEMA = new Schema([]);
+
+/**
+ * Dispatch a stream RPC call (producer or exchange).
+ *
+ * Producer streams (empty input schema):
+ * - Client sends tick batches (empty schema, 0 rows)
+ * - Server reads each tick, calls produce(state, out)
+ * - Server writes output batch(es) for each tick
+ * - When produce() calls out.finish(), server closes output stream
+ *
+ * Exchange streams (real input schema):
+ * - Client sends data batches
+ * - Server reads each batch, calls exchange(state, input, out)
+ * - Server writes output batch(es) for each input
+ * - Stream ends when client closes input (EOS)
+ */
+export async function dispatchStream(
+  method: MethodDefinition,
+  params: Record<string, any>,
+  writer: IpcStreamWriter,
+  reader: IpcStreamReader,
+  serverId: string,
+  requestId: string | null,
+): Promise<void> {
+  const isProducer =
+    !method.inputSchema || method.inputSchema.fields.length === 0;
+
+  let state: any;
+  try {
+    if (isProducer) {
+      state = await method.producerInit!(params);
+    } else {
+      state = await method.exchangeInit!(params);
+    }
+  } catch (error: any) {
+    const errSchema = method.headerSchema ?? EMPTY_SCHEMA;
+    const errBatch = buildErrorBatch(errSchema, error, serverId, requestId);
+    writer.writeStream(errSchema, [errBatch]);
+    // Still need to consume the input stream from the client
+    const inputSchema = await reader.openNextStream();
+    if (inputSchema) {
+      while ((await reader.readNextBatch()) !== null) {
+        // drain
+      }
+    }
+    return;
+  }
+
+  // Support dynamic output schemas: init may return state with __outputSchema
+  // to override the method's registered output schema (needed for methods
+  // like VGI's "init" that produce different schemas per invocation).
+  const outputSchema = state?.__outputSchema ?? method.outputSchema!;
+
+  // Effective producer mode: check state override (VGI "init" is registered as
+  // exchange but may act as producer depending on the function type).
+  const effectiveProducer = state?.__isProducer ?? isProducer;
+
+  // 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 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);
+    } catch (error: any) {
+      const errBatch = buildErrorBatch(method.headerSchema, error, serverId, requestId);
+      writer.writeStream(method.headerSchema, [errBatch]);
+      // Drain input stream so client doesn't hang
+      const inputSchema = await reader.openNextStream();
+      if (inputSchema) {
+        while ((await reader.readNextBatch()) !== null) {}
+      }
+      return;
+    }
+  }
+
+  // Open the input IPC stream (ticks or data from client)
+  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]);
+    return;
+  }
+
+  // Use a single continuous IPC stream for all output (matching Python vgi-rpc).
+  // DuckDB exchanges are ping-pong: one input batch → one output batch on the
+  // same stream. We use IncrementalStream which writes bytes synchronously.
+  const stream = writer.openStream(outputSchema);
+
+  try {
+    while (true) {
+      const inputBatch = await reader.readNextBatch();
+      if (!inputBatch) break;
+
+      const out = new OutputCollector(outputSchema, effectiveProducer, serverId, requestId);
+
+      if (isProducer) {
+        await method.producerFn!(state, out);
+      } else {
+        await method.exchangeFn!(state, inputBatch, out);
+      }
+
+      for (const emitted of out.batches) {
+        stream.write(emitted.batch);
+      }
+
+      if (out.finished) {
+        break;
+      }
+    }
+  } catch (error: any) {
+    stream.write(buildErrorBatch(outputSchema, error, serverId, requestId));
+  }
+
+  stream.close();
+
+  // Drain remaining input so transport stays synchronized for next request.
+  // Matches Python's _drain_stream() called after every streaming method.
+  // Needed when the loop exits early (out.finished, error) while client
+  // is still sending batches.
+  try {
+    while ((await reader.readNextBatch()) !== null) {}
+  } catch {
+    // Suppress errors during drain (broken pipe, etc.)
+  }
+}

package/src/dispatch/unary.ts

@@ -0,0 +1,35 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+import { Schema } from "apache-arrow";
+import type { MethodDefinition } from "../types.js";
+import { OutputCollector } from "../types.js";
+import type { IpcStreamWriter } from "../wire/writer.js";
+import { buildResultBatch, buildErrorBatch } from "../wire/response.js";
+
+/**
+ * Dispatch a unary RPC call.
+ * Calls the handler with parsed params, writes result or error batch.
+ * Supports client-directed logging via ctx.clientLog().
+ */
+export async function dispatchUnary(
+  method: MethodDefinition,
+  params: Record<string, any>,
+  writer: IpcStreamWriter,
+  serverId: string,
+  requestId: string | null,
+): Promise<void> {
+  const schema = method.resultSchema;
+  const out = new OutputCollector(schema, true, serverId, requestId);
+
+  try {
+    const result = await method.handler!(params, out);
+    const resultBatch = buildResultBatch(schema, result, serverId, requestId);
+    // Collect log batches (from clientLog) + result batch
+    const batches = [...out.batches.map((b) => b.batch), resultBatch];
+    writer.writeStream(schema, batches);
+  } catch (error: any) {
+    const batch = buildErrorBatch(schema, error, serverId, requestId);
+    writer.writeStream(schema, [batch]);
+  }
+}