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}`
    )}`
)