feat(server): For dynamic usage, context option can be a function too

README.md

@@ -490,16 +490,18 @@ server.listen(443);
 </details>
 
 <details>
-<summary>Server usage with custom static GraphQL arguments</summary>
+<summary>Server usage with custom context value</summary>
 
 ```typescript
 import { validate, execute, subscribe } from 'graphql';
 import { createServer } from 'graphql-ws';
-import { schema, roots, getStaticContext } from 'my-graphql';
+import { schema, roots, getDynamicContext } from 'my-graphql';
 
 createServer(
   {
-    context: getStaticContext(),
+    context: (ctx, msg, args) => {
+      return getDynamicContext(ctx, msg, args);
+    }, // or static context by supplying the value direcly
     schema,
     roots,
     execute,
@@ -515,12 +517,12 @@ createServer(
 </details>
 
 <details>
-<summary>Server usage with custom dynamic GraphQL arguments and validation</summary>
+<summary>Server usage with custom execution arguments and validation</summary>
 
 ```typescript
 import { parse, validate, execute, subscribe } from 'graphql';
 import { createServer } from 'graphql-ws';
-import { schema, getDynamicContext, myValidationRules } from 'my-graphql';
+import { schema, myValidationRules } from 'my-graphql';
 
 createServer(
   {
@@ -529,7 +531,6 @@ createServer(
     onSubscribe: (ctx, msg) => {
       const args = {
         schema,
-        contextValue: getDynamicContext(ctx, msg),
         operationName: msg.payload.operationName,
         document: parse(msg.payload.query),
         variableValues: msg.payload.variables,

docs/interfaces/_server_.serveroptions.md

@@ -48,15 +48,20 @@ ___
 
 ### context
 
-• `Optional` **context**: unknown
+• `Optional` **context**: [GraphQLExecutionContextValue](../modules/_server_.md#graphqlexecutioncontextvalue) \| (ctx: [Context](_server_.context.md), message: [SubscribeMessage](_message_.subscribemessage.md), args: ExecutionArgs) => [GraphQLExecutionContextValue](../modules/_server_.md#graphqlexecutioncontextvalue)
 
 A value which is provided to every resolver and holds
 important contextual information like the currently
 logged in user, or access to a database.
 
-If you return from the `onSubscribe` callback, this
-context value will NOT be injected. You should add it
-in the returned `ExecutionArgs` from the callback.
+If you return from `onSubscribe`, and the returned value is
+missing the `contextValue` field, this context will be used
+instead.
+
+If you use the function signature, the final execution arguments
+will be passed in (also the returned value from `onSubscribe`).
+Since the context is injected on every subscribe, the `SubscribeMessage`
+with the regular `Context` will be passed in through the arguments too.
 
 ___
 
@@ -199,7 +204,10 @@ If you return `ExecutionArgs` from the callback,
 it will be used instead of trying to build one
 internally. In this case, you are responsible
 for providing a ready set of arguments which will
-be directly plugged in the operation execution.
+be directly plugged in the operation execution. Beware,
+the `context` server option is an exception. Only if you
+dont provide a context alongside the returned value
+here, the `context` server option will be used instead.
 
 To report GraphQL errors simply return an array
 of them from the callback, they will be reported

docs/modules/_server_.md

@@ -14,6 +14,7 @@
 
 ### Type aliases
 
+* [GraphQLExecutionContextValue](_server_.md#graphqlexecutioncontextvalue)
 * [OperationResult](_server_.md#operationresult)
 
 ### Functions
@@ -22,6 +23,19 @@
 
 ## Type aliases
 
+### GraphQLExecutionContextValue
+
+Ƭ  **GraphQLExecutionContextValue**: object \| symbol \| number \| string \| boolean \| null \| undefined
+
+A concrete GraphQL execution context value type.
+
+Mainly used because TypeScript collapes unions
+with `any` or `unknown` to `any` or `unknown`. So,
+we use a custom type to allow definitions such as
+the `context` server option.
+
+___
+
 ### OperationResult
 
 Ƭ  **OperationResult**: Promise\<AsyncIterableIterator\<ExecutionResult> \| ExecutionResult> \| AsyncIterableIterator\<ExecutionResult> \| ExecutionResult

src/server.ts

@@ -43,6 +43,23 @@ export type OperationResult =
   | AsyncIterableIterator<ExecutionResult>
   | ExecutionResult;
 
+/**
+ * A concrete GraphQL execution context value type.
+ *
+ * Mainly used because TypeScript collapes unions
+ * with `any` or `unknown` to `any` or `unknown`. So,
+ * we use a custom type to allow definitions such as
+ * the `context` server option.
+ */
+export type GraphQLExecutionContextValue =
+  // eslint-disable-next-line @typescript-eslint/ban-types
+  | object // you can literally pass "any" JS object as the context value
+  | symbol
+  | number
+  | string
+  | boolean
+  | null;
+
 export interface ServerOptions {
   /**
    * The GraphQL schema on which the operations
@@ -58,11 +75,22 @@ export interface ServerOptions {
    * important contextual information like the currently
    * logged in user, or access to a database.
    *
-   * If you return from the `onSubscribe` callback, this
-   * context value will NOT be injected. You should add it
-   * in the returned `ExecutionArgs` from the callback.
+   * If you return from `onSubscribe`, and the returned value is
+   * missing the `contextValue` field, this context will be used
+   * instead.
+   *
+   * If you use the function signature, the final execution arguments
+   * will be passed in (also the returned value from `onSubscribe`).
+   * Since the context is injected on every subscribe, the `SubscribeMessage`
+   * with the regular `Context` will be passed in through the arguments too.
    */
-  context?: unknown;
+  context?:
+    | GraphQLExecutionContextValue
+    | ((
+        ctx: Context,
+        message: SubscribeMessage,
+        args: ExecutionArgs,
+      ) => GraphQLExecutionContextValue);
   /**
    * The GraphQL root fields or resolvers to go
    * alongside the schema. Learn more about them
@@ -149,7 +177,10 @@ export interface ServerOptions {
    * it will be used instead of trying to build one
    * internally. In this case, you are responsible
    * for providing a ready set of arguments which will
-   * be directly plugged in the operation execution.
+   * be directly plugged in the operation execution. Beware,
+   * the `context` server option is an exception. Only if you
+   * dont provide a context alongside the returned value
+   * here, the `context` server option will be used instead.
    *
    * To report GraphQL errors simply return an array
    * of them from the callback, they will be reported
@@ -538,7 +569,6 @@ export function createServer(
               const { operationName, query, variables } = message.payload;
               const document = typeof query === 'string' ? parse(query) : query;
               execArgs = {
-                contextValue: context,
                 schema,
                 operationName,
                 document,
@@ -569,6 +599,15 @@ export function createServer(
               execArgs.rootValue = roots?.[operationAST.operation];
             }
 
+            // inject the context, if provided, before the operation.
+            // but, only if the `onSubscribe` didnt provide one already
+            if (context !== undefined && !execArgs.contextValue) {
+              execArgs.contextValue =
+                typeof context === 'function'
+                  ? context(ctx, message, execArgs)
+                  : context;
+            }
+
             // the execution arguments have been prepared
             // perform the operation and act accordingly
             let operationResult;

src/tests/server.ts

@@ -318,6 +318,95 @@ it('should pass in the context value from the config', async () => {
   expect(subscribeFn.mock.calls[0][0].contextValue).toBe(context);
 });
 
+it('should pass the `onSubscribe` exec args to the `context` option and use it', async (done) => {
+  const context = {};
+  const execArgs = {
+    // no context here
+    schema,
+    document: parse(`query { getValue }`),
+  };
+
+  const { url } = await startTServer({
+    onSubscribe: () => {
+      return execArgs;
+    },
+    context: (_ctx, _msg, args) => {
+      expect(args).toBe(args); // from `onSubscribe`
+      return context; // will be injected
+    },
+    execute: (args) => {
+      expect(args).toBe(execArgs); // from `onSubscribe`
+      expect(args.contextValue).toBe(context); // injected by `context`
+      done();
+      return execute(args);
+    },
+    subscribe,
+  });
+
+  const client = await createTClient(url);
+  client.ws.send(
+    stringifyMessage<MessageType.ConnectionInit>({
+      type: MessageType.ConnectionInit,
+    }),
+  );
+  await client.waitForMessage(({ data }) => {
+    expect(parseMessage(data).type).toBe(MessageType.ConnectionAck);
+  });
+
+  client.ws.send(
+    stringifyMessage<MessageType.Subscribe>({
+      id: '1',
+      type: MessageType.Subscribe,
+      payload: {
+        query: `{ getValue }`,
+      },
+    }),
+  );
+});
+
+it('should prefer the `onSubscribe` context value even if `context` option is set', async (done) => {
+  const context = 'not-me';
+  const execArgs = {
+    contextValue: 'me-me', // my custom context
+    schema,
+    document: parse(`query { getValue }`),
+  };
+
+  const { url } = await startTServer({
+    onSubscribe: () => {
+      return execArgs;
+    },
+    context, // should be ignored because there is one in `execArgs`
+    execute: (args) => {
+      expect(args).toBe(execArgs); // from `onSubscribe`
+      expect(args.contextValue).not.toBe(context); // from `onSubscribe`
+      done();
+      return execute(args);
+    },
+    subscribe,
+  });
+
+  const client = await createTClient(url);
+  client.ws.send(
+    stringifyMessage<MessageType.ConnectionInit>({
+      type: MessageType.ConnectionInit,
+    }),
+  );
+  await client.waitForMessage(({ data }) => {
+    expect(parseMessage(data).type).toBe(MessageType.ConnectionAck);
+  });
+
+  client.ws.send(
+    stringifyMessage<MessageType.Subscribe>({
+      id: '1',
+      type: MessageType.Subscribe,
+      payload: {
+        query: `{ getValue }`,
+      },
+    }),
+  );
+});
+
 describe('Connect', () => {
   it('should refuse connection and close socket if returning `false`', async () => {
     const { url } = await startTServer({