next-ws 0.0.0-beta-20250821061718 → 0.0.0-beta-20250823112013

package/README.md

@@ -50,14 +50,14 @@ To set up a WebSocket server with `next-ws`, you need to patch your local Next.j
 
 ## 🚀 Usage
 
-Using WebSocket connections in your Next.js app directory is simple with `next-ws`. You can handle WebSocket connections directly in your API routes via exported `SOCKET` functions. Here's an example of a simple WebSocket echo server:
+Using WebSocket connections in your Next.js app directory is simple with `next-ws`. You can handle WebSocket connections directly in your API routes via exported `UPGRADE` functions.
 
 ```js
-export function SOCKET(
+export function UPGRADE(
   client: import('ws').WebSocket,
-  request: import('http').IncomingMessage,
   server: import('ws').WebSocketServer,
-  context: { params: Record<string, string | string[]> },
+  request: import('next/server').NextRequest,
+  context: import('next-ws/server').RouteContext<'/api/ws'>,
 ) {
   // ...
 }
@@ -70,25 +70,24 @@ export function SOCKET(
 
 ### Echo Server
 
-This example demonstrates a simple WebSocket echo server that sends back any message it receives. Create a new API route file anywhere in your app directory and export a `SOCKET` function to handle WebSocket connections:
+This example demonstrates a simple WebSocket echo server that sends back any message it receives. Create a new API route file anywhere in your app directory and export a `UPGRADE` function to handle WebSocket connections:
 
 ```ts
 // app/api/ws/route.ts (can be any route file in the app directory)
 
