rivetkit 2.3.12-rc.4 → 2.3.12

package/src/registry/index.ts

@@ -16,8 +16,11 @@ import {
 } from "./config";
 import { logger } from "./log";
 import { buildConfiguredRegistry } from "./native";
+import { convertNativeHttpResponse } from "./native-http";
 import type {
+	CoreRuntime,
 	RuntimeApplicationFetch,
+	RuntimeHttpResponseBodyStream,
 	RuntimeServerlessResponseHead,
 } from "./runtime";
 
@@ -43,18 +46,33 @@ function finishShutdownSignal(signal: ShutdownSignal): void {
 
 function createApplicationFetch(
 	application: RegistryApplication,
+	runtime: CoreRuntime,
 ): RuntimeApplicationFetch {
-	return async (request) => {
+	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)) {
@@ -62,10 +80,25 @@ function createApplicationFetch(
 				"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: response.status,
-			headers: Object.fromEntries(response.headers.entries()),
-			body: new Uint8Array(await response.arrayBuffer()),
+			status: conversion.response.status ?? response.status,
+			headers:
+				conversion.response.headers ??
+				Object.fromEntries(response.headers.entries()),
+			body: conversion.response.body,
+			stream: conversion.response.stream,
 		};
 	};
 }
@@ -406,11 +439,8 @@ export class Registry<A extends RegistryActors> {
 		const port = opts.port ?? parsePortEnv(process.env.RIVET_PORT) ?? 3000;
 		const publicDir = opts.publicDir ?? getRivetkitPublicDir();
 		const config = this.parseConfig();
-		const application = opts.application
-			? createApplicationFetch(opts.application)
-			: undefined;
 
-		if (application && getRivetkitRuntimeMode() !== "serverless") {
+		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",
@@ -431,6 +461,11 @@ export class Registry<A extends RegistryActors> {
 			}
 			const { runtime, registry, serveConfig } =
 				await configuredRegistryPromise;
+			const application = createApplicationFetch(
+				opts.application,
+				runtime,
+			);
+			await readyPromise;
 			const listenerPromise = runtime.serveApplicationListener(
 				registry,
 				{
@@ -448,16 +483,12 @@ export class Registry<A extends RegistryActors> {
 					"registry envoy serve lifecycle did not initialize",
 				);
 			}
-			await Promise.all([
-				readyPromise,
-				listenerPromise,
-				serveLifecyclePromise,
-			]);
+			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);
@@ -473,6 +504,9 @@ 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,
 			{

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 {
@@ -21,6 +22,7 @@ import type {
 	RuntimeCronFire,
 	RuntimeCronJobInfo,
 	RuntimeHttpRequest,
+	RuntimeHttpResponseBodyStream,
 	RuntimeKvEntry,
 	RuntimeKvListOptions,
 	RuntimeListenerConfig,
@@ -154,23 +156,48 @@ function toNapiApplication(application: RuntimeApplicationFetch) {
 			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),
-		});
+		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: toNapiBuffer(response.body),
+			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] {

package/src/registry/native-http.ts

@@ -0,0 +1,225 @@
+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) {
+				// Web Streams closes pending reads before waiting for the source's cancel
+				// callback. Do not let an arbitrary user cancel promise retain the native
+				// request and connection cleanup indefinitely.
+				void reader
+					.cancel(
+						new Error(
+							"native http response stream receiver dropped",
+						),
+					)
+					.catch((error) => {
+						logger().debug({
+							msg: "failed to cancel native http response reader",
+							error,
+						});
+					});
+				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,
+			});
+		}
+		void reader.cancel(error).catch(() => {
+			// Reader may already be closed after a native-side disconnect.
+		});
+	} finally {
+		try {
+			reader.releaseLock();
+		} catch (error) {
+			logger().debug({
+				msg: "failed to release native http response reader lock",
+				error,
+			});
+		}
+	}
+}
+
+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,
+};