rivetkit 2.3.12-rc.3 → 2.3.12

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,
+};

package/src/registry/native.ts

@@ -57,7 +57,7 @@ import {
 } from "@/common/database/mod";
 import { wrapJsNativeDatabase } from "@/common/database/native-database";
 import { assertJsonCompatValue, type JsonCompatValue } from "@/common/encoding";
-import { isResponseLike, type ResponseLike } from "@/common/fetch-like";
+import { isResponseLike } from "@/common/fetch-like";
 import {
 	decodeWorkflowHistoryTransport,
 	type WorkflowHistoryBytes,
@@ -85,6 +85,13 @@ import {
 } from "@/utils/node";
 import { logger } from "./log";
 import { loadNapiRuntime } from "./napi-runtime";
+import {
+	buildNativeHttpRequest,
+	cancelNativeHttpRequestBody,
+	convertNativeHttpResponse,
+	type NativeHttpRequestBodyStream,
+	type NativeHttpResponseBodyStream,
+} from "./native-http";
 import {
 	type NativeValidationConfig,
 	validateActionArgs,
@@ -105,7 +112,6 @@ import type {
 	RuntimeBytes,
 	RuntimeCronFire,
 	RuntimeCronJobInfo,
-	RuntimeHttpResponse,
 	RuntimeInspectorTabEntry,
 	RuntimeQueueMessage,
 	RuntimeScheduledEventInfo,
@@ -746,6 +752,12 @@ async function loadEngineCli(): Promise<typeof import("@rivetkit/engine-cli")> {
 	return import(["@rivetkit", "engine-cli"].join("/"));
 }
 
+async function loadServices(): Promise<typeof import("@rivet-dev/services")> {
+	// Keep the native binary opaque to edge bundlers for the same reason as
+	// engine-cli above.
+	return import(["@rivet-dev", "services"].join("/"));
+}
+
 function decodeValue<T>(value?: RuntimeBytes | null): T {
 	if (!value || value.length === 0) {
 		return undefined as T;
@@ -1250,38 +1262,6 @@ function decodeArgs(value?: RuntimeBytes | null): unknown[] {
 	return normalizeArgs(decoded);
 }
 
-function buildRequest(init: {
-	method: string;
-	uri: string;
-	headers?: Record<string, string>;
-	body?: RuntimeBytes;
-}): Request {
-	const url = init.uri.startsWith("http")
-		? init.uri
-		: new URL(init.uri, "http://127.0.0.1").toString();
-	const body =
-		init.body && init.body.length > 0
-			? runtimeBytesToArrayBuffer(init.body)
-			: undefined;
-	return new Request(url, {
-		method: init.method,
-		headers: init.headers,
-		body,
-	});
-}
-
-async function toRuntimeHttpResponse(
-	response: ResponseLike,
-): Promise<RuntimeHttpResponse> {
-	const headers = Object.fromEntries(response.headers.entries());
-	const body = new Uint8Array(await response.arrayBuffer());
-	return {
-		status: response.status,
-		headers,
-		body,
-	};
-}
-
 function toActorKey(
 	segments: Array<{
 		kind: string;
@@ -4068,19 +4048,22 @@ export function buildNativeFactory(
 		);
 	const maybeHandleNativeInspectorRequest = async (
 		ctx: ActorContextHandle,
-		_rawRequest: {
+		rawRequest: {
 			method: string;
 			uri: string;
 			headers?: Record<string, string>;
 			body?: RuntimeBytes;
 		},
-		jsRequest: Request,
 	): Promise<Response | undefined> => {
-		const url = new URL(jsRequest.url);
+		const rawUrl = rawRequest.uri.startsWith("http")
+			? rawRequest.uri
+			: new URL(rawRequest.uri, "http://127.0.0.1").toString();
+		const url = new URL(rawUrl);
 		if (!url.pathname.startsWith("/inspector/")) {
 			return undefined;
 		}
 
+		const jsRequest = buildNativeHttpRequest(rawRequest);
 		const jsonResponse = (body: unknown, init?: ResponseInit) =>
 			new Response(JSON.stringify(body), {
 				status: init?.status ?? 200,
@@ -4772,7 +4755,9 @@ export function buildNativeFactory(
 							);
 							const actorCtx = makeActorCtx(
 								ctx,
-								request ? buildRequest(request) : undefined,
+								request
+									? buildNativeHttpRequest(request)
+									: undefined,
 							);
 							try {
 								await config.onBeforeConnect(
@@ -4809,7 +4794,9 @@ export function buildNativeFactory(
 								unwrapTsfnPayload(error, payload);
 							const actorCtx = makeActorCtx(
 								ctx,
-								request ? buildRequest(request) : undefined,
+								request
+									? buildNativeHttpRequest(request)
+									: undefined,
 							);
 							const connAdapter = new NativeConnAdapter(
 								runtime,
@@ -4865,7 +4852,9 @@ export function buildNativeFactory(
 							);
 							const actorCtx = makeActorCtx(
 								ctx,
-								request ? buildRequest(request) : undefined,
+								request
+									? buildNativeHttpRequest(request)
+									: undefined,
 							);
 							const connAdapter = new NativeConnAdapter(
 								runtime,
@@ -5017,38 +5006,63 @@ export function buildNativeFactory(
 						uri: string;
 						headers?: Record<string, string>;
 						body?: RuntimeBytes;
+						bodyStream?: NativeHttpRequestBodyStream;
 					};
 					cancelToken?: CancellationTokenHandle;
+					responseBodyStream?: NativeHttpResponseBodyStream;
 				},
 			) => {
 				try {
-					const { ctx, request, cancelToken } = unwrapTsfnPayload(
-						error,
-						payload,
-					);
-					const jsRequest = buildRequest(request);
+					const { ctx, request, cancelToken, responseBodyStream } =
+						unwrapTsfnPayload(error, payload);
 					const inspectorResponse =
-						await maybeHandleNativeInspectorRequest(
-							ctx,
-							request,
-							jsRequest,
-						);
+						await maybeHandleNativeInspectorRequest(ctx, request);
 					if (inspectorResponse) {
-						return await toRuntimeHttpResponse(inspectorResponse);
+						await cancelNativeHttpRequestBody(request.bodyStream);
+						return (
+							await convertNativeHttpResponse(
+								inspectorResponse,
+								responseBodyStream,
+							)
+						).response;
 					}
 
 					if (typeof config.onRequest !== "function") {
-						return await toRuntimeHttpResponse(
-							new Response(null, { status: 404 }),
-						);
+						await cancelNativeHttpRequestBody(request.bodyStream);
+						return (
+							await convertNativeHttpResponse(
+								new Response(null, { status: 404 }),
+								responseBodyStream,
+							)
+						).response;
 					}
 
+					const requestAbortController = new AbortController();
+					const handlerRequest = buildNativeHttpRequest({
+						...request,
+						abortController: requestAbortController,
+					});
 					const rawConnParams =
-						jsRequest.headers.get(HEADER_CONN_PARAMS);
+						handlerRequest.headers.get(HEADER_CONN_PARAMS);
 					let requestCtx:
 						| ReturnType<typeof withConnContext>
 						| undefined;
 					let conn: ConnHandle | undefined;
+					let removeRequestAbortListener: (() => void) | undefined;
+					let cleanupDeferredToBody = false;
+					let cleanedUp = false;
+					const cleanupRequest = async () => {
+						if (cleanedUp) return;
+						cleanedUp = true;
+						removeRequestAbortListener?.();
+						try {
+							await requestCtx?.dispose();
+						} finally {
+							if (conn) {
+								await runtime.connDisconnect(conn);
+							}
+						}
+					};
 					try {
 						const connParams = validateConnParams(
 							schemaConfig.connParamsSchema,
@@ -5066,23 +5080,62 @@ export function buildNativeFactory(
 						requestCtx = makeConnCtx(
 							ctx,
 							conn,
-							jsRequest,
+							handlerRequest,
 							cancelToken,
 						);
+						const ctxAbortSignal = requestCtx.abortSignal;
+						const abortRequest = () =>
+							requestAbortController.abort(ctxAbortSignal.reason);
+						if (ctxAbortSignal.aborted) {
+							abortRequest();
+						} else {
+							ctxAbortSignal.addEventListener(
+								"abort",
+								abortRequest,
+								{ once: true },
+							);
+							removeRequestAbortListener = () =>
+								ctxAbortSignal.removeEventListener(
+									"abort",
+									abortRequest,
+								);
+						}
 						const response = await config.onRequest(
 							requestCtx,
-							jsRequest,
+							handlerRequest,
 						);
 						if (!isResponseLike(response)) {
 							throw new Error(
 								"onRequest handler must return a Response",
 							);
 						}
-						return await toRuntimeHttpResponse(response);
+						const conversion = await convertNativeHttpResponse(
+							response,
+							responseBodyStream,
+						);
+						if (conversion.bodyCompletion) {
+							cleanupDeferredToBody = true;
+							void conversion.bodyCompletion
+								.then(cleanupRequest)
+								.catch((cleanupError) => {
+									logger().error({
+										msg: "failed to clean up native streaming http request",
+										error: cleanupError,
+									});
+								});
+						}
+						return conversion.response;
 					} finally {
-						await requestCtx?.dispose();
-						if (conn) {
-							await runtime.connDisconnect(conn);
+						try {
+							// Handler completion ends upload ownership even when
+							// the Web Request body is locked or partly consumed.
+							await cancelNativeHttpRequestBody(
+								request.bodyStream,
+							);
+						} finally {
+							if (!cleanupDeferredToBody) {
+								await cleanupRequest();
+							}
 						}
 					}
 				} catch (error) {
@@ -5114,7 +5167,7 @@ export function buildNativeFactory(
 							const { ctx, conn, ws, request } =
 								unwrapTsfnPayload(error, payload);
 							const jsRequest = request
-								? buildRequest(request)
+								? buildNativeHttpRequest(request)
 								: undefined;
 							const actorCtx = makeConnCtx(ctx, conn, jsRequest);
 							try {
@@ -5281,7 +5334,7 @@ export function buildNativeFactory(
 					timeoutMs,
 					cancelToken,
 				} = unwrapTsfnPayload(error, payload);
-				const jsRequest = buildRequest(request);
+				const jsRequest = buildNativeHttpRequest(request);
 				const actorCtx = withConnContext(
 					runtime,
 					ctx,
@@ -5380,6 +5433,7 @@ export function buildNativeFactory(
 
 export async function buildServeConfig(
 	config: RegistryConfig,
+	resolveServicesBinary = true,
 ): Promise<RuntimeServeConfig> {
 	if (!config.endpoint) {
 		throw nativeEndpointNotConfiguredError();
@@ -5391,6 +5445,7 @@ export async function buildServeConfig(
 		token: config.token,
 		namespace: config.namespace,
 		poolName: config.envoy.poolName,
+		startServices: resolveServicesBinary && config.startServices,
 		handleInspectorHttpInRuntime: true,
 		serverlessBasePath: config.serverless.basePath,
 		serverlessPackageVersion: VERSION,
@@ -5422,6 +5477,17 @@ export async function buildServeConfig(
 			error: stringifyError(error),
 		});
 	}
+	if (resolveServicesBinary && config.startServices) {
+		try {
+			const { getServicesPath } = await loadServices();
+			serveConfig.servicesBinaryPath = getServicesPath();
+		} catch (error) {
+			logger().warn({
+				msg: "could not resolve the Services binary; a local engine will fail with services.binary_unavailable — set RIVET_RUN_SERVICES=0 to disable Services or RIVET_SERVICES_BINARY to a local rivet-services binary",
+				error: stringifyError(error),
+			});
+		}
+	}
 	serveConfig.engineHost = config.engineHost;
 	serveConfig.enginePort = config.enginePort;
 	if (config.test?.enabled) {
@@ -5467,7 +5533,7 @@ export async function buildRegistryWithRuntime(
 	return {
 		runtime,
 		registry,
-		serveConfig: await buildServeConfig(config),
+		serveConfig: await buildServeConfig(config, runtime.kind === "napi"),
 	};
 }