-export function SOCKET(
-  client: import("ws").WebSocket,
-  request: import("http").IncomingMessage,
-  server: import("ws").WebSocketServer
+export function UPGRADE(
+  client: import('ws').WebSocket,
+  server: import('ws').WebSocketServer
 ) {
-  console.log("A client connected");
+  console.log('A client connected');
 
-  client.on("message", (message) => {
-    console.log("Received message:", message);
+  client.on('message', (message) => {
+    console.log('Received message:', message);
     client.send(message);
   });
 
-  client.on("close", () => {
-    console.log("A client disconnected");
+  client.once('close', () => {
+    console.log('A client disconnected');
   });
 }
 ```

package/dist/cli.cjs

@@ -27,7 +27,7 @@ var __toESM = (mod, isNodeMode, target) => (target = mod != null ? __create(__ge
 var import_minimist = __toESM(require("minimist"));
 
 // package.json
-var version = "0.0.0-beta-20250821061718";
+var version = "0.0.0-beta-20250823112013";
 
 // src/commands/helpers/define.ts
 function defineCommandGroup(definition) {
@@ -79,7 +79,39 @@ Options:
 `;
 }
 
+// src/patches/helpers/next.ts
+var import_promises = require("fs/promises");
+var import_node_path = require("path");
+function resolveNextWsDirectory() {
+  const id = (
+    //
+    require.resolve("next-ws/package.json", { paths: [process.cwd()] })
+  );
+  return (0, import_node_path.dirname)(id);
+}
+function resolveNextDirectory() {
+  const id = (
+    //
+    require.resolve("next/package.json", { paths: [process.cwd()] })
+  );
+  return (0, import_node_path.dirname)(id);
+}
+async function getInstalledNextVersion() {
+  const id = (0, import_node_path.join)(resolveNextDirectory(), "package.json");
+  const pkg = await (0, import_promises.readFile)(id, "utf8").then(JSON.parse);
+  return String(pkg.version.split("-")[0]);
+}
+async function readTrace() {
+  const id = (0, import_node_path.join)(resolveNextDirectory(), ".next-ws-trace.json");
+  return (0, import_promises.readFile)(id, "utf-8").then(JSON.parse).catch(() => null);
+}
+async function writeTrace(trace) {
+  const id = (0, import_node_path.join)(resolveNextDirectory(), ".next-ws-trace.json");
+  await (0, import_promises.writeFile)(id, JSON.stringify(trace, null, 2));
+}
+
 // src/patches/patch-1.ts
+var import_node_path3 = require("path");
 var import_jscodeshift = __toESM(require("jscodeshift"));
 
 // src/patches/helpers/define.ts
@@ -627,7 +659,7 @@ async function task(promise, ...message) {
   let spinnerIndex = 0;
   const spinnerInterval = setInterval(() => {
     import_node_readline.default.cursorTo(process.stdout, 0);
-    const spinnerChar = spinnerChars[spinnerIndex++ % spinnerChars.length];
+    const spinnerChar = spinnerChars[spinnerIndex++ % spinnerChars.length] ?? " ";
     process.stdout.write(source_default.cyan("[next-ws]", spinnerChar, ...message));
   }, 100);
   return promise.then((value) => {
@@ -647,40 +679,6 @@ async function task(promise, ...message) {
   });
 }
 
-// src/patches/helpers/next.ts
-var import_promises = require("fs/promises");
-var import_node_path = require("path");
-function getDistDirname() {
-  const resolveOptions = { paths: [process.cwd()] };
-  const nextWsPackagePath = (
-    //
-    require.resolve("next-ws/package.json", resolveOptions)
-  );
-  const nextWsDirName = (0, import_node_path.dirname)(nextWsPackagePath);
-  return `${nextWsDirName}/dist`.replace(/\\/g, "/").replace(/'/g, "\\'");
-}
-function findNextDirectory() {
-  const resolveOptions = { paths: [process.cwd()] };
-  const nextPackagePath = require.resolve("next/package.json", resolveOptions);
-  return (0, import_node_path.dirname)(nextPackagePath);
-}
-async function getNextVersion() {
-  const nextDirectory = findNextDirectory();
-  const nextPackagePath = (0, import_node_path.join)(nextDirectory, "package.json");
-  const nextPackage = await (0, import_promises.readFile)(nextPackagePath, "utf-8").then(JSON.parse);
-  return String(nextPackage.version.split("-")[0]);
-}
-async function getTrace() {
-  const nextDirectory = findNextDirectory();
-  const tracePath = (0, import_node_path.join)(nextDirectory, ".next-ws-trace.json");
-  return (0, import_promises.readFile)(tracePath, "utf-8").then(JSON.parse).catch(() => null);
-}
-async function setTrace(trace) {
-  const nextDirectory = findNextDirectory();
-  const tracePath = (0, import_node_path.join)(nextDirectory, ".next-ws-trace.json");
-  await (0, import_promises.writeFile)(tracePath, JSON.stringify(trace, null, 2));
-}
-
 // src/patches/helpers/define.ts
 function definePatch(definition) {
   return {
@@ -707,7 +705,7 @@ function definePatchStep(definition) {
 function resolvePath(path) {
   switch (path.split(":")[0]) {
     case "next": {
-      const nextDirectory = findNextDirectory();
+      const nextDirectory = resolveNextDirectory();
       const realPath = path.slice(5);
       return (0, import_node_path2.resolve)(nextDirectory, realPath);
     }
@@ -718,38 +716,27 @@ function resolvePath(path) {
 }
 
 // src/patches/patch-1.ts
+var CommentLine = import_jscodeshift.default.Comment;
 var patchNextNodeServer = definePatchStep({
   title: "Add WebSocket server setup script to NextNodeServer constructor",
   path: "next:dist/server/next-server.js",
   async transform(code) {
-    const PATCH_MARKER = "@patch patchNextNodeServer";
-    const patchProgram = (0, import_jscodeshift.default)(`
-      // ${PATCH_MARKER}
+    const marker = "@patch attach-websocket-server";
+    const snippet = (0, import_jscodeshift.default)(`
+      // ${marker}
       let nextWs;
-      try { nextWs ??= require('next-ws/dist/server/index.cjs') } catch {}
-      try { nextWs ??= require('${getDistDirname()}/server/index.cjs') } catch {}
+      try { nextWs ??= require('next-ws/server') } catch {}
+      try { nextWs ??= require(require.resolve('next-ws/server', { paths: [process.cwd()] }) )} catch {}
+      try { nextWs ??= require('${resolveNextWsDirectory().replaceAll(import_node_path3.sep, "/").replaceAll("'", "\\'")}/dist/server/index.cjs') } catch {}
       nextWs?.setupWebSocketServer(this);
     `);
-    const patchBlock = (
-      //
-      import_jscodeshift.default.blockStatement(patchProgram.nodes()[0].program.body)
-    );
-    return (0, import_jscodeshift.default)(code).find(import_jscodeshift.default.ClassDeclaration, {
-      id: { name: "NextNodeServer" }
-    }).find(import_jscodeshift.default.MethodDefinition, {
-      kind: "constructor",
-      value: { body: { type: "BlockStatement" } }
-    }).forEach(({ node: method }) => {
-      const bodyStatements = method.value.body.body;
-      const existingPatchPath = (0, import_jscodeshift.default)(bodyStatements).find(import_jscodeshift.default.Comment, {
-        value: ` ${PATCH_MARKER}`
-      }).paths()[0];
-      const existingPatchIndex = bodyStatements.findIndex(
-        (s) => s === existingPatchPath?.parent.node
-      );
-      if (existingPatchIndex > -1)
-        bodyStatements[existingPatchIndex] = patchBlock;
-      else bodyStatements.push(patchBlock);
+    const block = import_jscodeshift.default.blockStatement(snippet.nodes()[0].program.body);
+    return (0, import_jscodeshift.default)(code).find(import_jscodeshift.default.ClassDeclaration, { id: { name: "NextNodeServer" } }).find(import_jscodeshift.default.MethodDefinition, { kind: "constructor" }).forEach(({ node }) => {
+      const body = node.value.body.body;
+      const existing = (0, import_jscodeshift.default)(body).find(CommentLine, { value: ` ${marker}` }).paths()[0];
+      const idx = body.indexOf(existing?.parent.node);
+      if (existing && idx > -1) body[idx] = block;
+      else body.push(block);
     }).toSource();
   }
 });
@@ -763,7 +750,11 @@ var patchRouterServer = definePatchStep({
         object: { type: "Identifier", name: "socket" },
         property: { type: "Identifier", name: "end" }
       }
-    }).replaceWith((path) => `null // ${(0, import_jscodeshift.default)(path.node).toSource()}`).toSource();
+    }).replaceWith((path) => {
+      const expr = import_jscodeshift.default.unaryExpression("void", import_jscodeshift.default.literal(0));
+      expr.comments = [import_jscodeshift.default.commentLine(` ${(0, import_jscodeshift.default)(path.node).toSource()}`)];
+      return expr;
+    }).toSource();
   }
 });
 var patch_1_default = definePatch({
@@ -776,13 +767,13 @@ var patch_1_default = definePatch({
 var patches_default = [patch_1_default];
 
 // src/commands/helpers/semver.ts
-var import_range = __toESM(require("semver/classes/range"));
-var import_semver = __toESM(require("semver/classes/semver"));
-var import_gt = __toESM(require("semver/functions/gt"));
-var import_min_version = __toESM(require("semver/ranges/min-version"));
-var import_ltr = __toESM(require("semver/ranges/ltr"));
-var import_gtr = __toESM(require("semver/ranges/gtr"));
+var import_range = __toESM(require("semver/classes/range.js"));
+var import_semver = __toESM(require("semver/classes/semver.js"));
+var import_gt = __toESM(require("semver/functions/gt.js"));
 var import_satisfies = __toESM(require("semver/functions/satisfies"));
+var import_gtr = __toESM(require("semver/ranges/gtr"));
+var import_ltr = __toESM(require("semver/ranges/ltr"));
+var import_min_version = __toESM(require("semver/ranges/min-version"));
 function maxVersion(range, loose) {
   range = new import_range.default(range, loose);
   let maximumVersion = null;
@@ -819,7 +810,7 @@ var patch_default = defineCommand({
     const supported = patches_default.map((p) => p.versions).join(" || ");
     const minimum = import_min_version.default(supported)?.version ?? supported;
     const maximum = maxVersion(supported)?.version ?? supported;
-    const current = await getNextVersion();
+    const current = await getInstalledNextVersion();
     if (import_ltr.default(current, minimum)) {
       error(
         `Next.js v${current} is not supported, a minimum of v${minimum} is required`
@@ -852,7 +843,7 @@ var patch_default = defineCommand({
     info(`Patching Next.js v${current} with '${patch.versions}'`);
     await patch.execute();
     info("Saving patch trace file...");
-    await setTrace({ patch: patch.versions, version: current });
+    await writeTrace({ patch: patch.versions, version: current });
     info("All done!");
   }
 });
@@ -869,7 +860,7 @@ var verify_default = defineCommand({
     }
   ],
   async action(options) {
-    const trace = await getTrace();
+    const trace = await readTrace();
     if (!trace) {
       if (options.ensure) {
         warn("Next.js has not been patched, running the patch command");
@@ -881,7 +872,7 @@ var verify_default = defineCommand({
         process.exit(1);
       }
     }
-    const current = await getNextVersion();
+    const current = await getInstalledNextVersion();
     if (current !== trace.version) {
       error(
         "Next.js has been patched with a different version, you'll need to run the patch command"

package/dist/client/index.cjs

@@ -39,13 +39,12 @@ module.exports = __toCommonJS(client_exports);
 
 // src/client/context.tsx
 var import_react = __toESM(require("react"));
-var import_react2 = require("react");
-var WebSocketContext = (0, import_react2.createContext)(null);
+var WebSocketContext = (0, import_react.createContext)(null);
 WebSocketContext.displayName = "WebSocketContext";
 var WebSocketConsumer = WebSocketContext.Consumer;
 function WebSocketProvider(p) {
   const clientRef = (0, import_react.useRef)(null);
-  (0, import_react2.useEffect)(() => {
+  (0, import_react.useEffect)(() => {
     if (typeof window === "undefined") return;
     if (clientRef.current) {
       clientRef.current.close();
@@ -62,7 +61,7 @@ function WebSocketProvider(p) {
   return /* @__PURE__ */ import_react.default.createElement(WebSocketContext.Provider, { value: clientRef.current }, p.children);
 }
 function useWebSocket() {
-  const context = (0, import_react2.useContext)(WebSocketContext);
+  const context = (0, import_react.useContext)(WebSocketContext);
   if (context === void 0)
     throw new Error("useWebSocket must be used within a WebSocketProvider");
   return context;

package/dist/server/index.cjs

@@ -38,11 +38,7 @@ __export(server_exports, {
 });
 module.exports = __toCommonJS(server_exports);
 
-// src/server/setup.ts
-var logger3 = __toESM(require("next/dist/build/output/log.js"));
-var import_ws = require("ws");
-
-// src/server/helpers/persistent.ts
+// src/server/persistent.ts
 var logger = __toESM(require("next/dist/build/output/log.js"));
 function getEnvironmentMeta() {
   const isCustomServer = !process.title.startsWith("next-");
@@ -50,91 +46,103 @@ function getEnvironmentMeta() {
   const isDevelopment = process.env.NODE_ENV === "development";
   return { isCustomServer, isMainProcess, isDevelopment };
 }
-function mainProcessOnly(fnName) {
+function mainProcessOnly(callerName) {
   if (process.env.NEXT_WS_SKIP_ENVIRONMENT_CHECK === "1") return;
   const meta = getEnvironmentMeta();
-  if (!meta.isMainProcess) {
+  if (!meta.isCustomServer && !meta.isMainProcess) {
     throw new Error(
-      `[next-ws] Attempt to invoke '${fnName}' outside the main process.
-You may be attempting to interact with the WebSocket server outside of a SOCKET handler. This will fail in production, as Next.js employs a worker process for routing, which do not have access to the WebSocket server on the main process.
+      `[next-ws] Attempt to call '${callerName}' outside the main process.
+You may be attempting to interact with the WebSocket server outside of an UPGRADE handler. This will fail in production, as Next.js employs a worker process for routing, which do not have access to the WebSocket server on the main process.
 You can resolve this by using a custom server.`
     );
   } else if (!meta.isCustomServer) {
     logger.warnOnce(
-      `[next-ws] Caution: The function '${fnName}' was invoked without a custom server.
-This could lead to unintended behaviour, especially if you're attempting to interact with the WebSocket server outside of a SOCKET handler.
+      `[next-ws] The function '${callerName}' was called without a custom server.
+This could lead to unintended behaviour, especially if you're attempting to interact with the WebSocket server outside of an UPGRADE handler.
 Please note, while such configurations might function during development, they will fail in production. This is because Next.js employs a worker process for routing in production, which do not have access to the WebSocket server on the main process.
 You can resolve this by using a custom server.`
     );
   }
 }
-var NextWsHttpServer = Symbol.for("NextWs_HttpServer");
-function setHttpServer(server) {
-  Reflect.set(globalThis, NextWsHttpServer, server);
-}
+var kHttpServer = Symbol.for("kHttpServer");
 function getHttpServer() {
   mainProcessOnly("getHttpServer");
-  return Reflect.get(globalThis, NextWsHttpServer);
+  return Reflect.get(globalThis, kHttpServer);
+}
+function setHttpServer(server) {
+  mainProcessOnly("setHttpServer");
+  Reflect.set(globalThis, kHttpServer, server);
+  return getHttpServer();
 }
 function useHttpServer(server) {
+  mainProcessOnly("useHttpServer");
   const existing = getHttpServer();
   if (existing) return existing;
-  if (server) setHttpServer(server);
-  return server;
-}
-var NextWsWebSocketServer = Symbol.for("NextWs_WebSocketServer");
-function setWebSocketServer(wsServer) {
-  Reflect.set(globalThis, NextWsWebSocketServer, wsServer);
+  return setHttpServer(server());
 }
+var kWebSocketServer = Symbol.for("kWebSocketServer");
 function getWebSocketServer() {
   mainProcessOnly("getWebSocketServer");
-  return Reflect.get(globalThis, NextWsWebSocketServer);
+  return Reflect.get(globalThis, kWebSocketServer);
 }
-function useWebSocketServer(wsServer) {
+function setWebSocketServer(server) {
+  mainProcessOnly("setWebSocketServer");
+  Reflect.set(globalThis, kWebSocketServer, server);
+  return getWebSocketServer();
+}
+function useWebSocketServer(server) {
+  mainProcessOnly("useWebSocketServer");
   const existing = getWebSocketServer();
   if (existing) return existing;
-  if (wsServer) setWebSocketServer(wsServer);
-  return wsServer;
+  return setWebSocketServer(server());
 }
 
-// src/server/helpers/route.ts
-var logger2 = __toESM(require("next/dist/build/output/log.js"));
-function createRouteRegex(routePattern) {
+// src/server/setup.ts
+var logger3 = __toESM(require("next/dist/build/output/log.js"));
+var import_ws = require("ws");
+
+// src/server/helpers/match.ts
+function compileRoutePattern(routePattern) {
   const escapedPattern = routePattern.replace(/[-/\\^$*+?.()|[\]{}]/g, "\\$&");
-  const paramRegex = escapedPattern.replace(/\\\[([a-zA-Z0-9_]+)\\\]/g, "(?<$1>[^/]+)").replace(/\\\[(?:\\\.){3}([a-zA-Z0-9_]+)\\\]/g, "(?<rest_$1>.+)");
+  const paramRegex = escapedPattern.replace(/\\\[\\\[(?:\\\.){3}([a-z0-9_]+)\\\]\\\]/gi, "?(?<r_o_$1>.+)?").replace(/\\\[(?:\\\.){3}([a-z0-9_]+)\\\]/gi, "(?<r_$1>.+)").replace(/\\\[([a-z0-9_]+)\\\]/gi, "(?<$1>[^/]+)");
   return new RegExp(`^${paramRegex}$`);
 }
-function getRouteParams(routePattern, routePath) {
-  const routeRegex = createRouteRegex(routePattern);
-  const match = routePath.replace(/\/+$/, "").match(routeRegex);
+function getRouteParams(routePattern, requestPathname) {
+  const routeRegex = compileRoutePattern(routePattern);
+  const match = requestPathname.replace(/\/+$/, "").match(routeRegex);
   if (!match) return null;
   if (!match.groups) return {};
   const params = {};
   for (let [k, v] of Object.entries(match.groups)) {
-    if (k.startsWith("rest_")) {
-      k = k.slice(5);
-      v = v.split("/");
+    if (k.startsWith("r_")) {
+      const optional = k.startsWith("r_o_");
+      k = k.slice(optional ? 4 : 2);
+      v = v?.split("/");
     }
-    Reflect.set(params, k, v);
+    if (v) Reflect.set(params, k, v);
   }
   return params;
 }
-function resolvePathToRoute(nextServer, requestPath) {
-  const basePath = nextServer.serverOptions.conf.basePath;
-  const routes = {
+function findMatchingRoute(nextServer, requestPathname) {
+  const basePath = nextServer.serverOptions?.conf.basePath || "";
+  const appPathRoutes = {
     // @ts-expect-error - appPathRoutes is protected
     ...nextServer.appPathRoutes,
     // @ts-expect-error - getAppPathRoutes is protected
     ...nextServer.getAppPathRoutes()
   };
-  let pathToRoute = null;
-  for (const [routePath, [filePath]] of Object.entries(routes)) {
+  let matchedRoute = void 0;
+  for (const [routePath, [filePath]] of Object.entries(appPathRoutes)) {
+    if (!routePath || !filePath) continue;
     const realPath = `${basePath}${routePath}`;
-    const routeParams = getRouteParams(realPath, requestPath);
-    if (routeParams) pathToRoute = { filePath, routeParams };
+    const routeParams = getRouteParams(realPath, requestPathname);
+    if (routeParams) matchedRoute = { filename: filePath, params: routeParams };
   }
-  return pathToRoute || null;
+  return matchedRoute;
 }
+
+// src/server/helpers/module.ts
+var logger2 = __toESM(require("next/dist/build/output/log.js"));
 async function importRouteModule(nextServer, filePath) {
   try {
     if ("hotReloader" in nextServer) {
@@ -151,52 +159,104 @@ async function importRouteModule(nextServer, filePath) {
     }
   } catch {
   }
-  const buildPagePath = nextServer.getPagePath(filePath);
-  return require(buildPagePath);
+  try {
+    const buildPagePath = nextServer.getPagePath(filePath);
+    return require(buildPagePath).routeModule;
+  } catch (cause) {
+    console.error(cause);
+    return void 0;
+  }
 }
-async function getSocketHandler(routeModule) {
-  if (routeModule.default instanceof Promise)
-    return getSocketHandler(await routeModule.default);
-  return routeModule?.default?.routeModule?.userland?.SOCKET ?? routeModule?.routeModule?.userland?.SOCKET ?? routeModule?.default?.handlers?.SOCKET ?? routeModule?.handlers?.SOCKET;
+
+// src/server/helpers/request.ts
+var import_server = require("next/server");
+function toNextRequest(message) {
+  const controller = new AbortController();
+  const headers = new Headers(message.headers);
+  const protocol = "encrypted" in message.socket ? "https" : "http";
+  const url = `${protocol}://${headers.get("host")}${message.url}`;
+  message.once("aborted", () => controller.abort());
+  return new import_server.NextRequest(url, {
+    method: message.method,
+    headers,
+    body: message.method === "GET" || message.method === "HEAD" ? void 0 : message,
+    signal: controller.signal,
+    referrer: headers.get("referer") || void 0
+  });
 }
 
 // src/server/setup.ts
 function setupWebSocketServer(nextServer) {
   process.env.NEXT_WS_MAIN_PROCESS = String(1);
   process.env.NEXT_WS_SKIP_ENVIRONMENT_CHECK = String(1);
-  const httpServer = useHttpServer(nextServer.serverOptions?.httpServer);
-  const wsServer = useWebSocketServer(new import_ws.WebSocketServer({ noServer: true }));
-  process.env.NEXT_WS_SKIP_ENVIRONMENT_CHECK = String(0);
+  const httpServer = (
+    //
+    // @ts-expect-error - serverOptions is protected
+    useHttpServer(() => nextServer.serverOptions?.httpServer)
+  );
   if (!httpServer)
     return logger3.error("[next-ws] was not able to find the HTTP server");
+  const wsServer = (
+    //
+    useWebSocketServer(() => new import_ws.WebSocketServer({ noServer: true }))
+  );
   if (!wsServer)
     return logger3.error("[next-ws] was not able to find the WebSocket server");
+  process.env.NEXT_WS_SKIP_ENVIRONMENT_CHECK = String(0);
   logger3.ready("[next-ws] has started the WebSocket server");
-  httpServer.on("upgrade", async (request, socket, head) => {
-    const url = new URL(request.url ?? "", "ws://next");
-    const pathname = url.pathname;
+  const kInstalled = Symbol.for("kInstalled");
+  if (Reflect.has(httpServer, kInstalled)) return;
+  Reflect.set(httpServer, kInstalled, true);
+  httpServer.on("upgrade", async (message, socket, head) => {
+    const request = toNextRequest(message);
+    const pathname = request.nextUrl.pathname;
     if (pathname.includes("/_next")) return;
-    const routeInfo = resolvePathToRoute(nextServer, pathname);
-    if (!routeInfo) {
-      logger3.error(`[next-ws] could not find module for page ${pathname}`);
-      return socket.destroy();
+    const route = findMatchingRoute(nextServer, pathname);
+    if (!route) {
+      logger3.error(`[next-ws] could not find route for page ${pathname}`);
+      return socket.end();
     }
-    const routeModule = await importRouteModule(nextServer, routeInfo.filePath);
-    if (!routeModule) {
-      logger3.error(`[next-ws] could not find module for page ${pathname}`);
-      return socket.destroy();
+    const module2 = await importRouteModule(nextServer, route.filename);
+    if (!module2) {
+      logger3.error(`[next-ws] could not import module for page ${pathname}`);
+      return socket.end();
     }
-    const socketHandler = await getSocketHandler(routeModule);
-    if (!socketHandler || typeof socketHandler !== "function") {
-      logger3.error(`[next-ws] ${pathname} does not export a SOCKET handler`);
-      return socket.destroy();
+    const handleUpgrade = module2.userland.UPGRADE;
+    const handleSocket = module2.userland.SOCKET;
+    if ((!handleUpgrade || typeof handleUpgrade !== "function") && (!handleSocket || typeof handleSocket !== "function")) {
+      logger3.error(`[next-ws] route '${pathname}' does not export a handler`);
+      return socket.end();
     }
-    return wsServer.handleUpgrade(request, socket, head, async (c, r) => {
-      const routeContext = { params: routeInfo.routeParams };
-      const handleClose = await socketHandler(c, r, wsServer, routeContext);
-      if (typeof handleClose === "function")
-        c.once("close", () => handleClose());
+    if (handleSocket)
+      logger3.warnOnce(
+        "DeprecationWarning: [next-ws] SOCKET is deprecated, use UPGRADE instead, see https://github.com/apteryxxyz/next-ws#-usage"
+      );
+    wsServer.handleUpgrade(message, socket, head, async (client) => {
+      wsServer.emit("connection", client, message);
+      try {
+        const context = { params: route.params };
+        if (handleUpgrade) {
+          await handleUpgrade(client, wsServer, request, context);
+        } else if (handleSocket) {
+          const handleClose = (
+            //
+            await handleSocket(client, message, wsServer, context)
+          );
+          if (typeof handleClose === "function")
+            client.once("close", () => handleClose());
+        }
+      } catch (cause) {
+        logger3.error(
+          `[next-ws] error in socket handler for '${pathname}'`,
+          cause
+        );
+        try {
+          client.close(1011, "Internal Server Error");
+        } catch {
+        }
+      }
     });
+    return;
   });
 }
 // Annotate the CommonJS export names for ESM import in node:

package/dist/server/index.d.ts

@@ -1,32 +1,32 @@
-import NextNodeServer from 'next/dist/server/next-server.js';
-import { Server } from 'node:http';
+import * as next_server from 'next/server';
+import * as http from 'http';
+import * as ws from 'ws';
 import { WebSocketServer } from 'ws';
+import { Server } from 'node:http';
+import NextNodeServer from 'next/dist/server/next-server.js';
 
-declare function setupWebSocketServer(nextServer: NextNodeServer): void;
+/** @deprecated Prefer UPGRADE and {@link UpgradeHandler} */
+type SocketHandler = (client: ws.WebSocket, request: http.IncomingMessage, server: ws.WebSocketServer, context: {
+    params: Record<string, string | string[]>;
+}) => unknown;
+type UpgradeHandler = (client: ws.WebSocket, server: ws.WebSocketServer, request: next_server.NextRequest, context: RouteContext<string>) => unknown;
+
+type RouteParams<Path extends string> = Path extends `${infer Before}[[...${infer Param}]]${infer After}` ? RouteParams<Before> & {
+    [K in Param]?: string[];
+} & RouteParams<After> : Path extends `${infer Before}[...${infer Param}]${infer After}` ? RouteParams<Before> & {
+    [K in Param]: string[];
+} & RouteParams<After> : Path extends `${infer Before}[${infer Param}]${infer After}` ? RouteParams<Before> & {
+    [K in Param]: string;
+} & RouteParams<After> : {};
+type RouteContext<Path extends string> = {
+    params: Record<string, string | string[] | undefined> & RouteParams<Path> & RouteParams<Path>;
+};
 
-/**
- * Set the HTTP server that the WebSocket server should listen on, must be called before the WebSocket server is created.
- * @param server The HTTP server.
- */
-declare function setHttpServer(server: Server): void;
-/**
- * Get the HTTP server that the WebSocket server is listening on.
- * @remark If you want to access the HTTP server outside of a SOCKET handler, you must be using a custom server.
- * @returns The HTTP server.
- * @throws If attempting to access the HTTP server outside of the main process.
- */
-declare function getHttpServer(): Server;
-/**
- * Set the WebSocket server that the WebSocket server should listen on, must be called before the WebSocket server is created.
- * @param wsServer The WebSocket server.
- */
-declare function setWebSocketServer(wsServer: WebSocketServer): void;
-/**
- * Get the WebSocket server that the WebSocket server is listening on.
- * @remark If you want to access the WebSocket server outside of a SOCKET handler, you must be using a custom server.
- * @returns The WebSocket server.
- * @throws If attempting to access the WebSocket server outside of the main process.
- */
-declare function getWebSocketServer(): WebSocketServer;
+declare function getHttpServer(): Server | undefined;
+declare function setHttpServer<T extends Server | undefined>(server: T): T;
+declare function getWebSocketServer(): WebSocketServer | undefined;
+declare function setWebSocketServer<T extends WebSocketServer | undefined>(server: T): T;
+
+declare function setupWebSocketServer(nextServer: NextNodeServer): void;
 
-export { getHttpServer, getWebSocketServer, setHttpServer, setWebSocketServer, setupWebSocketServer };
+export { type RouteContext, type SocketHandler, type UpgradeHandler, getHttpServer, getWebSocketServer, setHttpServer, setWebSocketServer, setupWebSocketServer };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "next-ws",
-  "version": "0.0.0-beta-20250821061718",
+  "version": "0.0.0-beta-20250823112013",
   "description": "Add support for WebSockets in the Next.js app directory",
   "license": "MIT",
   "keywords": [
@@ -38,7 +38,7 @@
   "dependencies": {
     "jscodeshift": "^17.3.0",
     "minimist": "^1.2.8",
-    "semver": "^7.7.1"
+    "semver": "^7.7.2"
   },
   "peerDependencies": {
     "next": ">=13.5.1",
@@ -46,24 +46,24 @@
     "ws": "*"
   },
   "devDependencies": {
-    "@biomejs/biome": "^1.9.4",
-    "@changesets/changelog-git": "^0.2.0",
+    "@biomejs/biome": "^2.2.0",
+    "@changesets/changelog-git": "^0.2.1",
     "@changesets/cli": "^2.27.12",
-    "@playwright/test": "^1.50.1",
+    "@playwright/test": "^1.55.0",
     "@types/jscodeshift": "^17.3.0",
     "@types/minimist": "^1.2.5",
-    "@types/node": "^22.13.1",
-    "@types/react": "19.0.0",
-    "@types/semver": "^7.5.8",
-    "@types/ws": "^8.5.14",
-    "chalk": "^5.4.1",
+    "@types/node": "^24.3.0",
+    "@types/react": "19.1.10",
+    "@types/semver": "^7.7.0",
+    "@types/ws": "^8.18.1",
+    "chalk": "^5.6.0",
     "husky": "^9.1.7",
     "next": "15.5.0",
     "pinst": "^3.0.0",
-    "react": "19.0.0",
-    "react-dom": "19.0.0",
-    "tsup": "^8.3.6",
-    "typescript": "^5.7.3"
+    "react": "19.1.1",
+    "react-dom": "19.1.1",
+    "tsup": "^8.5.0",
+    "typescript": "^5.9.2"
   },
   "scripts": {
     "check": "tsc --noEmit",