@query-farm/vgi-rpc 0.6.3 → 0.7.0

package/src/server.ts

@@ -1,21 +1,44 @@
 // © Copyright 2025-2026, Query.Farm LLC - https://query.farm
 // SPDX-License-Identifier: Apache-2.0
 
-import { Schema } from "@query-farm/apache-arrow";
-import { DESCRIBE_METHOD_NAME } from "./constants.js";
+import { schema as makeSchema, serializeBatch } from "./arrow/index.js";
+import { DESCRIBE_METHOD_NAME, PROTOCOL_VERSION_KEY } from "./constants.js";
 import { buildDescribeBatch } from "./dispatch/describe.js";
 import { dispatchStream } from "./dispatch/stream.js";
 import { dispatchUnary } from "./dispatch/unary.js";
-import { RpcError, VersionError } from "./errors.js";
+import {
+  MethodNotImplementedError,
+  ProtocolVersionError,
+  parseProtocolVersion,
+  RpcError,
+  VersionError,
+} from "./errors.js";
 import type { ExternalLocationConfig } from "./external.js";
 import type { Protocol } from "./protocol.js";
-import { type CallStatistics, type DispatchHook, type DispatchInfo, MethodType } from "./types.js";
+import {
+  type CallStatistics,
+  type DispatchHook,
+  type DispatchInfo,
+  MethodType,
+  type ServeStartHook,
+  TransportKind,
+} from "./types.js";
 import { IpcStreamReader } from "./wire/reader.js";
-import { parseRequest } from "./wire/request.js";
+import { applyDefaults, parseRequest } from "./wire/request.js";
 import { buildErrorBatch } from "./wire/response.js";
 import { IpcStreamWriter } from "./wire/writer.js";
 
