rivetkit 2.3.7 → 2.3.8

package/dist/tsup/mod.js

@@ -1,6 +1,6 @@
 import {
   decodeWorkflowHistoryTransport
-} from "./chunk-6D2WRHJV.js";
+} from "./chunk-HWBISPRP.js";
 import {
   ACTOR_CONTEXT_INTERNAL_SYMBOL,
   ActorConfigSchema,
@@ -16,18 +16,20 @@ import {
 import {
   ALLOWED_PUBLIC_HEADERS,
   HEADER_CONN_PARAMS,
+  HEADER_ORIGINAL_REQUEST_URL,
   RemoteEngineControlClient,
   convertRegistryConfigToClientConfig,
   createClientWithDriver,
   getDatacenters,
+  isResponseLike,
   tryParseEndpoint,
   updateRunnerConfig
-} from "./chunk-YU3C5OCU.js";
+} from "./chunk-BL6ERZQO.js";
 import {
   makePrefixedKey,
   removePrefixFromKey
 } from "./chunk-JI6GZ2C2.js";
-import "./chunk-YHPFCYL3.js";
+import "./chunk-V7PJYBA5.js";
 import "./chunk-7AFZMFPQ.js";
 import {
   LogLevelSchema,
@@ -61,7 +63,7 @@ import {
   parsePortEnv,
   stringifyError,
   toUint8Array
-} from "./chunk-OG4OR2CM.js";
+} from "./chunk-ZKDZW5DQ.js";
 import {
   INTERNAL_ERROR_CODE,
   RivetError,
@@ -1433,6 +1435,36 @@ function toNapiSqlBatchStatement(statement) {
 function toNapiBuffer(value) {
   return Buffer.from(value);
 }
+function toNapiApplication(application) {
+  return async (error, request) => {
+    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
+      },
+      request.responseBodyStream ? fromNapiHttpResponseBodyStream(
+        request.responseBodyStream
+      ) : void 0
+    );
+    return {
+      ...response,
+      body: response.body === void 0 ? void 0 : toNapiBuffer(response.body)
+    };
+  };
+}
+function fromNapiHttpResponseBodyStream(stream) {
+  return {
+    cancelled: () => stream.cancelled(),
+    write: (chunk) => stream.write(Buffer.from(chunk)),
+    end: () => stream.end(),
+    error: (message) => stream.error(message)
+  };
+}
 function toNapiHttpRequest(request) {
   if (!request) {
     return request;
@@ -1539,12 +1571,25 @@ var NapiCoreRuntime = class {
     );
   }
   async serveListener(registry, listener, config) {
+    const application = listener.application ? toNapiApplication(listener.application) : void 0;
     await asNativeRegistry(registry).serveListener(
       {
         port: listener.port,
         host: listener.host,
         publicDir: listener.publicDir
       },
+      config,
+      application
+    );
+  }
+  async serveApplicationListener(registry, listener, config) {
+    await asNativeRegistry(registry).serveApplicationListener(
+      {
+        port: listener.port,
+        host: listener.host,
+        publicDir: listener.publicDir
+      },
+      toNapiApplication(listener.application),
       config
     );
   }
@@ -1932,6 +1977,137 @@ async function loadNapiRuntime() {
   };
 }
 
+// src/registry/native-http.ts
+var HTTP_BODY_CHUNK_SIZE = 64 * 1024;
+function runtimeBytesToArrayBuffer(value) {
+  return value.buffer.slice(
+    value.byteOffset,
+    value.byteOffset + value.byteLength
+  );
+}
+function buildNativeHttpRequest(init) {
+  var _a;
+  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 method2 = init.method.toUpperCase();
+  const bodyForbidden = method2 === "GET" || method2 === "HEAD";
+  const body = bodyForbidden ? void 0 : init.bodyStream ? new ReadableStream({
+    async pull(controller) {
+      var _a2, _b;
+      try {
+        if (init.body && init.body.length > 0) {
+          controller.enqueue(new Uint8Array(init.body));
+          init.body = void 0;
+          return;
+        }
+        const chunk = await ((_a2 = init.bodyStream) == null ? void 0 : _a2.read());
+        if (!chunk) {
+          controller.close();
+        } else {
+          controller.enqueue(new Uint8Array(chunk));
+        }
+      } catch (error) {
+        (_b = init.abortController) == null ? void 0 : _b.abort(error);
+        controller.error(error);
+      }
+    },
+    async cancel() {
+      var _a2;
+      await ((_a2 = init.bodyStream) == null ? void 0 : _a2.cancel());
+    }
+  }) : init.body && init.body.length > 0 ? runtimeBytesToArrayBuffer(init.body) : void 0;
+  const streamInit = init.bodyStream && !bodyForbidden ? { duplex: "half" } : {};
+  return new Request(url, {
+    method: method2,
+    headers,
+    body,
+    signal: ((_a = init.abortController) == null ? void 0 : _a.signal) ?? init.signal,
+    ...streamInit
+  });
+}
+async function cancelNativeHttpRequestBody(bodyStream) {
+  await (bodyStream == null ? void 0 : bodyStream.cancel());
+}
+async function writeResponseChunk(stream, chunk) {
+  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, stream) {
+  var _a;
+  const cancelled = stream.cancelled().then(() => ({ cancelled: true }));
+  try {
+    for (; ; ) {
+      const read = reader.read().then((next2) => ({ cancelled: false, next: next2 }));
+      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 ((_a = next.value) == null ? void 0 : _a.byteLength) {
+        await writeResponseChunk(stream, next.value);
+      }
+    }
+  } catch (error) {
+    try {
+      await stream.error(stringifyError(error));
+    } catch (streamError) {
+      logger2().debug({
+        msg: "failed to report native http response stream error",
+        error: streamError
+      });
+    }
+    try {
+      await reader.cancel(error);
+    } catch {
+    }
+  } finally {
+    reader.releaseLock();
+  }
+}
+async function convertNativeHttpResponse(response, responseBodyStream) {
+  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())
+    }
+  };
+}
+
 // src/registry/native-validation.ts
 var CONN_PARAMS_KEY = "__conn_params__";
 function validateActionArgs(schemas, name, args) {
@@ -2204,6 +2380,11 @@ var WasmCoreRuntime = class {
       "registry.listen() is not supported on the wasm runtime; use registry.serve() and mount the handler in your platform's HTTP server instead"
     );
   }
