rivetkit 2.3.6 → 2.3.8

package/src/registry/index.ts

@@ -1,4 +1,5 @@
 import { isLocalEngineEndpoint } from "@/common/engine";
+import { isResponseLike } from "@/common/fetch-like";
 import { configureBaseLogger } from "@/common/log";
 import { configureServerlessPool } from "@/serverless/configure";
 import { VERSION } from "@/utils";
@@ -14,8 +15,16 @@ import {
 	RegistryConfigSchema,
 } from "./config";
 import { logger } from "./log";
-import { buildConfiguredRegistry } from "./native";
-import type { RuntimeServerlessResponseHead } from "./runtime";
+import {
+	buildConfiguredRegistry,
+} from "./native";
+import { convertNativeHttpResponse } from "./native-http";
+import type {
+	CoreRuntime,
+	RuntimeApplicationFetch,
+	RuntimeHttpResponseBodyStream,
+	RuntimeServerlessResponseHead,
+} from "./runtime";
 
 type ShutdownSignal = "SIGINT" | "SIGTERM";
 
@@ -37,6 +46,66 @@ function finishShutdownSignal(signal: ShutdownSignal): void {
 	process.kill(process.pid, signal);
 }
 
+function createApplicationFetch(
+	application: RegistryApplication,
+	runtime: CoreRuntime,
+): RuntimeApplicationFetch {
+	return async (
+		request,
+		responseBodyStream?: RuntimeHttpResponseBodyStream,
+	) => {
+		const method = request.method.toUpperCase();
+		const body =
+			method === "GET" || method === "HEAD" || request.body.length === 0
+				? undefined
+				: Uint8Array.from(request.body).buffer;
+		const abortController = new AbortController();
+		if (request.cancelToken) {
+			if (runtime.cancellationTokenAborted(request.cancelToken)) {
+				abortController.abort();
+			} else {
+				runtime.onCancellationTokenCancelled(
+					request.cancelToken,
+					() => abortController.abort(),
+				);
+			}
+		}
+		const response = await application.fetch(
+			new Request(request.url, {
+				method,
+				headers: request.headers,
+				body,
+				signal: abortController.signal,
+			}),
+		);
+		if (!isResponseLike(response)) {
+			throw new TypeError(
+				"registry.listen() application fetch must return a Response",
+			);
+		}
+		const conversion = await convertNativeHttpResponse(
+			response,
+			responseBodyStream,
+		);
+		if (conversion.bodyCompletion) {
+			void conversion.bodyCompletion.catch((error) => {
+				logger().error({
+					msg: "application response stream failed",
+					error,
+				});
+			});
+		}
+		return {
+			status: conversion.response.status ?? response.status,
+			headers:
+				conversion.response.headers ??
+				Object.fromEntries(response.headers.entries()),
+			body: conversion.response.body,
+			stream: conversion.response.stream,
+		};
+	};
+}
+
 export type FetchHandler = (
 	request: Request,
 	...args: any
@@ -46,6 +115,17 @@ export interface ServerlessHandler {
 	fetch: FetchHandler;
 }
 
+export interface RegistryApplication {
+	fetch: FetchHandler;
+}
+
+export interface RegistryListenOptions {
+	port?: number;
+	host?: string;
+	publicDir?: string;
+	application?: RegistryApplication;
+}
+
 export interface RegistryRoutes {
 	health(): Promise<Response>;
 	metadata(): Promise<Response>;
@@ -79,6 +159,7 @@ export class Registry<A extends RegistryActors> {
 	#runtimeReadyPromise?: Promise<void>;
 	#runtimeServeConfiguredPromise?: ReturnType<typeof buildConfiguredRegistry>;
 	#runtimeServerlessPromise?: ReturnType<typeof buildConfiguredRegistry>;
+	#applicationListenerPromise?: Promise<void>;
 	#configureServerlessPoolPromise?: Promise<void>;
 	#welcomePrinted = false;
 	#shutdownInstalled = false;
@@ -348,22 +429,72 @@ export class Registry<A extends RegistryActors> {
 	 * @param opts.host      Address to bind. Defaults to `0.0.0.0`.
 	 * @param opts.publicDir If set, serves static files from this directory
 	 *                       as a fallback below the framework routes.
+	 * @param opts.application If set, handles requests that do not match a
+	 *                         framework route.
 	 *
 	 * @example
 	 * ```ts
 	 * await registry.listen();
-	 * await registry.listen({ port: 8080, publicDir: "./public" });
+	 * await registry.listen({ application: app });
 	 * ```
 	 */
-	public async listen(
-		opts: { port?: number; host?: string; publicDir?: string } = {},
-	): Promise<void> {
+	public async listen(opts: RegistryListenOptions = {}): Promise<void> {
 		const port = opts.port ?? parsePortEnv(process.env.RIVET_PORT) ?? 3000;
 		const publicDir = opts.publicDir ?? getRivetkitPublicDir();
 		const config = this.parseConfig();
 
+		if (opts.application && getRivetkitRuntimeMode() !== "serverless") {
+			if (config.runtime === "wasm") {
+				throw new Error(
+					"registry.listen() requires the native runtime; use an application-owned HTTP server with WebAssembly",
+				);
+			}
+			this.#installSignalHandlers(config);
+			this.#printWelcome(config, "serverful", {
+				port,
+				host: opts.host,
+				publicDir,
+			});
+			this.#startEnvoy(config, true);
+			const readyPromise = this.startAndWait();
+			const configuredRegistryPromise =
+				this.#runtimeServeConfiguredPromise;
+			if (!configuredRegistryPromise) {
+				throw new Error("registry envoy startup did not initialize");
+			}
+			const { runtime, registry, serveConfig } =
+				await configuredRegistryPromise;
+			const application = createApplicationFetch(
+				opts.application,
+				runtime,
+			);
+			await readyPromise;
+			const listenerPromise = runtime.serveApplicationListener(
+				registry,
+				{
+					port,
+					host: opts.host,
+					publicDir,
+					application,
+				},
+				serveConfig,
+			);
+			this.#applicationListenerPromise = listenerPromise;
+			const serveLifecyclePromise = this.#runtimeServeLifecyclePromise;
+			if (!serveLifecyclePromise) {
+				throw new Error(
+					"registry envoy serve lifecycle did not initialize",
+				);
+			}
+			await Promise.all([
+				listenerPromise,
+				serveLifecyclePromise,
+			]);
+			return;
+		}
+
 		// Cache on both promise fields so the shutdown drain sees Mode A and B.
-		const configuredRegistryPromise = buildConfiguredRegistry(config);
+		const configuredRegistryPromise = this.#buildConfiguredRegistry(config);
 		this.#runtimeServeConfiguredPromise = configuredRegistryPromise;
 		this.#runtimeServerlessPromise = configuredRegistryPromise;
 		this.#installSignalHandlers(config);
@@ -379,9 +510,17 @@ export class Registry<A extends RegistryActors> {
 
 		const { runtime, registry, serveConfig } =
 			await configuredRegistryPromise;
+		const application = opts.application
+			? createApplicationFetch(opts.application, runtime)
+			: undefined;
 		await runtime.serveListener(
 			registry,
-			{ port, host: opts.host, publicDir },
+			{
+				port,
+				host: opts.host,
+				publicDir,
+				application,
+			},
 			serveConfig,
 		);
 	}
@@ -499,15 +638,14 @@ export class Registry<A extends RegistryActors> {
 					await runtime.serveRegistry(registry, serveConfig);
 				},
 			);
-			this.#runtimeServePromise = this.#runtimeServeLifecyclePromise.catch(
-				(error) => {
+			this.#runtimeServePromise =
+				this.#runtimeServeLifecyclePromise.catch((error) => {
 					// Always-attached catch so the stored promise never leaves a
 					// rejection unhandled. Downstream awaits (e.g. #runShutdown's
 					// Promise.race) attach their own catches and still observe
 					// resolution via the race.
 					logger().warn({ error }, "runtime registry serve errored");
-				},
-			);
+				});
 			// Install signal handlers once an envoy lifecycle has begun. Only
 			// Mode A ever reaches here. Mode B (handler(request)) intentionally
 			// does not install handlers because it runs on Workers/Vercel/Deno
