Implementing GraphQL Subscriptions with Websockets in a Bun Server using graphql-yoga

server.ts

import Bun from 'bun'
import { createYoga, YogaInitialContext, YogaServerInstance } from 'graphql-yoga'
import { makeHandler } from "graphql-ws/lib/use/bun";
import { ExecutionArgs } from "@envelop/types";

import { schema } from './graphql/schema';

interface IUserContext {
    token?: string;
}

const PORT = process.env.PORT || 4000;

const yoga: YogaServerInstance<{}, YogaInitialContext> = createYoga({
    graphqlEndpoint: '/graphql',
    schema,
    graphiql: {
        subscriptionsProtocol: 'WS',
    },
    context: async ({ request }): Promise<IUserContext> => {
          return {
              token: request.headers.get("token") ?? ""
          }
    }
})

const websocketHandler = makeHandler({
    schema,
    execute: (args: ExecutionArgs) => args.rootValue.execute(args),
    subscribe: (args: ExecutionArgs) => args.rootValue.subscribe(args),
    onSubscribe: async (ctx, msg) => {
        const {schema, execute, subscribe, contextFactory, parse, validate} = yoga.getEnveloped({
            ...ctx,
            req: ctx.extra.request,
            socket: ctx.extra.socket,
            params: msg.payload
        })

        const args = {
            schema,
            operationName: msg.payload.operationName,
            document: parse(msg.payload.query),
            variableValues: msg.payload.variables,
            contextValue: await contextFactory(),
            rootValue: {
                execute,
                subscribe
            }
        }

        const errors = validate(args.schema, args.document)
        if (errors.length) return errors
        return args
    },
})

const server: Bun.Server = Bun.serve({
    fetch: (request: Request, server: Bun.Server): Promise<Response> | Response => {
        // Upgrade the request to a WebSocket
        if (server.upgrade(request)) {
            return new Response()
        }
        return yoga.fetch(request, server)
    },
    port: PORT,
    websocket: websocketHandler,
})

console.info(
    `🚀 Server is running on ${new URL(
        yoga.graphqlEndpoint,
        `http://${server.hostname}:${server.port}`
    )}`
)

For those who faced issues with accessing the passed headers for custom context providing, like authorisation header, here's a useful note on the Authorization token example:

export function getContextJWTToken(context: AppQraphQLContext): string | null {
    if (context.params.extensions?.headers?.Authorization) {
        return getJWTFromStr(context.params.extensions.headers.Authorization);
    }
    if (context.request) {
        const request = context.request;
        const authHeader = request.headers.get('Authorization');
        if (authHeader) return getJWTFromStr(authHeader);
    }

    return null;
}

The typical context.request headers are not present in WS connection but appear in context.params.extensions?.headers

Full example in context of global user context pass:

type AppQraphQLContext = YogaInitialContext & {
  user: User | null;
};
function getContextJWTToken(context: AppQraphQLContext): string | null {
    if (context.params.extensions?.headers?.Authorization) {
        return getJWTFromStr(context.params.extensions.headers.Authorization);
    }
    if (context.request) {
        const request = context.request;
        const authHeader = request.headers.get('Authorization');
        if (authHeader) return getJWTFromStr(authHeader);
    }

    return null;
}
const getContextUser = async (context: AppQraphQLContext) => {
    let token = getContextJWTToken(context);
    try {
        const decoded: any = jwt.verify(token || '', process.env.JWT_SECRET);
        const user = await findUser({ id: decoded.userId });

        return user;
    }
    catch (err) {
        return null;
    }
};

const yoga = createYoga<AppQraphQLContext>({
    schema,
    graphiql: {
        subscriptionsProtocol: 'WS',
    },
    context: async (context: AppQraphQLContext) => {
        return {
            ...context,
            user: await getContextUser(context)
        }
    },
    landingPage: false,
    graphqlEndpoint: '/',
});