+  async serveApplicationListener(_registry, _listener, _config) {
+    throw new Error(
+      "registry.listen() is not supported on the wasm runtime; use an application-owned HTTP server instead"
+    );
+  }
   createActorFactory(callbacks, config) {
     return callWasmSync(
       () => asActorFactoryHandle2(
@@ -3061,7 +3242,7 @@ function toRuntimeBytes(value) {
 function arrayBufferViewToRuntimeBytes(value) {
   return new Uint8Array(value.buffer, value.byteOffset, value.byteLength);
 }
-function runtimeBytesToArrayBuffer(value) {
+function runtimeBytesToArrayBuffer2(value) {
   return value.buffer.slice(
     value.byteOffset,
     value.byteOffset + value.byteLength
@@ -3391,24 +3572,6 @@ function decodeArgs(value) {
   const decoded = decodeValue(value);
   return normalizeArgs(decoded);
 }
-function buildRequest(init) {
-  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) : void 0;
-  return new Request(url, {
-    method: init.method,
-    headers: init.headers,
-    body
-  });
-}
-async function toRuntimeHttpResponse(response) {
-  const headers = Object.fromEntries(response.headers.entries());
-  const body = new Uint8Array(await response.arrayBuffer());
-  return {
-    status: response.status,
-    headers,
-    body
-  };
-}
 function toActorKey(segments) {
   return segments.map(
     (segment) => segment.kind === "number" ? String(segment.numberValue ?? 0) : segment.stringValue ?? ""
@@ -4073,7 +4236,7 @@ var NativeWebSocketAdapter = class {
     this.#runtime.webSocketSetEventCallback(this.#ws, (event2) => {
       if (event2.kind === "message") {
         this.#virtual.triggerMessage(
-          event2.binary ? runtimeBytesToArrayBuffer(event2.data) : event2.data,
+          event2.binary ? runtimeBytesToArrayBuffer2(event2.data) : event2.data,
           event2.messageIndex
         );
         return;
@@ -5305,12 +5468,14 @@ function buildNativeFactory(runtime, registryConfig, definition) {
     onStateChange,
     cancelToken
   );
-  const maybeHandleNativeInspectorRequest = async (ctx, _rawRequest, jsRequest) => {
+  const maybeHandleNativeInspectorRequest = async (ctx, rawRequest) => {
     var _a2, _b, _c, _d;
-    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 void 0;
     }
+    const jsRequest = buildNativeHttpRequest(rawRequest);
     const jsonResponse = (body, init) => new Response(JSON.stringify(body), {
       status: (init == null ? void 0 : init.status) ?? 200,
       headers: {
@@ -5841,7 +6006,7 @@ function buildNativeFactory(runtime, registryConfig, definition) {
         );
         const actorCtx = makeActorCtx(
           ctx,
-          request ? buildRequest(request) : void 0
+          request ? buildNativeHttpRequest(request) : void 0
         );
         try {
           await config.onBeforeConnect(
@@ -5861,7 +6026,7 @@ function buildNativeFactory(runtime, registryConfig, definition) {
         const { ctx, conn, params, request } = unwrapTsfnPayload(error, payload);
         const actorCtx = makeActorCtx(
           ctx,
-          request ? buildRequest(request) : void 0
+          request ? buildNativeHttpRequest(request) : void 0
         );
         const connAdapter = new NativeConnAdapter(
           runtime,
@@ -5898,7 +6063,7 @@ function buildNativeFactory(runtime, registryConfig, definition) {
         );
         const actorCtx = makeActorCtx(
           ctx,
-          request ? buildRequest(request) : void 0
+          request ? buildNativeHttpRequest(request) : void 0
         );
         const connAdapter = new NativeConnAdapter(
           runtime,
@@ -6003,27 +6168,45 @@ function buildNativeFactory(runtime, registryConfig, definition) {
     onRequest: wrapNativeCallback(
       async (error, payload) => {
         try {
-          const { ctx, request, cancelToken } = unwrapTsfnPayload(
-            error,
-            payload
-          );
-          const jsRequest = buildRequest(request);
-          const inspectorResponse = await maybeHandleNativeInspectorRequest(
-            ctx,
-            request,
-            jsRequest
-          );
+          const { ctx, request, cancelToken, responseBodyStream } = unwrapTsfnPayload(error, payload);
+          const inspectorResponse = 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 rawConnParams = jsRequest.headers.get(HEADER_CONN_PARAMS);
+          const requestAbortController = new AbortController();
+          const handlerRequest = buildNativeHttpRequest({
+            ...request,
+            abortController: requestAbortController
+          });
+          const rawConnParams = handlerRequest.headers.get(HEADER_CONN_PARAMS);
           let requestCtx;
           let conn;
+          let removeRequestAbortListener;
+          let cleanupDeferredToBody = false;
+          let cleanedUp = false;
+          const cleanupRequest = async () => {
+            if (cleanedUp) return;
+            cleanedUp = true;
+            removeRequestAbortListener == null ? void 0 : removeRequestAbortListener();
+            try {
+              await (requestCtx == null ? void 0 : requestCtx.dispose());
+            } finally {
+              if (conn) {
+                await runtime.connDisconnect(conn);
+              }
+            }
+          };
           try {
             const connParams = validateConnParams(
               schemaConfig.connParamsSchema,
@@ -6039,23 +6222,56 @@ function buildNativeFactory(runtime, registryConfig, definition) {
             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 (!(response instanceof Response)) {
+            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) => {
+                logger2().error({
+                  msg: "failed to clean up native streaming http request",
+                  error: cleanupError
+                });
+              });
+            }
+            return conversion.response;
           } finally {
-            await (requestCtx == null ? void 0 : requestCtx.dispose());
-            if (conn) {
-              await runtime.connDisconnect(conn);
+            try {
+              await cancelNativeHttpRequestBody(
+                request.bodyStream
+              );
+            } finally {
+              if (!cleanupDeferredToBody) {
+                await cleanupRequest();
+              }
             }
           }
         } catch (error2) {
@@ -6070,7 +6286,7 @@ function buildNativeFactory(runtime, registryConfig, definition) {
     onWebSocket: typeof config.onWebSocket === "function" ? wrapNativeCallback(
       async (error, payload) => {
         const { ctx, conn, ws, request } = unwrapTsfnPayload(error, payload);
-        const jsRequest = request ? buildRequest(request) : void 0;
+        const jsRequest = request ? buildNativeHttpRequest(request) : void 0;
         const actorCtx = makeConnCtx(ctx, conn, jsRequest);
         try {
           await config.onWebSocket(
@@ -6167,7 +6383,7 @@ function buildNativeFactory(runtime, registryConfig, definition) {
           timeoutMs,
           cancelToken
         } = unwrapTsfnPayload(error, payload);
-        const jsRequest = buildRequest(request);
+        const jsRequest = buildNativeHttpRequest(request);
         const actorCtx = withConnContext(
           runtime,
           ctx,
@@ -6318,6 +6534,54 @@ function finishShutdownSignal(signal) {
   }
   process.kill(process.pid, signal);
 }
+function createApplicationFetch(application, runtime) {
+  return async (request, responseBodyStream) => {
+    const method2 = request.method.toUpperCase();
+    const body = method2 === "GET" || method2 === "HEAD" || request.body.length === 0 ? void 0 : 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: method2,
+        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) => {
+        logger2().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
+    };
+  };
+}
 var Registry = class {
   #config;
   #buildConfiguredRegistry;
@@ -6333,6 +6597,7 @@ var Registry = class {
   #runtimeReadyPromise;
   #runtimeServeConfiguredPromise;
   #runtimeServerlessPromise;
+  #applicationListenerPromise;
   #configureServerlessPoolPromise;
   #welcomePrinted = false;
   #shutdownInstalled = false;
@@ -6558,18 +6823,67 @@ var Registry = class {
    * @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 });
    * ```
    */
   async listen(opts = {}) {
     const port = opts.port ?? parsePortEnv(process.env.RIVET_PORT) ?? 3e3;
     const publicDir = opts.publicDir ?? getRivetkitPublicDir();
     const config = this.parseConfig();
-    const configuredRegistryPromise = buildConfiguredRegistry(config);
+    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 configuredRegistryPromise2 = this.#runtimeServeConfiguredPromise;
+      if (!configuredRegistryPromise2) {
+        throw new Error("registry envoy startup did not initialize");
+      }
+      const { runtime: runtime2, registry: registry2, serveConfig: serveConfig2 } = await configuredRegistryPromise2;
+      const application2 = createApplicationFetch(
+        opts.application,
+        runtime2
+      );
+      await readyPromise;
+      const listenerPromise = runtime2.serveApplicationListener(
+        registry2,
+        {
+          port,
+          host: opts.host,
+          publicDir,
+          application: application2
+        },
+        serveConfig2
+      );
+      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;
+    }
+    const configuredRegistryPromise = this.#buildConfiguredRegistry(config);
     this.#runtimeServeConfiguredPromise = configuredRegistryPromise;
     this.#runtimeServerlessPromise = configuredRegistryPromise;
     this.#installSignalHandlers(config);
@@ -6580,9 +6894,15 @@ var Registry = class {
     });
     this.#ensureServerlessPoolConfigured(config);
     const { runtime, registry, serveConfig } = await configuredRegistryPromise;
+    const application = opts.application ? createApplicationFetch(opts.application, runtime) : void 0;
     await runtime.serveListener(
       registry,
-      { port, host: opts.host, publicDir },
+      {
+        port,
+        host: opts.host,
+        publicDir,
+        application
+      },
       serveConfig
     );
   }
@@ -6682,11 +7002,9 @@ var Registry = class {
           await runtime.serveRegistry(registry, serveConfig);
         }
       );
-      this.#runtimeServePromise = this.#runtimeServeLifecyclePromise.catch(
-        (error) => {
-          logger2().warn({ error }, "runtime registry serve errored");
-        }
-      );
+      this.#runtimeServePromise = this.#runtimeServeLifecyclePromise.catch((error) => {
+        logger2().warn({ error }, "runtime registry serve errored");
+      });
       this.#installSignalHandlers(config);
     }
     if (printWelcome) {
@@ -6793,6 +7111,9 @@ var Registry = class {
       if (runtimeServePromise !== void 0) {
         await runtimeServePromise.catch(() => void 0);
       }
+      if (this.#applicationListenerPromise !== void 0) {
+        await this.#applicationListenerPromise.catch(() => void 0);
+      }
     };
     await Promise.race([
       drain(),
@@ -6843,7 +7164,9 @@ var Registry = class {
   startAndWait() {
     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;
@@ -6884,7 +7207,9 @@ var Registry = class {
     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),