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

if anyone happens to find this, here is the setup for graphql-yoga v5, with bun

import { makeHandler } from 'graphql-ws/use/bun';
import { createYoga } from 'graphql-yoga';


const yogaApp = createYoga({ ... config here... })

  const bunServer = Bun.serve({
    port: PORT,
    hostname: '0.0.0.0',
    websocket: makeHandler<{ token: string }, { user: UserSession.DefaultTokenPayload }>({
      // extra is available here so we don't have to revalidate the JWT
      context: setWebhookContext,
      // return false or throw to prevent connection
      async onConnect({ connectionParams, extra }) {
        const token = connectionParams?.token;

        if (typeof token !== 'string') return false;

        try {
          extra.user = authenticateJWT(token);
          return true;
        } catch {
          return false;
        }
      },
      schema: executableSchema,
    }),
    async fetch(request, server) {
      const url = new URL(request.url);

      switch (url.pathname) {
        case '/health': {
          return new Response('OK', { status: 200 });
        }
        case '/subscriptions': {
          if (server.upgrade(request)) {
            return new Response('OK', { status: 200 });
          }
          return new Response('Unable to upgrade to websockets', { status: 400 });
        }
        case yogaApp.graphqlEndpoint: {
          return yogaApp.fetch(request, server);
        }
       ... etc
      }
    },
  });