@query-farm/vgi-rpc 0.6.4 → 0.7.0

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. */
@@ -70,7 +92,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 +104,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 +117,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 +135,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 +153,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 +173,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 +195,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 +215,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/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,66 @@
 // © 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";
 
 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 +71,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 };
 }