@@ -653,6 +791,9 @@ export class Registry<A extends RegistryActors> {
 				// already logged any serve-side error.
 				await runtimeServePromise.catch(() => undefined);
 			}
+			if (this.#applicationListenerPromise !== undefined) {
+				await this.#applicationListenerPromise.catch(() => undefined);
+			}
 		};
 		await Promise.race([
 			drain(),
@@ -712,7 +853,9 @@ export class Registry<A extends RegistryActors> {
 	public startAndWait(): Promise<void> {
 		if (this.#shutdownInFlight !== null) {
 			return Promise.reject(
-				new Error("registry.startAndWait() cannot run after shutdown has begun"),
+				new Error(
+					"registry.startAndWait() cannot run after shutdown has begun",
+				),
 			);
 		}
 		if (this.#runtimeReadyPromise) return this.#runtimeReadyPromise;
@@ -755,7 +898,9 @@ export class Registry<A extends RegistryActors> {
 		const readinessPromise = (async () => {
 			const { runtime, registry } = await configuredRegistryPromise;
 			const stoppedBeforeReady = serveLifecyclePromise.then(() => {
-				throw new Error("RivetKit registry stopped before becoming ready");
+				throw new Error(
+					"RivetKit registry stopped before becoming ready",
+				);
 			});
 			await Promise.race([
 				runtime.waitRegistryReady(registry),

package/src/registry/napi-runtime.ts

@@ -4,6 +4,7 @@ import type {
 	CancellationToken as NativeCancellationToken,
 	ConnHandle as NativeConnHandle,
 	CoreRegistry as NativeCoreRegistry,
+	HttpResponseBodyStream as NativeHttpResponseBodyStream,
 	WebSocket as NativeWebSocket,
 } from "@rivetkit/rivetkit-napi";
 import type {
@@ -14,10 +15,13 @@ import type {
 	CoreRuntime,
 	RegistryHandle,
 	RuntimeActorConfig,
+	RuntimeApplicationFetch,
+	RuntimeApplicationListenerConfig,
 	RuntimeBytes,
 	RuntimeCronFire,
 	RuntimeCronJobInfo,
 	RuntimeHttpRequest,
+	RuntimeHttpResponseBodyStream,
 	RuntimeKvEntry,
 	RuntimeKvListOptions,
 	RuntimeListenerConfig,
@@ -135,6 +139,56 @@ function toNapiBuffer(value: RuntimeBytes): Buffer {
 	return Buffer.from(value);
 }
 
+function toNapiApplication(application: RuntimeApplicationFetch) {
+	return async (
+		error: unknown,
+		request?: {
+			method: string;
+			url: string;
+			headers: Record<string, string>;
+			body: Buffer;
+			cancelToken?: NativeCancellationToken;
+			responseBodyStream?: RuntimeHttpResponseBodyStream;
+		},
+	) => {
+		if (error) throw error;
+		if (!request) {
+			throw new Error("application fetch callback received no request");
+		}
+		const response = await application(
+			{
+				...request,
+				body: new Uint8Array(request.body),
+				cancelToken:
+					request.cancelToken as unknown as CancellationTokenHandle,
+			},
+			request.responseBodyStream
+				? fromNapiHttpResponseBodyStream(
+						request.responseBodyStream as unknown as NativeHttpResponseBodyStream,
+					)
+				: undefined,
+		);
+		return {
+			...response,
+			body:
+				response.body === undefined
+					? undefined
+					: toNapiBuffer(response.body),
+		};
+	};
+}
+
+function fromNapiHttpResponseBodyStream(
+	stream: NativeHttpResponseBodyStream,
+): RuntimeHttpResponseBodyStream {
+	return {
+		cancelled: () => stream.cancelled(),
+		write: (chunk) => stream.write(Buffer.from(chunk)),
+		end: () => stream.end(),
+		error: (message) => stream.error(message),
+	};
+}
+
 function toNapiHttpRequest(
 	request?: RuntimeHttpRequest | undefined | null,
 ): Parameters<NativeActorContext["connectConn"]>[1] {
@@ -296,6 +350,9 @@ export class NapiCoreRuntime implements CoreRuntime {
 		listener: RuntimeListenerConfig,
 		config: RuntimeServeConfig,
 	): Promise<void> {
+		const application = listener.application
+			? toNapiApplication(listener.application)
+			: undefined;
 		await asNativeRegistry(registry).serveListener(
 			{
 				port: listener.port,
@@ -303,6 +360,23 @@ export class NapiCoreRuntime implements CoreRuntime {
 				publicDir: listener.publicDir,
 			},
 			config,
+			application,
+		);
+	}
+
+	async serveApplicationListener(
+		registry: RegistryHandle,
+		listener: RuntimeApplicationListenerConfig,
+		config: RuntimeServeConfig,
+	): Promise<void> {
+		await asNativeRegistry(registry).serveApplicationListener(
+			{
+				port: listener.port,
+				host: listener.host,
+				publicDir: listener.publicDir,
+			},
+			toNapiApplication(listener.application),
+			config,
 		);
 	}
 

package/src/registry/native-http.ts

@@ -0,0 +1,208 @@
+import type { ResponseLike } from "@/common/fetch-like";
+import { HEADER_ORIGINAL_REQUEST_URL } from "@/common/actor-router-consts";
+import { stringifyError } from "@/common/utils";
+import { logger } from "./log";
+import type { RuntimeBytes, RuntimeHttpResponse } from "./runtime";
+
+const HTTP_BODY_CHUNK_SIZE = 64 * 1024;
+
+// Structural types keep the HTTP adapter independent of the N-API package.
+export interface NativeHttpResponseBodyStream {
+	cancelled(): Promise<void>;
+	write(chunk: Uint8Array): Promise<void>;
+	end(): Promise<void>;
+	error(message: string): Promise<void>;
+}
+
+export interface NativeHttpRequestBodyStream {
+	read(): Promise<Uint8Array | null | undefined>;
+	cancel(): Promise<void>;
+}
+
+interface NativeHttpRequestInit {
+	method: string;
+	uri: string;
+	headers?: Record<string, string>;
+	body?: RuntimeBytes;
+	bodyStream?: NativeHttpRequestBodyStream;
+	signal?: AbortSignal;
+	abortController?: AbortController;
+}
+
+interface NativeHttpResponseConversion {
+	response: RuntimeHttpResponse;
+	// Keeps request-scoped actor resources alive until the response body closes.
+	bodyCompletion?: Promise<void>;
+}
+
+function runtimeBytesToArrayBuffer(value: RuntimeBytes): ArrayBuffer {
+	return value.buffer.slice(
+		value.byteOffset,
+		value.byteOffset + value.byteLength,
+	) as ArrayBuffer;
+}
+
+export function buildNativeHttpRequest(init: NativeHttpRequestInit): Request {
+	const headers = new Headers(init.headers);
+	const originalUrl = headers.get(HEADER_ORIGINAL_REQUEST_URL);
+	headers.delete(HEADER_ORIGINAL_REQUEST_URL);
+	const url =
+		originalUrl ??
+		(init.uri.startsWith("http")
+			? init.uri
+			: new URL(init.uri, "http://127.0.0.1").toString());
+	const method = init.method.toUpperCase();
+	const bodyForbidden = method === "GET" || method === "HEAD";
+	const body = bodyForbidden
+		? undefined
+		: init.bodyStream
+			? new ReadableStream<Uint8Array>({
+					async pull(controller) {
+						try {
+							if (init.body && init.body.length > 0) {
+								controller.enqueue(new Uint8Array(init.body));
+								init.body = undefined;
+								return;
+							}
+							const chunk = await init.bodyStream?.read();
+							if (!chunk) {
+								controller.close();
+							} else {
+								controller.enqueue(new Uint8Array(chunk));
+							}
+						} catch (error) {
+							init.abortController?.abort(error);
+							controller.error(error);
+						}
+					},
+					async cancel() {
+						await init.bodyStream?.cancel();
+					},
+				})
+			: init.body && init.body.length > 0
+				? runtimeBytesToArrayBuffer(init.body)
+				: undefined;
+	const streamInit =
+		init.bodyStream && !bodyForbidden ? { duplex: "half" } : {};
+	return new Request(url, {
+		method,
+		headers,
+		body,
+		signal: init.abortController?.signal ?? init.signal,
+		...streamInit,
+	} as RequestInit);
+}
+
+export async function cancelNativeHttpRequestBody(
+	bodyStream?: NativeHttpRequestBodyStream,
+) {
+	await bodyStream?.cancel();
+}
+
+async function writeResponseChunk(
+	stream: NativeHttpResponseBodyStream,
+	chunk: Uint8Array,
+) {
+	for (
+		let offset = 0;
+		offset < chunk.byteLength;
+		offset += HTTP_BODY_CHUNK_SIZE
+	) {
+		await stream.write(
+			chunk.subarray(offset, offset + HTTP_BODY_CHUNK_SIZE),
+		);
+	}
+}
+
+async function pumpResponseBody(
+	reader: ReadableStreamDefaultReader<Uint8Array>,
+	stream: NativeHttpResponseBodyStream,
+) {
+	// Awaiting each write applies native-side backpressure. Racing the next read
+	// against cancellation also stops an otherwise unbounded SSE response when
+	// its downstream receiver disappears.
+	const cancelled = stream
+		.cancelled()
+		.then(() => ({ cancelled: true }) as const);
+	try {
+		for (;;) {
+			const read = reader
+				.read()
+				.then((next) => ({ cancelled: false, next }) as const);
+			const result = await Promise.race([read, cancelled]);
+			if (result.cancelled) {
+				await reader.cancel(
+					new Error("native http response stream receiver dropped"),
+				);
+				return;
+			}
+			const { next } = result;
+			if (next.done) {
+				await stream.end();
+				return;
+			}
+			if (next.value?.byteLength) {
+				await writeResponseChunk(stream, next.value);
+			}
+		}
+	} catch (error) {
+		try {
+			await stream.error(stringifyError(error));
+		} catch (streamError) {
+			logger().debug({
+				msg: "failed to report native http response stream error",
+				error: streamError,
+			});
+		}
+		try {
+			await reader.cancel(error);
+		} catch {
+			// Reader may already be closed after a native-side disconnect.
+		}
+	} finally {
+		reader.releaseLock();
+	}
+}
+
+export async function convertNativeHttpResponse(
+	response: ResponseLike,
+	responseBodyStream?: NativeHttpResponseBodyStream,
+): Promise<NativeHttpResponseConversion> {
+	const headers = Object.fromEntries(response.headers.entries());
+	if (!response.body) {
+		return {
+			response: {
+				status: response.status,
+				headers,
+				body: new Uint8Array(),
+			},
+		};
+	}
+
+	if (responseBodyStream) {
+		const reader = response.body.getReader();
+		const bodyCompletion = pumpResponseBody(reader, responseBodyStream);
+		return {
+			response: {
+				status: response.status,
+				headers,
+				stream: true,
+			},
+			bodyCompletion,
+		};
+	}
+
+	return {
+		response: {
+			status: response.status,
+			headers,
+			body: new Uint8Array(await response.arrayBuffer()),
+		},
+	};
+}
+
+export const nativeHttpTestInternals = {
+	buildRequest: buildNativeHttpRequest,
+	cancelRequestBody: cancelNativeHttpRequestBody,
+	convertRuntimeHttpResponse: convertNativeHttpResponse,
+};