@query-farm/vgi-rpc 0.6.4 → 0.7.1

package/src/errors.ts

@@ -4,8 +4,11 @@
 /** Error thrown when the server encounters an RPC protocol error. */
 export class RpcError extends Error {
   constructor(
+    /** Remote error class name (e.g. `"ValueError"`). */
     public readonly errorType: string,
+    /** Human-readable message from the remote error. */
     public readonly errorMessage: string,
+    /** Remote stack-trace text, or an empty string when unavailable. */
     public readonly remoteTraceback: string,
   ) {
     super(`${errorType}: ${errorMessage}`);
@@ -20,3 +23,87 @@ export class VersionError extends Error {
     this.name = "VersionError";
   }
 }
+
+/** `vgi_rpc.error_kind` batch-metadata value for {@link MethodNotImplementedError}.
+ *  Mirrors Python's `vgi_rpc.metadata.ERROR_KIND_*` constants. */
+export const ERROR_KIND_METHOD_NOT_IMPLEMENTED = "method_not_implemented";
+/** `vgi_rpc.error_kind` batch-metadata value for {@link SessionLostError}. */
+export const ERROR_KIND_SESSION_LOST = "session_lost";
+/** `vgi_rpc.error_kind` batch-metadata value for {@link ServerDrainingError}. */
+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 {
+  /** Typed `vgi_rpc.error_kind` marker for this error class. */
+  static readonly errorKind = ERROR_KIND_METHOD_NOT_IMPLEMENTED;
+  /** Typed `vgi_rpc.error_kind` marker hoisted onto the error batch metadata. */
+  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 {
+  /** Typed `vgi_rpc.error_kind` marker for this error class. */
+  static readonly errorKind = ERROR_KIND_SESSION_LOST;
+  /** Typed `vgi_rpc.error_kind` marker hoisted onto the error batch metadata. */
+  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 {
+  /** Typed `vgi_rpc.error_kind` marker for this error class. */
+  static readonly errorKind = ERROR_KIND_SERVER_DRAINING;
+  /** Typed `vgi_rpc.error_kind` marker hoisted onto the error batch metadata. */
+  readonly errorKind = ERROR_KIND_SERVER_DRAINING;
+  constructor(message: string) {
+    super(message);
+    this.name = "ServerDrainingError";
+  }
+}

package/src/external.ts

@@ -10,7 +10,7 @@
  * download URL and SHA-256 checksum in metadata.
  */
 
-import { type RecordBatch, RecordBatchReader, RecordBatchStreamWriter, type Schema } from "@query-farm/apache-arrow";
+import { deserializeBatch, serializeBatch, type VgiBatch, type VgiSchema } from "./arrow/index.js";
 import { LOCATION_KEY, LOCATION_SHA256_KEY, LOG_LEVEL_KEY } from "./constants.js";
 import { zstdCompress, zstdDecompress } from "./util/zstd.js";
 import { buildEmptyBatch } from "./wire/response.js";
@@ -25,6 +25,28 @@ export interface ExternalStorage {
   upload(data: Uint8Array, contentEncoding: string): Promise<string>;
 }
 
+/** A pre-signed PUT/GET URL pair for client-side data upload. */
+export interface UploadUrl {
+  /** Pre-signed PUT URL the client uploads to. */
+  uploadUrl: string;
+  /** Pre-signed GET URL the server fetches from. */
+  downloadUrl: string;
+  /** Expiration time (UTC) for the URL pair. */
+  expiresAt: Date;
+}
+
+/**
+ * Generates pre-signed upload URL pairs for client-vended externalization.
+ *
+ * Implementations must be safe to call from multiple concurrent requests.
+ * Object lifecycle is the operator's responsibility — uploaded objects are
+ * not automatically deleted by vgi-rpc.
+ */
+export interface UploadUrlProvider {
+  /** Allocate one upload/download URL pair. */
+  generateUploadUrl(): Promise<UploadUrl> | UploadUrl;
+}
+
 /** Configuration for external storage of large batches. */
 export interface ExternalLocationConfig {
   /** Storage backend for uploading. */
@@ -32,7 +54,12 @@ export interface ExternalLocationConfig {
   /** Minimum batch byte size to trigger externalization. Default: 1MB. */
   externalizeThresholdBytes?: number;
   /** Optional zstd compression for uploaded data. */
-  compression?: { algorithm: "zstd"; level?: number };
+  compression?: {
+    /** Compression algorithm; only `"zstd"` is currently supported. */
+    algorithm: "zstd";
+    /** zstd compression level. Default: 3. */
+    level?: number;
+  };
   /** URL validator called before fetching. Throw to reject. Default: HTTPS-only. */
   urlValidator?: ((url: string) => void) | null;
 }
@@ -70,7 +97,7 @@ async function sha256Hex(data: Uint8Array): Promise<string> {
 // ---------------------------------------------------------------------------
 
 /** Returns true if the batch is a zero-row pointer to external data. */
-export function isExternalLocationBatch(batch: RecordBatch): boolean {
+export function isExternalLocationBatch(batch: VgiBatch): boolean {
   if (batch.numRows !== 0) return false;
   const meta = batch.metadata;
   if (!meta) return false;
@@ -82,7 +109,7 @@ export function isExternalLocationBatch(batch: RecordBatch): boolean {
 // ---------------------------------------------------------------------------
 
 /** Create a zero-row pointer batch with location URL and optional SHA-256. */
-export function makeExternalLocationBatch(schema: Schema, url: string, sha256?: string): RecordBatch {
+export function makeExternalLocationBatch(schema: VgiSchema, url: string, sha256?: string): VgiBatch {
   const metadata = new Map<string, string>();
   metadata.set(LOCATION_KEY, url);
   if (sha256) {
@@ -95,22 +122,13 @@ export function makeExternalLocationBatch(schema: Schema, url: string, sha256?:
 // IPC serialization helpers
 // ---------------------------------------------------------------------------
 
-function serializeBatchToIpc(batch: RecordBatch): Uint8Array {
-  const writer = new RecordBatchStreamWriter();
-  writer.reset(undefined, batch.schema);
-  writer.write(batch);
-  writer.close();
-  return writer.toUint8Array(true);
+function serializeBatchToIpc(batch: VgiBatch): Uint8Array {
+  return serializeBatch(batch);
 }
 
-function batchByteSize(batch: RecordBatch): number {
-  // Arrow TS data.byteLength doesn't reflect actual data size.
+function batchByteSize(batch: VgiBatch): number {
   // Estimate from IPC serialization size for threshold check.
-  const writer = new RecordBatchStreamWriter();
-  writer.reset(undefined, batch.schema);
-  writer.write(batch);
-  writer.close();
-  return writer.toUint8Array(true).byteLength;
+  return serializeBatch(batch).byteLength;
 }
 
 // ---------------------------------------------------------------------------
@@ -122,9 +140,9 @@ function batchByteSize(batch: RecordBatch): number {
  * Returns the original batch unchanged if below threshold or no config.
  */
 export async function maybeExternalizeBatch(
-  batch: RecordBatch,
+  batch: VgiBatch,
   config?: ExternalLocationConfig | null,
-): Promise<RecordBatch> {
+): Promise<VgiBatch> {
   if (!config?.storage) return batch;
   if (batch.numRows === 0) return batch;
 
@@ -140,7 +158,7 @@ export async function maybeExternalizeBatch(
   // Optionally compress
   let contentEncoding = "";
   if (config.compression?.algorithm === "zstd") {
-    ipcData = zstdCompress(ipcData, config.compression.level ?? 3) as Uint8Array;
+    ipcData = (await zstdCompress(ipcData, config.compression.level ?? 3)) as Uint8Array;
     contentEncoding = "zstd";
   }
 
@@ -160,9 +178,9 @@ export async function maybeExternalizeBatch(
  * Returns the original batch unchanged if not a pointer or no config.
  */
 export async function resolveExternalLocation(
-  batch: RecordBatch,
+  batch: VgiBatch,
   config?: ExternalLocationConfig | null,
-): Promise<RecordBatch> {
+): Promise<VgiBatch> {
   if (!config) return batch;
   if (!isExternalLocationBatch(batch)) return batch;
 
@@ -182,10 +200,14 @@ export async function resolveExternalLocation(
   }
   let data = new Uint8Array(await response.arrayBuffer());
 
-  // Decompress if needed
+  // Decompress if needed.  Cap the decompressed size at 16x the
+  // compressed body — generous for typical Arrow IPC zstd ratios but
+  // tight enough that a tiny response cannot inflate to multi-GB.
+  // Mirrors Python's external_fetch.fetch_url.
   const contentEncoding = response.headers.get("Content-Encoding");
   if (contentEncoding === "zstd") {
-    data = new Uint8Array(zstdDecompress(data));
+    const cap = data.byteLength * 16;
+    data = new Uint8Array(await zstdDecompress(data, cap));
   }
 
   // Verify SHA-256 if present
@@ -198,12 +220,9 @@ export async function resolveExternalLocation(
   }
 
   // Parse IPC stream
-  const reader = await RecordBatchReader.from(data);
-  await reader.open();
-  const resolved = reader.next();
-  if (!resolved || resolved.done || !resolved.value) {
+  const resolved = deserializeBatch(data);
+  if (resolved.numRows === 0 && resolved.schema.fields.length === 0) {
     throw new Error(`No data batch found in external IPC stream from ${url}`);
   }
-
-  return resolved.value;
+  return resolved;
 }

package/src/http/auth.ts

@@ -8,14 +8,27 @@ export type AuthenticateFn = (request: Request) => AuthContext | Promise<AuthCon
 
 /** RFC 9728 OAuth Protected Resource Metadata. */
 export interface OAuthResourceMetadata {
+  /** The protected resource's canonical URL. Doubles as the base for the
+   *  `/_oauth/callback` redirect URI. */
   resource: string;
+  /** Authorization-server issuer URLs. The PKCE flow uses
+   *  `authorizationServers[0]` for OIDC discovery. */
   authorizationServers: string[];
+  /** Scopes the resource advertises. When non-empty these become the PKCE
+   *  authorization request's space-joined `scope`, taking precedence over
+   *  {@link HttpHandlerOptions.oauthPkceScope}. */
   scopesSupported?: string[];
+  /** Advertised bearer methods (e.g. `["header"]`). */
   bearerMethodsSupported?: string[];
+  /** JWS algorithms the resource accepts. */
   resourceSigningAlgValuesSupported?: string[];
+  /** Human-readable resource name. */
   resourceName?: string;
+  /** Documentation URL for the resource. */
   resourceDocumentation?: string;
+  /** Policy URL for the resource. */
   resourcePolicyUri?: string;
+  /** Terms-of-service URL for the resource. */
   resourceTosUri?: string;
   /** OAuth client_id that clients should use with the authorization server. */
   clientId?: string;

package/src/http/bearer.ts

@@ -1,8 +1,8 @@
 // © Copyright 2025-2026, Query.Farm LLC - https://query.farm
 // SPDX-License-Identifier: Apache-2.0
 
-import { timingSafeEqual } from "node:crypto";
 import type { AuthContext } from "../auth.js";
+import { constantTimeEqual } from "../util/web-crypto.js";
 import type { AuthenticateFn } from "./auth.js";
 
 /** Receives the raw bearer token string, returns an AuthContext on success. Must throw on failure. */
@@ -30,10 +30,7 @@ export function bearerAuthenticate(options: { validate: BearerValidateFn }): Aut
 /** Constant-time string comparison to prevent timing attacks on token lookup. */
 function safeEqual(a: string, b: string): boolean {
   const enc = new TextEncoder();
-  const bufA = enc.encode(a);
-  const bufB = enc.encode(b);
-  if (bufA.byteLength !== bufB.byteLength) return false;
-  return timingSafeEqual(bufA, bufB);
+  return constantTimeEqual(enc.encode(a), enc.encode(b));
 }
 
 /**

package/src/http/common.ts

@@ -1,11 +1,67 @@
 // © Copyright 2025-2026, Query.Farm LLC - https://query.farm
 // SPDX-License-Identifier: Apache-2.0
 
-import { type RecordBatch, RecordBatchReader, RecordBatchStreamWriter, type Schema } from "@query-farm/apache-arrow";
-import { conformBatchToSchema } from "../util/conform.js";
+import {
+  conformBatchToSchema,
+  deserializeBatch,
+  serializeBatches,
+  type VgiBatch,
+  type VgiSchema,
+} from "../arrow/index.js";
+import { RPC_ERROR_HEADER } from "../constants.js";
+import type { CookieSpec } from "../types.js";
 
+/** MIME type for Arrow IPC stream request and response bodies. */
 export const ARROW_CONTENT_TYPE = "application/vnd.apache.arrow.stream";
 
+// Sticky session header conventions (HTTP-only). Mirrors Python's
+// `vgi_rpc.http._common`. Headers — not cookies — so multiple concurrent
+// sessions to one host from a single client multiplex correctly.
+export const SESSION_HEADER = "VGI-Session";
+export const SESSION_ACCEPT_HEADER = "VGI-Session-Accept";
+export const SESSION_CLOSE_HEADER = "VGI-Session-Close";
+export const STICKY_ENABLED_HEADER = "VGI-Sticky-Enabled";
+export const STICKY_DEFAULT_TTL_HEADER = "VGI-Sticky-Default-TTL";
+export const STICKY_ECHO_HEADERS_HEADER = "VGI-Sticky-Echo-Headers";
+
+/** Prefix the server uses to tell the client "echo this header on subsequent
+ *  requests in this session". Clients capture and replay
+ *  `VGI-Echo-<name>: <value>` as plain `<name>: <value>` for the session
+ *  lifetime — used for client-driven routing (e.g. `fly-force-instance-id`). */
+export const ECHO_HEADER_PREFIX = "VGI-Echo-";
+
+/** Framework-managed sticky session teardown endpoint path component.
+ *  `DELETE {prefix}/__session__` idempotently closes the session referenced
+ *  by the request's `VGI-Session` header. */
+export const SESSION_ENDPOINT = "__session__";
+
+/** Serialize a CookieSpec into a Set-Cookie header value. */
+export function formatSetCookieHeader(c: CookieSpec): string {
+  const parts: string[] = [];
+  if (c.delete) {
+    parts.push(`${c.name}=`);
+    parts.push("Max-Age=0");
+  } else {
+    parts.push(`${c.name}=${c.value}`);
+    if (c.maxAge !== undefined) parts.push(`Max-Age=${c.maxAge}`);
+    if (c.expires) parts.push(`Expires=${c.expires.toUTCString()}`);
+  }
+  if (c.path) parts.push(`Path=${c.path}`);
+  if (c.domain) parts.push(`Domain=${c.domain}`);
+  if (c.secure) parts.push("Secure");
+  if (c.httpOnly) parts.push("HttpOnly");
+  if (c.sameSite) parts.push(`SameSite=${c.sameSite}`);
+  if (c.partitioned) parts.push("Partitioned");
+  return parts.join("; ");
+}
+
+/** Append Set-Cookie headers for each queued CookieSpec onto an existing Headers object. */
+export function appendCookieHeaders(headers: Headers, cookies: readonly CookieSpec[]): void {
+  for (const c of cookies) {
+    headers.append("Set-Cookie", formatSetCookieHeader(c));
+  }
+}
+
 export class HttpRpcError extends Error {
   constructor(
     message: string,
@@ -16,35 +72,47 @@ export class HttpRpcError extends Error {
   }
 }
 
-/** Serialize a schema + batches into a complete IPC stream as Uint8Array. */
-export function serializeIpcStream(schema: Schema, batches: RecordBatch[]): Uint8Array {
-  const writer = new RecordBatchStreamWriter();
-  writer.reset(undefined, schema);
-  for (const batch of batches) {
-    writer.write(conformBatchToSchema(batch, schema));
-  }
-  writer.close();
-  return writer.toUint8Array(true);
+/**
+ * Serialize a schema + batches into a complete IPC stream as Uint8Array.
+ *
+ * A single IPC stream is `[schema_msg, batch_msg, batch_msg, ..., EOS]`.
+ * Each backend implements `serializeBatches` to write that atomically —
+ * arrow-js via `RecordBatchStreamWriter`, flechette via `tablesToIPC`
+ * (added in our flechette fork). Naive concatenation of per-batch streams
+ * produces multiple EOS markers and breaks readers.
+ */
+export function serializeIpcStream(schema: VgiSchema, batches: VgiBatch[]): Uint8Array {
+  const conformed = batches.map((b) => conformBatchToSchema(b, schema));
+  return serializeBatches(schema, conformed);
 }
 
-/** Create a Response with Arrow IPC content type. Casts Uint8Array for TS lib compat. */
+/**
+ * Create a Response with Arrow IPC content type.
+ *
+ * Server errors (status 500) are translated to HTTP 200 with an
+ * ``X-VGI-RPC-Error: true`` header so that clients which discard
+ * response bodies on 5xx still receive the Arrow IPC error metadata.
+ * Client errors (400, 401, 404, 415) are passed through unchanged.
+ */
 export function arrowResponse(body: Uint8Array, status = 200, extraHeaders?: Headers): Response {
   const headers = extraHeaders ?? new Headers();
   headers.set("Content-Type", ARROW_CONTENT_TYPE);
+  if (status === 500) {
+    headers.set(RPC_ERROR_HEADER, "true");
+    return new Response(body as unknown as BodyInit, { status: 200, headers });
+  }
   return new Response(body as unknown as BodyInit, { status, headers });
 }
 
-/** Read schema + first batch from an IPC stream body. */
-export async function readRequestFromBody(body: Uint8Array): Promise<{ schema: Schema; batch: RecordBatch }> {
-  const reader = await RecordBatchReader.from(body);
-  await reader.open();
-  const schema = reader.schema;
-  if (!schema) {
+/** Read schema + first batch from an IPC stream body via the facade. */
+export async function readRequestFromBody(body: Uint8Array): Promise<{ schema: VgiSchema; batch: VgiBatch }> {
+  const batch = deserializeBatch(body);
+  // Reject only truly empty bodies. A zero-field, zero-row batch with batch
+  // metadata is a legal exchange/cancel/continuation signal — the state
+  // token rides on `batch.metadata` and downstream code (cancel detection,
+  // schema conformance gating) is built to handle it.
+  if (batch.schema.fields.length === 0 && batch.numRows === 0 && (batch.metadata?.size ?? 0) === 0) {
     throw new HttpRpcError("Empty IPC stream: no schema", 400);
   }
-  const batches = reader.readAll();
-  if (batches.length === 0) {
-    throw new HttpRpcError("IPC stream contains no batches", 400);
-  }
-  return { schema, batch: batches[0] };
+  return { schema: batch.schema, batch };
 }