logo

Practical Implementation of WebSockets with Next.js and tRPC

Cover Simple depiction of Realtime WebSocket Communication. Generated by DALLE

In modern web development, real-time communication between clients and servers has become increasingly important. WebSockets provide a bidirectional communication channel that enables efficient data transfer without the overhead of traditional HTTP requests. In Next.js applications deployed in a serverless environment, using WebSockets directly is not possible due to the stateless and event-driven nature of serverless functions. WebSockets require a long-lived connection between the client and the server, which is incompatible with the serverless architecture.

Typically, to run WebSockets with Next.js, a custom web server needs to be configured alongside the Next.js server. However, this approach has a significant drawback: it involves a separate build step that references the installed node_modules directory. While node_modules contain dependencies required by the application, not all parts are necessary for the runtime execution. Consequently, for containerized and horizontally scalable deployments, using a custom web server significantly increases the size of the container image, consuming more drive space on the server and utilizing more RAM.

This article will demonstrate how to integrate a WebSocket server directly into Next.js and use the integrated app router to create WebSocket endpoints. By following this approach, you can create a tRPC WebSocket connection that allows for typesafe subscriptions and standard WebSockets for bidirectional communication and stateful interaction between the client and server.

One of the key advantages of this approach is that it maintains the optimized performance and resource utilization that Next.js provides out of the box. By avoiding the need for a custom server, developers can ensure a more streamlined and efficient deployment process, aligning with Next.js's design principles.

Throughout the article, you will learn how to seamlessly integrate WebSockets into your Next.js application, enabling real-time communication and synchronization of server state with clients. By leveraging the built-in capabilities of Next.js and tRPC, you can create robust and scalable applications while adhering to best practices and minimizing resource overhead.

Reader should be familiar with the following technologies:

  • Solid understanding of Next.js, a popular React framework for building server-rendered and static websites
  • Proficiency in TypeScript
  • Basic grasp of WebSocket principles and functionality
  • Understanding of tRPC library

WebSockets with Next.js native router

The next-ws package is a game-changer for Next.js applications, enabling developers to create WebSocket endpoints seamlessly within the app router. With the introduction of Next.js 14 and its stable app router, routing WebSocket endpoints has become a reality. This approach eliminates the need for a custom server, aligning with the principles of keeping Next.js projects streamlined and optimized.

The next-ws package seamlessly integrates with the Next.js app router, allowing developers to create WebSocket routes just like any other endpoint. Although a relatively new addition, the next-ws package is rapidly gaining popularity and adoption within the Next.js community, providing a robust and efficient solution for real-time communication.

To get started, you can install the required dependencies with the following command:

npm i next-ws ws

To ensure compatibility with the Next.js app router, the next-ws package needs to modify the Next.js core to support WebSockets. This modification is accomplished using the following patch command:

npx next-ws-cli@latest patch

Creating a WebSocket endpoint follows the same pattern as creating any other app router endpoint. For example, creating a file at src/app/api/fizzbuzz/route.ts would establish a WebSocket endpoint at /api/fizzbuzz.

To differentiate between WebSocket and HTTP endpoints, the route files contain different function names. For HTTP functions, you would use the familiar export default function handler syntax:

import type { NextApiRequest, NextApiResponse } from 'next'

export default function handler(
  req: NextApiRequest,
  res: NextApiResponse
) {
  // HTTP function
}

On the other hand, for WebSocket functions, you would use the export async function SOCKET syntax:

import { WebSocket, WebSocketServer } from 'ws'
import { IncomingMessage } from 'http'

export async function SOCKET(
  client: WebSocket, 
  request: IncomingMessage, 
  wss: WebSocketServer
) {
  // WebSocket function
}

The non-default export with the name SOCKET is crucial for Next.js to differentiate between the intended use of WebSockets and HTTP. The SOCKET function provides direct access to the WebSocket, IncomingMessage, and WebSocketServer instances. The WebSocket client is where most communication occurs, sending and receiving messages, and holding state. The IncomingMessage is also useful when authentication is needed to access the WebSocket. Below is an example of a simple WebSocket authenticated endpoint that holds some state:

export async function SOCKET(client: WebSocket, request: IncomingMessage, server: WebSocketServer) {
  console.log('A client connected!');
  const session = await getSession({ req: request });
  if (!session) {
    console.log('The client is unauthenticated...');
    client.send(JSON.stringify({ error: 'Unauthenticated' }));
    client.close(4001);
    return;
  }
  console.log('The client authenticated successfully!');
  let counter = 0;
  client.on('message', (message) => {
    counter++
    client.send(JSON.stringify({ message: `Received message ${counter} ${message}` }));
    if (counter === 5) {
      client.close(1000);
    }
  });
  client.on('close', () => {
    console.log('A client disconnected!');
  });
}

By leveraging the next-ws package and the Next.js app router, developers can integrate WebSocket functionality seamlessly into their Next.js applications without the need for a custom server, maintaining the optimized performance and resource utilization that Next.js provides out of the box.

tRPC with WebSockets in Next.js

With the ability to create WebSocket endpoints directly within the Next.js app router, the integration of tRPC with WebSockets becomes more streamlined. Previously, implementing tRPC with WebSockets in Next.js applications required the use of a custom server, which goes against the principles of keeping Next.js projects optimized and resource-efficient.

The tRPC package defines the applyWSSHandler function, which wires up the incoming WebSocketServer and creates a connection to the client. Fortunately, the next-ws package provides this connection within the SOCKET function, allowing tRPC to reuse the provided connection seamlessly.

Here's an example implementation in the file /src/app/api/trpcws/route.ts. It's important to note that if a pages router file /src/pages/api/trpc/[trpc].ts already exists, the app router tRPC implementation must use a different URL to avoid conflicts.

import { RawData, WebSocket, WebSocketServer } from 'ws';
import { IncomingMessage } from 'http';
import { createTrpcWsContext } from 'server/api/trpc';
import { appRouter } from 'server/api/root';
import {   
  JSONRPC2,
  parseTRPCMessage,
  TRPCClientOutgoingMessage,
  TRPCResponseMessage
} from '@trpc/server/rpc';
import { isObservable, Unsubscribable } from '@trpc/server/observable';
import { getErrorShape, transformTRPCResponse } from '@trpc/server/shared';
import { TRPCError, callProcedure, getTRPCErrorFromUnknown, inferRouterContext } from '@trpc/server';