-const EMPTY_SCHEMA = new Schema([]);
+const EMPTY_SCHEMA = makeSchema([]);
+
+function randomStreamId(): string {
+  const bytes = new Uint8Array(16);
+  crypto.getRandomValues(bytes);
+  let out = "";
+  for (let i = 0; i < bytes.length; i++) {
+    out += bytes[i].toString(16).padStart(2, "0");
+  }
+  return out;
+}
 
 /**
  * RPC server that reads Arrow IPC requests from stdin and writes responses to stdout.
@@ -25,9 +48,22 @@ export class VgiRpcServer {
   private protocol: Protocol;
   private enableDescribe: boolean;
   private serverId: string;
-  private describeBatch: import("@query-farm/apache-arrow").RecordBatch | null = null;
+  // Lazily initialized — `buildDescribeBatch` is async because the protocol
+  // hash is computed via `crypto.subtle.digest` (Web Crypto). The dispatch
+  // path awaits the cached promise on first use.
+  private _describePromise: Promise<{
+    batch: import("./arrow/index.js").VgiBatch;
+    protocolHash: string;
+  }> | null = null;
+  private protocolVersion: string;
   private dispatchHook: DispatchHook | null = null;
   private externalConfig: ExternalLocationConfig | undefined;
+  private onServeStart: ServeStartHook | null = null;
+  /** True once the on_serve_start hook has fired successfully. The bind
+   *  state is committed only after the hook returns, so a transient
+   *  failure on first request leaves it `false` and the next request
+   *  re-fires rather than silently skipping. Mirrors Python 7b3999c. */
+  private serveStartFired = false;
 
   constructor(
     protocol: Protocol,
@@ -36,6 +72,9 @@ export class VgiRpcServer {
       serverId?: string;
       dispatchHook?: DispatchHook;
       externalLocation?: ExternalLocationConfig;
+      protocolVersion?: string;
+      /** Lifecycle hook fired once before the first dispatched request. */
+      onServeStart?: ServeStartHook;
     },
   ) {
     this.protocol = protocol;
@@ -43,11 +82,83 @@ export class VgiRpcServer {
     this.serverId = options?.serverId ?? crypto.randomUUID().replace(/-/g, "").slice(0, 12);
     this.dispatchHook = options?.dispatchHook ?? null;
     this.externalConfig = options?.externalLocation;
+    this.protocolVersion = options?.protocolVersion ?? "";
+    this.onServeStart = options?.onServeStart ?? null;
+  }
 
-    if (this.enableDescribe) {
-      const { batch } = buildDescribeBatch(protocol.name, protocol.getMethods(), this.serverId);
-      this.describeBatch = batch;
+  /** Fire the on_serve_start hook once for this transport. Idempotent
+   *  on success — re-throws on failure without committing the bind. */
+  private async notifyTransport(kind: TransportKind): Promise<void> {
+    if (this.serveStartFired) return;
+    if (this.onServeStart) {
+      await this.onServeStart(kind);
     }
+    this.serveStartFired = true;
+  }
+
+  /** Build (or retrieve cached) describe batch + protocol hash. */
+  private async describeInfo(): Promise<{
+    batch: import("./arrow/index.js").VgiBatch;
+    protocolHash: string;
+  }> {
+    if (!this._describePromise) {
+      this._describePromise = buildDescribeBatch(
+        this.protocol.name,
+        this.protocol.getMethods(),
+        this.serverId,
+        this.protocol.protocolVersion || undefined,
+      ).then(({ batch, metadata }) => ({
+        batch,
+        protocolHash: metadata.get("vgi_rpc.protocol_hash") ?? "",
+      }));
+    }
+    return this._describePromise;
+  }
+
+  /** Validate a client's declared protocol_version against the Protocol's
+   *  declared version. Caller invokes only when
+   *  `protocol.protocolVersionParts` is non-null. Mirrors Python's
+   *  `RpcServer._check_protocol_version`: exact major+minor match, patch
+   *  ignored; directional error message names which side is older. */
+  private checkProtocolVersion(clientVersion: string | undefined): void {
+    const serverParts = this.protocol.protocolVersionParts!;
+    const serverVersion = this.protocol.protocolVersion;
+    if (clientVersion === undefined) {
+      throw new ProtocolVersionError(
+        "VGI client/worker protocol_version mismatch.\n" +
+          "  Client: <not declared>\n" +
+          `  Server: ${serverVersion}\n` +
+          "  Direction: the client did not send a vgi_rpc.protocol_version " +
+          "metadata key. This is either a vgi-rpc framework bug or a " +
+          "non-VGI client connecting to a VGI worker.",
+      );
+    }
+    let clientParts: readonly [number, number, number];
+    try {
+      clientParts = parseProtocolVersion(clientVersion);
+    } catch {
+      throw new ProtocolVersionError(
+        "VGI client/worker protocol_version mismatch.\n" +
+          `  Client: ${clientVersion}\n` +
+          `  Server: ${serverVersion}\n` +
+          "  Direction: client sent a malformed protocol_version. " +
+          "Expected canonical semver MAJOR.MINOR.PATCH.",
+      );
+    }
+    if (clientParts[0] === serverParts[0] && clientParts[1] === serverParts[1]) {
+      return;
+    }
+    const clientOlder =
+      clientParts[0] < serverParts[0] || (clientParts[0] === serverParts[0] && clientParts[1] < serverParts[1]);
+    const direction = clientOlder
+      ? `client is too old; upgrade the VGI extension/client to a version supporting protocol_version ${serverVersion}.`
+      : `server is too old; upgrade the VGI worker to a version supporting protocol_version ${clientVersion}.`;
+    throw new ProtocolVersionError(
+      "VGI client/worker protocol_version mismatch.\n" +
+        `  Client: ${clientVersion}\n` +
+        `  Server: ${serverVersion}\n` +
+        `  Direction: ${direction}`,
+    );
   }
 
   /** Start the server loop. Reads requests until stdin closes. */
@@ -69,6 +180,11 @@ export class VgiRpcServer {
 
     try {
       while (true) {
+        // Fire on_serve_start lazily so the hook can do work that
+        // depends on the transport binding (matches Python which fires
+        // it inside serve()). Inside the loop so a failure on the very
+        // first request can be retried.
+        await this.notifyTransport(TransportKind.PIPE);
         await this.serveOne(reader, writer);
       }
     } catch (e: any) {
@@ -101,7 +217,7 @@ export class VgiRpcServer {
     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]);
+      await writer.writeStream(EMPTY_SCHEMA, [errBatch]);
       return;
     }
 
@@ -118,16 +234,17 @@ export class VgiRpcServer {
     } catch (e: any) {
       // Write error response for protocol/version errors
       const errBatch = buildErrorBatch(EMPTY_SCHEMA, e, this.serverId, null);
-      writer.writeStream(EMPTY_SCHEMA, [errBatch]);
+      await 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]);
+    // Handle __describe__ — lazy-build on first request.
+    if (methodName === DESCRIBE_METHOD_NAME && this.enableDescribe) {
+      const { batch } = await this.describeInfo();
+      await writer.writeStream(batch.schema, [batch]);
       return;
     }
 
@@ -136,15 +253,63 @@ export class VgiRpcServer {
     const method = methods.get(methodName);
     if (!method) {
       const available = [...methods.keys()].sort();
-      const err = new Error(`Unknown method: '${methodName}'. Available methods: [${available.join(", ")}]`);
+      const err = new MethodNotImplementedError(
+        `Unknown method: '${methodName}'. Available methods: [${available.join(", ")}]`,
+      );
       const errBatch = buildErrorBatch(EMPTY_SCHEMA, err, this.serverId, requestId);
-      writer.writeStream(EMPTY_SCHEMA, [errBatch]);
+      await writer.writeStream(EMPTY_SCHEMA, [errBatch]);
       return;
     }
 
+    // Application-protocol-version gate. Fires only when the Protocol
+    // declared a `protocolVersion`. `__describe__` is exempt — it is the
+    // diagnostic path a mismatched client uses to introspect the server's
+    // version. Mirrors Python's serve_one dispatch-boundary check.
+    if (this.protocol.protocolVersionParts !== null) {
+      try {
+        const md = batch.metadata;
+        this.checkProtocolVersion(md?.get(PROTOCOL_VERSION_KEY));
+      } catch (exc) {
+        const errSchema = method.type === MethodType.UNARY ? method.resultSchema : EMPTY_SCHEMA;
+        const errBatch = buildErrorBatch(errSchema, exc as Error, this.serverId, requestId);
+        await writer.writeStream(errSchema, [errBatch]);
+        return;
+      }
+    }
+
     // Dispatch based on method type, with optional hook
     const methodType = method.type === MethodType.UNARY ? "unary" : "stream";
-    const info: DispatchInfo = { method: methodName, methodType, serverId: this.serverId, requestId };
+
+    // Capture self-contained IPC bytes of the request batch for the access log.
+    let requestData: Uint8Array | undefined;
+    try {
+      requestData = serializeBatch(batch as any);
+    } catch {
+      // best-effort; observability must not fail dispatch
+    }
+
+    let streamId: string | undefined;
+    if (methodType === "stream") {
+      streamId = randomStreamId();
+    }
+
+    const { protocolHash } = await this.describeInfo();
+    const info: DispatchInfo = {
+      method: methodName,
+      methodType,
+      serverId: this.serverId,
+      requestId,
+      protocol: this.protocol.name,
+      protocolHash,
+      protocolVersion: this.protocolVersion,
+      kind: TransportKind.PIPE,
+      principal: "",
+      authDomain: "",
+      authenticated: false,
+      remoteAddr: "",
+      requestData,
+      streamId,
+    };
     const stats: CallStatistics = {
       inputBatches: 0,
       outputBatches: 0,
@@ -157,11 +322,22 @@ export class VgiRpcServer {
     const token = this.dispatchHook?.onDispatchStart(info);
     let dispatchError: Error | undefined;
 
+    applyDefaults(params, method.defaults);
+
     try {
       if (method.type === MethodType.UNARY) {
-        await dispatchUnary(method, params, writer, this.serverId, requestId, this.externalConfig);
+        await dispatchUnary(method, params, writer, this.serverId, requestId, this.externalConfig, TransportKind.PIPE);
       } else {
-        await dispatchStream(method, params, writer, reader, this.serverId, requestId, this.externalConfig);
+        await dispatchStream(
+          method,
+          params,
+          writer,
+          reader,
+          this.serverId,
+          requestId,
+          this.externalConfig,
+          TransportKind.PIPE,
+        );
       }
     } catch (e) {
       dispatchError = e instanceof Error ? e : new Error(String(e));