Updating/mutating a session with next-auth and tRPC

I am building a multi-tenant NextJS app that uses next-auth for account authentication, tRPC for API's, and postgresql for a data store.

I am trying to find a way to dynamically update/set/mutate a session value based on some client-side interaction

The approach I am taking is similar to the one described in this article:

a User is granted access to an Organization through a Membership
a User may have a Membership to >1 Organization
a User can change which Organization they are "logged in" to through some client-side UI.

When the user authenticates, I want to:

set session.user.orgId to some orgId (if they belong to an org)

When the user changes the org they are accessing through some client-side UI, I want to:

update session.user.orgId = newOrgId (validating they have proper permissions before doing so, of course).

I have searched the net for ways to update/mutate session values, and as far as I can tell, it's only possible using next-auth's callbacks:

...
  callbacks: {
    async session({ session, user, token }) {
      // we can modify session here, i.e `session.orgId = 'blah'`
      // or look up a value in the db and attach it here.
      return session
    },
...
}

However, there is no clear way to trigger this update from the client, outside of the authentication flow. I.E, if the user clicks to change their org in some UI, how do I validate the change + update the session value, without requiring the user to re-authenticate?

Solution

Hack into NextAuth's PrismaAdapter and trpc's createContext.

For example:

File: src/pages/api/auth/[...nextauth].ts

import NextAuth, { Awaitable, type NextAuthOptions } from "next-auth";
import { PrismaAdapter } from "@next-auth/prisma-adapter";
import type { AdapterSession, AdapterUser } from "next-auth/adapters";
import { prisma } from "../../../server/db/client";
import { MembershipRole } from "@prisma/client";

...

const adapter = PrismaAdapter(prisma);
adapter.createSession = (session: {
  sessionToken: string;
  userId: string;
  expires: Date;
}): Awaitable<AdapterSession> => {
  return prisma.user
    .findUniqueOrThrow({
      where: {
        id: session.userId,
      },
      select: {
        memberships: {
          where: {
            isActiveOrg: true,
          },
          select: {
            role: true,
            organization: true,
          },
        },
      },
    })
    .then((userWithOrg) => {
      const membership = userWithOrg.memberships[0];
      const orgId = membership?.organization.id;

      return prisma.session.create({
        data: {
          expires: session.expires,
          sessionToken: session.sessionToken,
          user: {
            connect: { id: session.userId },
          },
          organization: {
            connect: {
              id: orgId,
            },
          },
          role: membership?.role as MembershipRole,
        },
      });
    });
};


// the authOptions to user with NextAuth
export const authOptions: NextAuthOptions = {
  // Include user.id on session
  callbacks: {
    session({ session, user }) {
      if (session.user) {
        session.user.id = user.id;
      }
      return session;
    },
  },

  adapter: adapter,

  // Configure one or more authentication providers
  providers: [
    ...
  ],
};

export default NextAuth(authOptions);

File: src/server/trpc/context.ts

import type { inferAsyncReturnType } from "@trpc/server";
import type { CreateNextContextOptions } from "@trpc/server/adapters/next";
import type { Session } from "next-auth";
import { getServerAuthSession } from "../common/get-server-auth-session";
import { prisma } from "../db/client";

type CreateContextOptions = {
  session: Session | null;
};

/** Use this helper for:
 * - testing, so we dont have to mock Next.js' req/res
 * - trpc's `createSSGHelpers` where we don't have req/res
 **/
export const createContextInner = async (opts: CreateContextOptions) => {
  return {
    session: opts.session,
    prisma,
  };
};

/**
 * This is the actual context you'll use in your router
 * @link https://trpc.io/docs/context
 **/
export const createContext = async (opts: CreateNextContextOptions) => {
  const { req, res } = opts;

  const sessionToken = req.cookies["next-auth.session-token"];
  const prismaSession = await prisma.session.findUniqueOrThrow({
    where: {
      sessionToken: sessionToken,
    },
    select: {
      orgId: true,
    },
  });
  const orgId = prismaSession.orgId;

  // Get the session from the server using the unstable_getServerSession wrapper function
  const session = (await getServerAuthSession({ req, res })) as Session;
  const sessionWithOrg = {
    session: {
      user: {
        // need this otherwise createContextInner doesn't accept for a possible null session.user.id
        id: session?.user?.id || "",
        orgId: orgId,
        ...session?.user,
      },
      expires: session?.expires,
    },
  };
  const context = await createContextInner(sessionWithOrg);

  return context;
};

export type Context = inferAsyncReturnType<typeof createContext>;