export async function SOCKET(client: WebSocket, request: IncomingMessage, wss: WebSocketServer) {
  const { transformer } = appRouter._def._config;
    const clientSubscriptions = new Map<number | string, Unsubscribable>();

    function respond(untransformedJSON: TRPCResponseMessage) {
      client.send(
        JSON.stringify(
          transformTRPCResponse(appRouter._def._config, untransformedJSON),
        ),
      );
    }

    function stopSubscription(
      subscription: Unsubscribable,
      { id, jsonrpc }: JSONRPC2.BaseEnvelope & { id: JSONRPC2.RequestId },
    ) {
      subscription.unsubscribe();

      respond({
        id,
        jsonrpc,
        result: {
          type: 'stopped',
        },
      });
    }

    const ctxPromise = createTrpcWsContext?.({ req: request, res: client });
    let ctx: inferRouterContext<typeof appRouter> | undefined = undefined;

    async function handleRequest(msg: TRPCClientOutgoingMessage) {
      const { id, jsonrpc } = msg;
      /* istanbul ignore next -- @preserve */
      if (id === null) {
        throw new TRPCError({
          code: 'BAD_REQUEST',
          message: '`id` is required',
        });
      }
      if (msg.method === 'subscription.stop') {
        const sub = clientSubscriptions.get(id);
        if (sub) {
          stopSubscription(sub, { id, jsonrpc });
        }
        clientSubscriptions.delete(id);
        return;
      }
      const { path, input } = msg.params;
      const type = msg.method;
      try {
        await ctxPromise; // asserts context has been set

        const result = await callProcedure({
          procedures: appRouter._def.procedures,
          path,
          rawInput: input,
          ctx,
          type,
        });

        if (type === 'subscription') {
          if (!isObservable(result)) {
            throw new TRPCError({
              message: `Subscription ${path} did not return an observable`,
              code: 'INTERNAL_SERVER_ERROR',
            });
          }
        } else {
          // send the value as data if the method is not a subscription
          respond({
            id,
            jsonrpc,
            result: {
              type: 'data',
              data: result,
            },
          });
          return;
        }

        const observable = result;
        const sub = observable.subscribe({
          next(data) {
            respond({
              id,
              jsonrpc,
              result: {
                type: 'data',
                data,
              },
            });
          },
          error(err) {
            const error = getTRPCErrorFromUnknown(err);
            respond({
              id,
              jsonrpc,
              error: getErrorShape({
                config: appRouter._def._config,
                error,
                type,
                path,
                input,
                ctx,
              }),
            });
          },
          complete() {
            respond({
              id,
              jsonrpc,
              result: {
                type: 'stopped',
              },
            });
          },
        });
        /* istanbul ignore next -- @preserve */
        if (client.readyState !== client.OPEN) {
          // if the client got disconnected whilst initializing the subscription
          // no need to send stopped message if the client is disconnected
          sub.unsubscribe();
          return;
        }

        /* istanbul ignore next -- @preserve */
        if (clientSubscriptions.has(id)) {
          // duplicate request ids for client
          stopSubscription(sub, { id, jsonrpc });
          throw new TRPCError({
            message: `Duplicate id ${id}`,
            code: 'BAD_REQUEST',
          });
        }
        clientSubscriptions.set(id, sub);

        respond({
          id,
          jsonrpc,
          result: {
            type: 'started',
          },
        });
      } catch (cause) /* istanbul ignore next -- @preserve */ {
        // procedure threw an error
        const error = getTRPCErrorFromUnknown(cause);
        respond({
          id,
          jsonrpc,
          error: getErrorShape({
            config: appRouter._def._config,
            error,
            type,
            path,
            input,
            ctx,
          }),
        });
      }
    }
    client.on('message', async (message:RawData) => {
      try {
        // eslint-disable-next-line @typescript-eslint/no-base-to-string
        const msgJSON: unknown = JSON.parse(message.toString());
        const msgs: unknown[] = Array.isArray(msgJSON) ? msgJSON : [msgJSON];
        const promises = msgs
          .map((raw) => parseTRPCMessage(raw, transformer as any))
          .map(handleRequest);
        await Promise.all(promises);
      } catch (cause) {
        const error = new TRPCError({
          code: 'PARSE_ERROR',
          cause,
        });

        respond({
          id: null,
          error: getErrorShape({
            config: appRouter._def._config,
            error,
            type: 'unknown',
            path: undefined,
            input: undefined,
            ctx: undefined,
          }),
        });
      }
    });

    // WebSocket errors should be handled, as otherwise unhandled exceptions will crash Node.js.
    // This line was introduced after the following error brought down production systems:
    // "RangeError: Invalid WebSocket frame: RSV2 and RSV3 must be clear"
    // Here is the relevant discussion: https://github.com/websockets/ws/issues/1354#issuecomment-774616962
    client.on('error', (cause) => {
    });

    client.once('close', () => {
      for (const sub of clientSubscriptions.values()) {
        sub.unsubscribe();
      }
      clientSubscriptions.clear();
    });
    async function createContextAsync() {
      try {
        ctx = await ctxPromise;
      } catch (cause) {
        const error = getTRPCErrorFromUnknown(cause);
        respond({
          id: null,
          error: getErrorShape({
            config: appRouter._def._config,
            error,
            type: 'unknown',
            path: undefined,
            input: undefined,
            ctx,
          }),
        });

        // close in next tick
        (global.setImmediate ?? global.setTimeout)(() => {
          client.close();
        });
      }
    }
    await createContextAsync();
}

