Jotai trpc client for app router

client.ts

import { httpBatchLink } from "@trpc/client";
import superjson from "superjson";
import { createTRPCJotaiClient } from "jotai-next-client";
import type { AppRouter } from "@/server/router";

const getBaseUrl = () => {
  if (process.env.VERCEL_URL) return `https://${process.env.VERCEL_URL}`;
  return `http://localhost:${process.env.PORT ?? 3000}`;
};

export const trpc = createTRPCJotaiClient<AppRouter>({
  transformer: superjson,
  links: [
    httpBatchLink({
      url: `${getBaseUrl()}/api/trpc`,
    }),
  ],
});

jotai-next-client.ts

import { createRecursiveProxy } from "@trpc/server/shared";
import {
  AnyMutationProcedure,
  AnyProcedure,
  AnyQueryProcedure,
  AnyRouter,
  ProcedureArgs,
  ProcedureRouterRecord,
  inferProcedureOutput,
} from "@trpc/server";
import { atom, Getter, WritableAtom } from "jotai";
import { CreateTRPCClientOptions, createTRPCProxyClient } from "@trpc/client";

// Simplified types from original jotai-trpc
type AsyncValueOrGetter<T> =
  | T
  | Promise<T>
  | ((get: Getter) => T)
  | ((get: Getter) => Promise<T>);

type QueryResolver<TProcedure extends AnyProcedure> = {
  (
    getInput: AsyncValueOrGetter<ProcedureArgs<TProcedure["_def"]>[0] | null>,
  ): WritableAtom<
    Promise<inferProcedureOutput<TProcedure> | undefined>,
    [],
    void
  >;
};

type MutationResolver<TProcedure extends AnyProcedure> = () => WritableAtom<
  inferProcedureOutput<TProcedure> | null,
  ProcedureArgs<TProcedure["_def"]>,
  Promise<inferProcedureOutput<TProcedure> | undefined>
>;

type DecorateProcedure<TProcedure extends AnyProcedure> =
  TProcedure extends AnyQueryProcedure
    ? {
        atomWithQuery: QueryResolver<TProcedure>;
      }
    : TProcedure extends AnyMutationProcedure
      ? {
          atomWithMutation: MutationResolver<TProcedure>;
        }
      : never;

type DecoratedProcedureRecord<TProcedures extends ProcedureRouterRecord> = {
  [TKey in keyof TProcedures]: TProcedures[TKey] extends AnyRouter
    ? DecoratedProcedureRecord<TProcedures[TKey]["_def"]["record"]>
    : TProcedures[TKey] extends AnyProcedure
      ? DecorateProcedure<TProcedures[TKey]>
      : never;
};

// Helper functions
const isGetter = <T>(v: T | ((get: Getter) => T)): v is (get: Getter) => T =>
  typeof v === "function";

const getProcedure = (obj: any, path: string[]) => {
  for (let i = 0; i < path.length; ++i) {
    obj = obj[path[i] as string];
  }
  return obj;
};
  
// This atom controls if any requests are done. Should be set to true when on client.
export const enableClientAtom = atom(false);

// Client implementation
export const createTRPCJotaiClient = <TRouter extends AnyRouter>(
  opts: CreateTRPCClientOptions<TRouter>,
) => {
  const client = createTRPCProxyClient<TRouter>(opts);

  return createRecursiveProxy(
    ({ path, args }: { path: string[]; args: unknown[] }) => {
      const type = path.pop() as "atomWithQuery" | "atomWithMutation";

      if (type === "atomWithQuery") {
        const [getInput] = args;
        const refreshAtom = atom(0);
        const queryAtom = atom(
          (get, { signal }) => {
            get(refreshAtom);

            const enabled = get(enableClientAtom);
            if (!enabled) {
              // We're probably on server, return early
              return;
            }
  
            const procedure = getProcedure(client, path);
            return Promise.resolve(
              isGetter(getInput) ? getInput(get) : getInput,
            ).then((input) =>
              input !== null ? procedure.query(input, { signal }) : undefined,
            );
          },
          (_, set) => set(refreshAtom, (counter) => counter + 1),
        );
        return queryAtom;
      } else if (type === "atomWithMutation") {
        const mutationAtom = atom(null, (get, set, ...args: unknown[]) => {
          const enabled = get(enableClientAtom);
          if (enabled) {
            // We're probably on server, return early
            return;
          }
  
          const procedure = getProcedure(client, path);
          return procedure.mutate(...args).then((result: unknown) => {
            set(mutationAtom, result);
            return result;
          });
        });
        return mutationAtom;
      }
    },
  ) as DecoratedProcedureRecord<TRouter["_def"]["record"]>;
};

jotai-next-provider.tsx

"use client";

import { useAtom } from "jotai";
import { useEffect } from "react";
import { enableClientAtom } from "./jotai-next-client";

export function Provider({ children }: { children: React.ReactNode }) {
  const [_, setEnabled] = useAtom(enableClientAtom);
  useEffect(() => {
    setEnabled(true);
  }, []);
  return children;
}

layout.tsx

import { Provider as JotaiProvider } from "./jotai-next-provider";

function RootLayout({ children }: { children: React.ReactNode }) {
  return (
    <html lang="en">
      <body>
        <JotaiProvider>
          <main>
            {children}
          </main>
        </JotaiProvider>
      </body>
    </html>
  );
}