By utilizing the next-ws package, tRPC can now reuse the WebSocket connection provided within the SOCKET function, eliminating the need for a custom server. This approach ensures that your Next.js application remains optimized and resource-efficient while leveraging the power of tRPC's typesafe subscriptions and WebSocket communication.

tRPC Client Configuration with WebSockets

tRPC can be used on the client-side in two different ways, offering flexibility and catering to different use cases.

The first approach is the higher-level client, createTRPCNext<AppRouter>, which provides an enhanced developer experience by integrating with TanStack React Query hooks. This implementation simplifies the process of fetching and managing data, making it easier to build reactive user interfaces.

Alternatively, there is a lower-level client, createTRPCProxyClient<AppRouter>, which focuses solely on providing direct calls to the server. While it lacks the additional features of the higher-level client, it offers a more lightweight and straightforward implementation.

Regardless of the chosen client implementation, both need to be configured with appropriate ending links. These links can be either the default httpBatchLink for traditional HTTP communication or a wsLink<AppRouter> link for WebSocket communication.

In the context of this article, we'll explore a way to leverage both client implementations simultaneously, catering to complex applications where both approaches might be beneficial. It's important to note that using both clients will result in opening and maintaining separate WebSocket connections throughout the lifetime of the client-side application.

Here's an example of how you can configure both clients:

import { httpBatchLink, loggerLink, createTRPCProxyClient, wsLink, createWSClient } from "@trpc/client";
import { createTRPCNext } from "@trpc/next";
import superjson from "superjson";
import { type AppRouter } from "../server/api/root";
import { NextPageContext } from "next";

const getBaseUrl = () => {
  if (typeof window !== "undefined") return ""; // browser should use relative url
  return "http://localhost:3000"; // dev SSR should use localhost
};

const getWsLink = () => {
  let socketUrl = "ws://localhost:3000/api/trpcws"
  if (typeof window !== "undefined") {
    const protocol = window.location.protocol === 'https:' ? 'wss' : 'ws';
    socketUrl = `${protocol}://${window.location.host}/api/trpcws`;
  }
  const wslink = wsLink<AppRouter>({
    client: createWSClient({
      url: socketUrl,
    })
  });
  return wslink
}

function getEndingLink(ctx: NextPageContext | undefined) {
  if (typeof window === 'undefined') {
    return httpBatchLink({
      url: `${getBaseUrl()}/api/trpc`,
      headers() {
        if (!ctx?.req?.headers) {
          return {};
        }
        return {
          ...ctx.req.headers,
          'x-ssr': '1',
        };
      },
    });
  }
  return getWsLink();
}

export const api = createTRPCNext<AppRouter>({
  config({ ctx }) {
    return {
      transformer: superjson,
      links: [
        loggerLink({
          enabled: (opts) =>
            process.env.NODE_ENV === "development" ||
            (opts.direction === "down" && opts.result instanceof Error),
        }),
        getEndingLink(ctx)
      ],
      queryClientConfig: {
        defaultOptions: {
          queries: {
            retry: false,
          },
        },
      },
    };
  },
  ssr: false,
});


export const trpcClient = createTRPCProxyClient<AppRouter>({
  links: [
    loggerLink({
      enabled: (opts) =>
        process.env.NODE_ENV === "development" ||
        (opts.direction === "down" && opts.result instanceof Error),
    }),
    getEndingLink(undefined)
  ],
  transformer: superjson,
});

By providing both client implementations, developers can choose the approach that best suits their specific use case within the same application. The higher-level api client leverages React Query hooks for efficient data management, while the lower-level trpcClient offers a more direct connection to the server.

It's worth noting that maintaining multiple WebSocket connections can introduce additional overhead and complexity. Therefore, it's essential to carefully evaluate the trade-offs and only implement both clients if the application's requirements truly justify the additional complexity.

Using tRPC Subscriptions with WebSockets

Building upon the previous sections, we'll explore how to define tRPC subscriptions on the server-side and consume them from the client-side. tRPC subscriptions follow a familiar syntax to queries and mutations, with a notable exception: subscriptions must return an observable. The observable maintains a stateful connection until it is closed, disconnected from the client-side, or a complete() message is emitted.

Here's an example of a subscription endpoint that sends 10 FizzBuzz iterations from a given starting point provided in the input over the following 10 seconds:

import { z } from 'zod';
import { createTRPCRouter, publicProcedure } from '../trpc';
import { observable } from '@trpc/server/observable';

export const fizzbuzzRouter = createTRPCRouter({
    fizzbuzzSubcription: publicProcedure
        .input(z.number())
        .subscription(async ({ input }) => {
            return observable((emit) => {
                let counter = 0
                const fizzbuzz = () => {
                    setTimeout(() => {
                        const number = counter + input
                        let output = ''
                        if (number % 3 === 0) {
                            output += 'Fizz'
                        }
                        if (number % 5 === 0) {
                            output += 'Buzz'
                        }
                        emit.next({ number, output });
                        counter++
                        if (counter < 10) {
                            fizzbuzz();
                        } else {
                            emit.complete()
                        }
                    }, 1000)
                }
                fizzbuzz();
                // Handle WebSocket closure and errors as needed
                return () => {
                    // Cleanup logic if necessary
                };
            });
        })
})

In this example, the fizzbuzzSubscription endpoint is defined using the publicProcedure and .subscription methods from tRPC. The input is validated using zod, ensuring that a number is provided. The subscription returns an observable that emits the FizzBuzz iterations over a period of 10 seconds.

To subscribe to this endpoint from the client-side and receive the 10 FizzBuzz iterations, you can use the useSubscription hook provided by tRPC:

    api.fizzbuzzRouter.fizzbuzzSubcription.useSubscription(15, {
        onData: (data: any) => {
            console.log(data)
        },
        onError: (error) => {
            console.log(error)
        }
    })

In this example, the useSubscription hook is called with the starting point (15) as the input, and two callback functions are provided: onData and onError. The onData callback will be invoked for each iteration, logging the data to the console. The onError callback will be invoked if an error occurs during the subscription.

Conclusion

By leveraging tRPC's subscription functionality in conjunction with WebSockets, you can establish real-time, stateful connections between the client and server, enabling efficient data streaming and synchronization. This approach provides a powerful and type-safe way to implement real-time features in your Next.js applications.

It's important to note that not only subscriptions but also other tRPC queries and mutations are sent through the same WebSocket connection. This behavior is a result of the ending link configuration, which sends all tRPC traffic through the WebSocket link instead of splitting it across HTTP and WebSocket connections.

Regardless of whether there are active subscriptions or not, the tRPC client maintains an always-on WebSocket connection to the server. This approach has several advantages over traditional HTTP connections.

While WebSockets introduce an initial overhead for establishing the connection, they excel in scenarios where frequent communication between the client and server is required. Unlike HTTP connections, which have higher latency due to the need for handshakes on every request, WebSocket connections enable more efficient data transfer by eliminating this overhead.

By maintaining a persistent WebSocket connection, subsequent requests and responses can be transmitted without the need for additional handshakes, resulting in lower latency and improved responsiveness. This is particularly beneficial for real-time applications or those that require frequent updates, as data can be streamed seamlessly without the added overhead of establishing new connections for each request.

Furthermore, WebSocket connections enable bidirectional communication, allowing both the client and server to initiate data transfer as needed. This flexibility is advantageous for scenarios where the server needs to push updates to the client without the client explicitly requesting them, such as in real-time notifications or live updates.

While the initial setup and configuration of WebSocket connections may be more involved compared to traditional HTTP connections, the benefits of reduced latency, efficient data transfer, and bidirectional communication often outweigh the overhead, especially in applications that require real-time or frequent data exchange.

By leveraging the tRPC library and the WebSocket capabilities provided by the next-ws package, developers can take advantage of these performance optimizations while maintaining a consistent and type-safe API, streamlining the development process and improving the overall user experience.