diff --git a/examples/build.gradle.kts b/examples/build.gradle.kts
index 84a9b53a22..0f58345f83 100644
--- a/examples/build.gradle.kts
+++ b/examples/build.gradle.kts
@@ -40,6 +40,7 @@ subprojects {
 
     val kotlinVersion: String by project
     val junitVersion: String by project
+    val kotlinCoroutinesVersion: String by project
 
     val detektVersion: String by project
     val ktlintVersion: String by project
@@ -50,6 +51,8 @@ subprojects {
 
     dependencies {
         implementation(kotlin("stdlib", kotlinVersion))
+        implementation(kotlin("reflect", kotlinVersion))
+        implementation("org.jetbrains.kotlinx:kotlinx-coroutines-jdk8:$kotlinCoroutinesVersion")
         testImplementation(kotlin("test-junit5", kotlinVersion))
         testImplementation("org.junit.jupiter:junit-jupiter-api:$junitVersion")
         testImplementation("org.junit.jupiter:junit-jupiter-engine:$junitVersion")
diff --git a/examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/KtorGraphQLContextFactory.kt b/examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/KtorGraphQLContextFactory.kt
index 54e2ef485b..79f1ddb9d8 100644
--- a/examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/KtorGraphQLContextFactory.kt
+++ b/examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/KtorGraphQLContextFactory.kt
@@ -25,7 +25,7 @@ import io.ktor.request.ApplicationRequest
  */
 class KtorGraphQLContextFactory : GraphQLContextFactory<AuthorizedContext, ApplicationRequest> {
 
-    override fun generateContext(request: ApplicationRequest): AuthorizedContext {
+    override suspend fun generateContext(request: ApplicationRequest): AuthorizedContext {
         val loggedInUser = User(
             email = "fake@site.com",
             firstName = "Someone",
diff --git a/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/context/MyGraphQLContext.kt b/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/context/MyGraphQLContext.kt
index 5af61e650f..daa5413bcf 100644
--- a/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/context/MyGraphQLContext.kt
+++ b/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/context/MyGraphQLContext.kt
@@ -24,8 +24,8 @@ import org.springframework.web.reactive.socket.WebSocketSession
  * Simple [GraphQLContext] that holds extra value and the [ServerRequest]
  */
 class MyGraphQLContext(
-    val myCustomValue: String,
-    val request: ServerRequest
+    val request: ServerRequest,
+    val myCustomValue: String
 ) : GraphQLContext
 
 /**
@@ -33,5 +33,5 @@ class MyGraphQLContext(
  */
 class MySubscriptionGraphQLContext(
     val request: WebSocketSession,
-    var subscriptionValue: String? = null
+    var auth: String? = null
 ) : GraphQLContext
diff --git a/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/context/MyGraphQLContextFactory.kt b/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/context/MyGraphQLContextFactory.kt
index 0ce55473d7..513d30fbdf 100644
--- a/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/context/MyGraphQLContextFactory.kt
+++ b/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/context/MyGraphQLContextFactory.kt
@@ -29,17 +29,20 @@ import org.springframework.web.reactive.socket.WebSocketSession
 @Component
 class MyGraphQLContextFactory : SpringGraphQLContextFactory<MyGraphQLContext>() {
 
-    override fun generateContext(request: ServerRequest): MyGraphQLContext = MyGraphQLContext(
-        myCustomValue = request.headers().firstHeader("MyHeader") ?: "defaultContext",
-        request = request
+    override suspend fun generateContext(request: ServerRequest): MyGraphQLContext = MyGraphQLContext(
+        request = request,
+        myCustomValue = request.headers().firstHeader("MyHeader") ?: "defaultContext"
     )
 }
 
+/**
+ * [GraphQLContextFactory] that generates [MySubscriptionGraphQLContext] that will be available when processing subscription operations.
+ */
 @Component
 class MySubscriptionGraphQLContextFactory : SpringSubscriptionGraphQLContextFactory<MySubscriptionGraphQLContext>() {
 
-    override fun generateContext(request: WebSocketSession): MySubscriptionGraphQLContext = MySubscriptionGraphQLContext(
+    override suspend fun generateContext(request: WebSocketSession): MySubscriptionGraphQLContext = MySubscriptionGraphQLContext(
         request = request,
-        subscriptionValue = null
+        auth = null
     )
 }
diff --git a/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/execution/MySubscriptionHooks.kt b/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/execution/MySubscriptionHooks.kt
index c883c21a32..52a5c2a28b 100644
--- a/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/execution/MySubscriptionHooks.kt
+++ b/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/execution/MySubscriptionHooks.kt
@@ -19,9 +19,7 @@ package com.expediagroup.graphql.examples.server.spring.execution
 import com.expediagroup.graphql.examples.server.spring.context.MySubscriptionGraphQLContext
 import com.expediagroup.graphql.generator.execution.GraphQLContext
 import com.expediagroup.graphql.server.spring.subscriptions.ApolloSubscriptionHooks
-import kotlinx.coroutines.reactor.mono
 import org.springframework.web.reactive.socket.WebSocketSession
-import reactor.core.publisher.Mono
 
 /**
  * A simple implementation of Apollo Subscription Lifecycle Events.
@@ -31,12 +29,11 @@ class MySubscriptionHooks : ApolloSubscriptionHooks {
     override fun onConnect(
         connectionParams: Map<String, String>,
         session: WebSocketSession,
-        graphQLContext: GraphQLContext
-    ): Mono<GraphQLContext> = mono {
-        if (graphQLContext is MySubscriptionGraphQLContext) {
-            val bearer = connectionParams["Authorization"] ?: "none"
-            graphQLContext.subscriptionValue = bearer
+        graphQLContext: GraphQLContext?
+    ): GraphQLContext? {
+        if (graphQLContext != null && graphQLContext is MySubscriptionGraphQLContext) {
+            graphQLContext.auth = connectionParams["Authorization"]
         }
-        graphQLContext
+        return graphQLContext
     }
 }
diff --git a/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/subscriptions/SimpleSubscription.kt b/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/subscriptions/SimpleSubscription.kt
index 24fdbf597f..0e32fa2a31 100644
--- a/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/subscriptions/SimpleSubscription.kt
+++ b/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/subscriptions/SimpleSubscription.kt
@@ -82,5 +82,5 @@ class SimpleSubscription : Subscription {
 
     @GraphQLDescription("Returns a value from the subscription context")
     fun subscriptionContext(myGraphQLContext: MySubscriptionGraphQLContext): Flux<String> =
-        Flux.just(myGraphQLContext.subscriptionValue ?: "", "value 2", "value3")
+        Flux.just(myGraphQLContext.auth ?: "no-auth")
 }
diff --git a/generator/graphql-kotlin-schema-generator/src/main/kotlin/com/expediagroup/graphql/generator/execution/FunctionDataFetcher.kt b/generator/graphql-kotlin-schema-generator/src/main/kotlin/com/expediagroup/graphql/generator/execution/FunctionDataFetcher.kt
index 9f26351533..775e1d1508 100644
--- a/generator/graphql-kotlin-schema-generator/src/main/kotlin/com/expediagroup/graphql/generator/execution/FunctionDataFetcher.kt
+++ b/generator/graphql-kotlin-schema-generator/src/main/kotlin/com/expediagroup/graphql/generator/execution/FunctionDataFetcher.kt
@@ -59,7 +59,7 @@ open class FunctionDataFetcher(
      * Invoke a suspend function or blocking function, passing in the [target] if not null or default to using the source from the environment.
      */
     override fun get(environment: DataFetchingEnvironment): Any? {
-        val instance = target ?: environment.getSource<Any?>()
+        val instance: Any? = target ?: environment.getSource<Any?>()
         val instanceParameter = fn.instanceParameter
 
         return if (instance != null && instanceParameter != null) {
diff --git a/generator/graphql-kotlin-schema-generator/src/main/kotlin/com/expediagroup/graphql/generator/execution/GraphQLContext.kt b/generator/graphql-kotlin-schema-generator/src/main/kotlin/com/expediagroup/graphql/generator/execution/GraphQLContext.kt
index a64acdc57e..5cd5336891 100644
--- a/generator/graphql-kotlin-schema-generator/src/main/kotlin/com/expediagroup/graphql/generator/execution/GraphQLContext.kt
+++ b/generator/graphql-kotlin-schema-generator/src/main/kotlin/com/expediagroup/graphql/generator/execution/GraphQLContext.kt
@@ -16,19 +16,8 @@
 
 package com.expediagroup.graphql.generator.execution
 
-import java.util.concurrent.ConcurrentHashMap
-import java.util.concurrent.ConcurrentMap
-
 /**
  * Marker interface to indicate that the implementing class should be considered
  * as the GraphQL context. This means the implementing class will not appear in the schema.
  */
 interface GraphQLContext
-
-/**
- * Default [GraphQLContext] that can be used if there is none provided. Exposes generic concurrent hash map
- * that can be populated with custom data.
- */
-class DefaultGraphQLContext : GraphQLContext {
-    val contents: ConcurrentMap<Any, Any> = ConcurrentHashMap()
-}
diff --git a/servers/graphql-kotlin-server/src/main/kotlin/com/expediagroup/graphql/server/execution/GraphQLContextFactory.kt b/servers/graphql-kotlin-server/src/main/kotlin/com/expediagroup/graphql/server/execution/GraphQLContextFactory.kt
index 34e5d4b5ad..8251d2d7d5 100644
--- a/servers/graphql-kotlin-server/src/main/kotlin/com/expediagroup/graphql/server/execution/GraphQLContextFactory.kt
+++ b/servers/graphql-kotlin-server/src/main/kotlin/com/expediagroup/graphql/server/execution/GraphQLContextFactory.kt
@@ -16,7 +16,6 @@
 
 package com.expediagroup.graphql.server.execution
 
-import com.expediagroup.graphql.generator.execution.DefaultGraphQLContext
 import com.expediagroup.graphql.generator.execution.GraphQLContext
 
 /**
@@ -28,12 +27,5 @@ interface GraphQLContextFactory<out Context : GraphQLContext, Request> {
      * Generate GraphQL context based on the incoming request and the corresponding response.
      * If no context should be generated and used in the request, return null.
      */
-    fun generateContext(request: Request): Context?
-}
-
-/**
- * Default context factory that generates GraphQL context with empty concurrent map that can store any elements.
- */
-class DefaultGraphQLContextFactory<T> : GraphQLContextFactory<DefaultGraphQLContext, T> {
-    override fun generateContext(request: T) = DefaultGraphQLContext()
+    suspend fun generateContext(request: Request): Context?
 }
diff --git a/servers/graphql-kotlin-server/src/main/kotlin/com/expediagroup/graphql/server/extensions/requestExtensions.kt b/servers/graphql-kotlin-server/src/main/kotlin/com/expediagroup/graphql/server/extensions/requestExtensions.kt
index c7bf93eb6e..3994401f57 100644
--- a/servers/graphql-kotlin-server/src/main/kotlin/com/expediagroup/graphql/server/extensions/requestExtensions.kt
+++ b/servers/graphql-kotlin-server/src/main/kotlin/com/expediagroup/graphql/server/extensions/requestExtensions.kt
@@ -16,7 +16,6 @@
 
 package com.expediagroup.graphql.server.extensions
 
-import com.expediagroup.graphql.generator.execution.DefaultGraphQLContext
 import com.expediagroup.graphql.types.GraphQLRequest
 import graphql.ExecutionInput
 import org.dataloader.DataLoaderRegistry
@@ -29,6 +28,6 @@ fun GraphQLRequest.toExecutionInput(graphQLContext: Any? = null, dataLoaderRegis
         .query(this.query)
         .operationName(this.operationName)
         .variables(this.variables ?: emptyMap())
-        .context(graphQLContext ?: DefaultGraphQLContext())
+        .context(graphQLContext)
         .dataLoaderRegistry(dataLoaderRegistry ?: DataLoaderRegistry())
         .build()
diff --git a/servers/graphql-kotlin-server/src/test/kotlin/com/expediagroup/graphql/server/extensions/RequestExtensionsKtTest.kt b/servers/graphql-kotlin-server/src/test/kotlin/com/expediagroup/graphql/server/extensions/RequestExtensionsKtTest.kt
index 55247719ad..74e8615838 100644
--- a/servers/graphql-kotlin-server/src/test/kotlin/com/expediagroup/graphql/server/extensions/RequestExtensionsKtTest.kt
+++ b/servers/graphql-kotlin-server/src/test/kotlin/com/expediagroup/graphql/server/extensions/RequestExtensionsKtTest.kt
@@ -16,14 +16,13 @@
 
 package com.expediagroup.graphql.server.extensions
 
-import com.expediagroup.graphql.generator.execution.GraphQLContext
 import com.expediagroup.graphql.types.GraphQLRequest
 import io.mockk.mockk
 import org.dataloader.DataLoaderRegistry
 import org.junit.jupiter.api.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertNotNull
-import kotlin.test.assertTrue
+import kotlin.test.assertNull
 
 class RequestExtensionsKtTest {
 
@@ -32,7 +31,7 @@ class RequestExtensionsKtTest {
         val request = GraphQLRequest(query = "query { whatever }")
         val executionInput = request.toExecutionInput()
         assertEquals(request.query, executionInput.query)
-        assertTrue(executionInput.context is GraphQLContext)
+        assertNull(executionInput.context)
         assertNotNull(executionInput.dataLoaderRegistry)
     }
 
diff --git a/servers/graphql-kotlin-spring-server/build.gradle.kts b/servers/graphql-kotlin-spring-server/build.gradle.kts
index 5d33dddc21..5346a92dda 100644
--- a/servers/graphql-kotlin-spring-server/build.gradle.kts
+++ b/servers/graphql-kotlin-spring-server/build.gradle.kts
@@ -35,7 +35,7 @@ tasks {
                 limit {
                     counter = "BRANCH"
                     value = "COVEREDRATIO"
-                    minimum = "0.74".toBigDecimal()
+                    minimum = "0.70".toBigDecimal()
                 }
             }
         }
diff --git a/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/SubscriptionAutoConfiguration.kt b/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/SubscriptionAutoConfiguration.kt
index 007945445e..96bbf2da82 100644
--- a/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/SubscriptionAutoConfiguration.kt
+++ b/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/SubscriptionAutoConfiguration.kt
@@ -67,7 +67,7 @@ class SubscriptionAutoConfiguration {
 
     @Bean
     @ConditionalOnMissingBean
-    fun springSubscriptionGraphQLContextFactory(): SpringSubscriptionGraphQLContextFactory<*> = DefaultSpringSubscriptionGraphQLContextFactory
+    fun springSubscriptionGraphQLContextFactory(): SpringSubscriptionGraphQLContextFactory<*> = DefaultSpringSubscriptionGraphQLContextFactory()
 
     @Bean
     fun apolloSubscriptionProtocolHandler(
diff --git a/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/execution/SpringGraphQLContextFactory.kt b/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/execution/SpringGraphQLContextFactory.kt
index 28ab01c7c9..438511be2f 100644
--- a/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/execution/SpringGraphQLContextFactory.kt
+++ b/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/execution/SpringGraphQLContextFactory.kt
@@ -16,7 +16,6 @@
 
 package com.expediagroup.graphql.server.spring.execution
 
-import com.expediagroup.graphql.generator.execution.DefaultGraphQLContext
 import com.expediagroup.graphql.generator.execution.GraphQLContext
 import com.expediagroup.graphql.server.execution.GraphQLContextFactory
 import org.springframework.web.reactive.function.server.ServerRequest
@@ -27,8 +26,8 @@ import org.springframework.web.reactive.function.server.ServerRequest
 abstract class SpringGraphQLContextFactory<out T : GraphQLContext> : GraphQLContextFactory<T, ServerRequest>
 
 /**
- * Basic implementation of [SpringGraphQLContextFactory] that just returns a [DefaultGraphQLContext]
+ * Basic implementation of [SpringGraphQLContextFactory] that just returns null
  */
-class DefaultSpringGraphQLContextFactory : SpringGraphQLContextFactory<DefaultGraphQLContext>() {
-    override fun generateContext(request: ServerRequest) = DefaultGraphQLContext()
+class DefaultSpringGraphQLContextFactory : SpringGraphQLContextFactory<GraphQLContext>() {
+    override suspend fun generateContext(request: ServerRequest): GraphQLContext? = null
 }
diff --git a/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionHooks.kt b/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionHooks.kt
index 080a1afa2c..bd6a8a26be 100644
--- a/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionHooks.kt
+++ b/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionHooks.kt
@@ -17,7 +17,6 @@ package com.expediagroup.graphql.server.spring.subscriptions
 
 import com.expediagroup.graphql.generator.execution.GraphQLContext
 import org.springframework.web.reactive.socket.WebSocketSession
-import reactor.core.publisher.Mono
 
 /**
  * Implementation of Apollo Subscription Server Lifecycle Events
@@ -33,8 +32,8 @@ interface ApolloSubscriptionHooks {
     fun onConnect(
         connectionParams: Map<String, String>,
         session: WebSocketSession,
-        graphQLContext: GraphQLContext
-    ): Mono<GraphQLContext> = Mono.just(graphQLContext)
+        graphQLContext: GraphQLContext?
+    ): GraphQLContext? = graphQLContext
 
     /**
      * Called when the client executes a GraphQL operation.
@@ -43,7 +42,7 @@ interface ApolloSubscriptionHooks {
     fun onOperation(
         operationMessage: SubscriptionOperationMessage,
         session: WebSocketSession,
-        graphQLContext: GraphQLContext
+        graphQLContext: GraphQLContext?
     ): Unit = Unit
 
     /**
diff --git a/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionProtocolHandler.kt b/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionProtocolHandler.kt
index 5e5e543de2..d30902466f 100644
--- a/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionProtocolHandler.kt
+++ b/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionProtocolHandler.kt
@@ -16,8 +16,6 @@
 
 package com.expediagroup.graphql.server.spring.subscriptions
 
-import com.expediagroup.graphql.generator.execution.DefaultGraphQLContext
-import com.expediagroup.graphql.generator.execution.GraphQLContext
 import com.expediagroup.graphql.server.spring.GraphQLConfigurationProperties
 import com.expediagroup.graphql.server.spring.subscriptions.SubscriptionOperationMessage.ClientMessages.GQL_CONNECTION_INIT
 import com.expediagroup.graphql.server.spring.subscriptions.SubscriptionOperationMessage.ClientMessages.GQL_CONNECTION_TERMINATE
@@ -33,6 +31,7 @@ import com.fasterxml.jackson.databind.ObjectMapper
 import com.fasterxml.jackson.module.kotlin.convertValue
 import com.fasterxml.jackson.module.kotlin.readValue
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.runBlocking
 import org.slf4j.LoggerFactory
 import org.springframework.web.reactive.socket.WebSocketSession
 import reactor.core.publisher.Flux
@@ -66,7 +65,7 @@ class ApolloSubscriptionProtocolHandler(
         return try {
             when (operationMessage.type) {
                 GQL_CONNECTION_INIT.type -> onInit(operationMessage, session)
-                GQL_START.type -> onStart(operationMessage, session)
+                GQL_START.type -> startSubscription(operationMessage, session)
                 GQL_STOP.type -> onStop(operationMessage, session)
                 GQL_CONNECTION_TERMINATE.type -> onDisconnect(session)
                 else -> onUnknownOperation(operationMessage, session)
@@ -104,9 +103,10 @@ class ApolloSubscriptionProtocolHandler(
     @Suppress("Detekt.TooGenericExceptionCaught")
     private fun startSubscription(
         operationMessage: SubscriptionOperationMessage,
-        session: WebSocketSession,
-        context: GraphQLContext
+        session: WebSocketSession
     ): Flux<SubscriptionOperationMessage> {
+        val context = sessionState.getContext(session)
+
         subscriptionHooks.onOperation(operationMessage, session, context)
 
         if (operationMessage.id == null) {
@@ -114,7 +114,7 @@ class ApolloSubscriptionProtocolHandler(
             return Flux.just(basicConnectionErrorMessage)
         }
 
-        if (sessionState.operationExists(session, operationMessage)) {
+        if (sessionState.doesOperationExist(session, operationMessage)) {
             logger.info("Already subscribed to operation ${operationMessage.id} for session ${session.id}")
             return Flux.empty()
         }
@@ -149,13 +149,23 @@ class ApolloSubscriptionProtocolHandler(
     }
 
     private fun onInit(operationMessage: SubscriptionOperationMessage, session: WebSocketSession): Flux<SubscriptionOperationMessage> {
-        val connectionParams = getConnectionParams(operationMessage.payload)
-        val graphQLContext = contextFactory.generateContext(session) ?: DefaultGraphQLContext()
-        val onConnect = subscriptionHooks.onConnect(connectionParams, session, graphQLContext)
-        sessionState.saveContext(session, onConnect)
-        val acknowledgeMessage = Flux.just(acknowledgeMessage)
+        saveContext(operationMessage, session)
+        val acknowledgeMessage = Mono.just(acknowledgeMessage)
         val keepAliveFlux = getKeepAliveFlux(session)
         return acknowledgeMessage.concatWith(keepAliveFlux)
+            .onErrorReturn(getConnectionErrorMessage(operationMessage))
+    }
+
+    /**
+     * Generate the context and save it for all future messages.
+     */
+    private fun saveContext(operationMessage: SubscriptionOperationMessage, session: WebSocketSession) {
+        runBlocking {
+            val connectionParams = getConnectionParams(operationMessage.payload)
+            val context = contextFactory.generateContext(session)
+            val onConnect = subscriptionHooks.onConnect(connectionParams, session, context)
+            sessionState.saveContext(session, onConnect)
+        }
     }
 
     /**
@@ -172,25 +182,6 @@ class ApolloSubscriptionProtocolHandler(
         return emptyMap()
     }
 
-    /**
-     * Called when the client sends the start message.
-     * It triggers the specific hooks first, runs the operation, and appends it with a complete message.
-     */
-    private fun onStart(
-        operationMessage: SubscriptionOperationMessage,
-        session: WebSocketSession
-    ): Flux<SubscriptionOperationMessage> {
-        val context = sessionState.getContext(session)
-
-        // If we do not have a context, that means the init message was never sent
-        return if (context != null) {
-            context.flatMapMany { startSubscription(operationMessage, session, it) }
-        } else {
-            val message = getConnectionErrorMessage(operationMessage)
-            Flux.just(message)
-        }
-    }
-
     /**
      * Called with the publisher has completed on its own.
      */
@@ -216,7 +207,7 @@ class ApolloSubscriptionProtocolHandler(
     private fun onDisconnect(session: WebSocketSession): Flux<SubscriptionOperationMessage> {
         subscriptionHooks.onDisconnect(session)
         sessionState.terminateSession(session)
-        return Flux.empty<SubscriptionOperationMessage>()
+        return Flux.empty()
     }
 
     private fun onUnknownOperation(operationMessage: SubscriptionOperationMessage, session: WebSocketSession): Flux<SubscriptionOperationMessage> {
diff --git a/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionSessionState.kt b/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionSessionState.kt
index 701a4ba659..0a3abe27f1 100644
--- a/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionSessionState.kt
+++ b/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionSessionState.kt
@@ -31,22 +31,24 @@ internal class ApolloSubscriptionSessionState {
     // Operations are saved by web socket session id, then operation id
     internal val activeOperations = ConcurrentHashMap<String, ConcurrentHashMap<String, Subscription>>()
 
-    // OnConnect hooks are saved by web socket session id, then operation id
-    private val onConnectHooks = ConcurrentHashMap<String, Mono<GraphQLContext>>()
+    // The context is saved by web socket session id
+    private val cachedContext = ConcurrentHashMap<String, GraphQLContext>()
 
     /**
      * Save the context created from the factory and possibly updated in the onConnect hook.
      * This allows us to include some intial state to be used when handling all the messages.
      * This will be removed in [terminateSession].
      */
-    fun saveContext(session: WebSocketSession, onConnect: Mono<GraphQLContext>) {
-        onConnectHooks[session.id] = onConnect
+    fun saveContext(session: WebSocketSession, graphQLContext: GraphQLContext?) {
+        if (graphQLContext != null) {
+            cachedContext[session.id] = graphQLContext
+        }
     }
 
     /**
-     * Return the onConnect mono so that the operation can wait to start until it has been resolved.
+     * Return the context for this session.
      */
-    fun getContext(session: WebSocketSession): Mono<GraphQLContext>? = onConnectHooks[session.id]
+    fun getContext(session: WebSocketSession): GraphQLContext? = cachedContext[session.id]
 
     /**
      * Save the session that is sending keep alive messages.
@@ -119,7 +121,7 @@ internal class ApolloSubscriptionSessionState {
     fun terminateSession(session: WebSocketSession) {
         activeOperations[session.id]?.forEach { (_, subscription) -> subscription.cancel() }
         activeOperations.remove(session.id)
-        onConnectHooks.remove(session.id)
+        cachedContext.remove(session.id)
         activeKeepAliveSessions[session.id]?.cancel()
         activeKeepAliveSessions.remove(session.id)
         session.close()
@@ -128,6 +130,6 @@ internal class ApolloSubscriptionSessionState {
     /**
      * Looks up the operation for the client, to check if it already exists
      */
-    fun operationExists(session: WebSocketSession, operationMessage: SubscriptionOperationMessage): Boolean =
+    fun doesOperationExist(session: WebSocketSession, operationMessage: SubscriptionOperationMessage): Boolean =
         activeOperations[session.id]?.containsKey(operationMessage.id) ?: false
 }
diff --git a/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/SpringGraphQLSubscriptionHandler.kt b/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/SpringGraphQLSubscriptionHandler.kt
index c5708de438..0e88f5f877 100644
--- a/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/SpringGraphQLSubscriptionHandler.kt
+++ b/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/SpringGraphQLSubscriptionHandler.kt
@@ -34,7 +34,7 @@ import reactor.kotlin.core.publisher.toFlux
  */
 open class SpringGraphQLSubscriptionHandler(private val graphQL: GraphQL) {
 
-    fun executeSubscription(graphQLRequest: GraphQLRequest, graphQLContext: GraphQLContext?): Flux<GraphQLResponse<*>> = Flux.deferContextual {
+    fun executeSubscription(graphQLRequest: GraphQLRequest, graphQLContext: GraphQLContext?): Flux<GraphQLResponse<*>> =
         graphQL.execute(graphQLRequest.toExecutionInput(graphQLContext))
             .getData<Publisher<ExecutionResult>>()
             .toFlux()
@@ -43,5 +43,4 @@ open class SpringGraphQLSubscriptionHandler(private val graphQL: GraphQL) {
                 val error = KotlinGraphQLError(throwable).toGraphQLKotlinType()
                 Flux.just(GraphQLResponse<Any>(errors = listOf(error)))
             }
-    }
 }
diff --git a/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/SpringSubscriptionGraphQLContextFactory.kt b/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/SpringSubscriptionGraphQLContextFactory.kt
index f94d2aa16e..b162515088 100644
--- a/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/SpringSubscriptionGraphQLContextFactory.kt
+++ b/servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/SpringSubscriptionGraphQLContextFactory.kt
@@ -16,7 +16,6 @@
 
 package com.expediagroup.graphql.server.spring.subscriptions
 
-import com.expediagroup.graphql.generator.execution.DefaultGraphQLContext
 import com.expediagroup.graphql.generator.execution.GraphQLContext
 import com.expediagroup.graphql.server.execution.GraphQLContextFactory
 import org.springframework.web.reactive.socket.WebSocketSession
@@ -26,6 +25,9 @@ import org.springframework.web.reactive.socket.WebSocketSession
  */
 abstract class SpringSubscriptionGraphQLContextFactory<out T : GraphQLContext> : GraphQLContextFactory<T, WebSocketSession>
 
-object DefaultSpringSubscriptionGraphQLContextFactory : SpringSubscriptionGraphQLContextFactory<DefaultGraphQLContext>() {
-    override fun generateContext(request: WebSocketSession) = DefaultGraphQLContext()
+/**
+ * Basic implementation of [SpringSubscriptionGraphQLContextFactory] that just returns null
+ */
+class DefaultSpringSubscriptionGraphQLContextFactory : SpringSubscriptionGraphQLContextFactory<GraphQLContext>() {
+    override suspend fun generateContext(request: WebSocketSession): GraphQLContext? = null
 }
diff --git a/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/context/GraphQLContextFactoryIT.kt b/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/context/GraphQLContextFactoryIT.kt
index 3496e94cad..4e448067ce 100644
--- a/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/context/GraphQLContextFactoryIT.kt
+++ b/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/context/GraphQLContextFactoryIT.kt
@@ -62,7 +62,7 @@ class GraphQLContextFactoryIT(@Autowired private val testClient: WebTestClient)
         @Bean
         @ExperimentalCoroutinesApi
         fun customContextFactory(): SpringGraphQLContextFactory<CustomContext> = object : SpringGraphQLContextFactory<CustomContext>() {
-            override fun generateContext(request: ServerRequest): CustomContext {
+            override suspend fun generateContext(request: ServerRequest): CustomContext {
                 return CustomContext(
                     first = request.headers().firstHeader("X-First-Header") ?: "DEFAULT_FIRST",
                     second = request.headers().firstHeader("X-Second-Header") ?: "DEFAULT_SECOND"
diff --git a/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/routes/RouteConfigurationIT.kt b/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/routes/RouteConfigurationIT.kt
index cf7cafcbc2..3e60762249 100644
--- a/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/routes/RouteConfigurationIT.kt
+++ b/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/routes/RouteConfigurationIT.kt
@@ -38,7 +38,10 @@ import org.springframework.web.reactive.function.server.ServerRequest
 
 @SpringBootTest(
     webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT,
-    properties = ["graphql.packages=com.expediagroup.graphql.server.spring.routes"]
+    properties = [
+        "graphql.packages=com.expediagroup.graphql.server.spring.routes",
+        "graphql.sdl.enabled=true"
+    ]
 )
 @EnableAutoConfiguration
 class RouteConfigurationIT(@Autowired private val testClient: WebTestClient) {
@@ -50,7 +53,7 @@ class RouteConfigurationIT(@Autowired private val testClient: WebTestClient) {
 
         @Bean
         fun customContextFactory(): SpringGraphQLContextFactory<CustomContext> = object : SpringGraphQLContextFactory<CustomContext>() {
-            override fun generateContext(request: ServerRequest): CustomContext = CustomContext(
+            override suspend fun generateContext(request: ServerRequest): CustomContext = CustomContext(
                 value = request.headers().firstHeader("X-Custom-Header") ?: "default"
             )
         }
diff --git a/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionProtocolHandlerTest.kt b/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionProtocolHandlerTest.kt
index 2f23aff98a..c55f373e1a 100644
--- a/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionProtocolHandlerTest.kt
+++ b/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionProtocolHandlerTest.kt
@@ -16,7 +16,6 @@
 
 package com.expediagroup.graphql.server.spring.subscriptions
 
-import com.expediagroup.graphql.generator.execution.DefaultGraphQLContext
 import com.expediagroup.graphql.server.spring.GraphQLConfigurationProperties
 import com.expediagroup.graphql.server.spring.subscriptions.SubscriptionOperationMessage.ClientMessages.GQL_CONNECTION_INIT
 import com.expediagroup.graphql.server.spring.subscriptions.SubscriptionOperationMessage.ClientMessages.GQL_CONNECTION_TERMINATE
@@ -25,6 +24,7 @@ import com.expediagroup.graphql.server.spring.subscriptions.SubscriptionOperatio
 import com.expediagroup.graphql.server.spring.subscriptions.SubscriptionOperationMessage.ServerMessages.GQL_COMPLETE
 import com.expediagroup.graphql.server.spring.subscriptions.SubscriptionOperationMessage.ServerMessages.GQL_CONNECTION_ACK
 import com.expediagroup.graphql.server.spring.subscriptions.SubscriptionOperationMessage.ServerMessages.GQL_CONNECTION_ERROR
+import com.expediagroup.graphql.server.spring.subscriptions.SubscriptionOperationMessage.ServerMessages.GQL_CONNECTION_KEEP_ALIVE
 import com.expediagroup.graphql.server.spring.subscriptions.SubscriptionOperationMessage.ServerMessages.GQL_DATA
 import com.expediagroup.graphql.server.spring.subscriptions.SubscriptionOperationMessage.ServerMessages.GQL_ERROR
 import com.expediagroup.graphql.types.GraphQLError
@@ -37,11 +37,9 @@ import io.mockk.mockk
 import io.mockk.verify
 import io.mockk.verifyOrder
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.reactor.mono
 import org.junit.jupiter.api.Test
 import org.springframework.web.reactive.socket.WebSocketSession
 import reactor.core.publisher.Flux
-import reactor.core.publisher.Mono
 import reactor.test.StepVerifier
 import java.time.Duration
 import kotlin.test.assertEquals
@@ -105,33 +103,13 @@ class ApolloSubscriptionProtocolHandlerTest {
         val handler = ApolloSubscriptionProtocolHandler(config, nullContextFactory, subscriptionHandler, objectMapper, subscriptionHooks)
         val flux = handler.handle(simpleInitMessage.toJson(), session)
 
-        val message = flux.blockFirst(Duration.ofSeconds(2))
-        assertNotNull(message)
-        assertEquals(expected = GQL_CONNECTION_ACK.type, actual = message.type)
-    }
-
-    @Test
-    fun `Return only GQL_CONNECTION_ACK when sending GQL_CONNECTION_INIT and keep alive is on but no id is sent`() {
-        val config: GraphQLConfigurationProperties = mockk {
-            every { subscriptions } returns mockk {
-                every { keepAliveInterval } returns 500
-            }
-        }
-        val session: WebSocketSession = mockk {
-            every { id } returns "123"
-        }
-        val subscriptionHandler: SpringGraphQLSubscriptionHandler = mockk()
-
-        val handler = ApolloSubscriptionProtocolHandler(config, nullContextFactory, subscriptionHandler, objectMapper, subscriptionHooks)
-        val flux = handler.handle(simpleInitMessage.toJson(), session)
-
-        val message = flux.blockFirst(Duration.ofSeconds(2))
-        assertNotNull(message)
-        assertEquals(expected = GQL_CONNECTION_ACK.type, actual = message.type)
+        StepVerifier.create(flux)
+            .expectNextMatches { it.type == GQL_CONNECTION_ACK.type }
+            .verifyComplete()
     }
 
     @Test
-    fun `Return GQL_CONNECTION_ACK + GQL_CONNECTION_KEEP_ALIVE when sent type is GQL_CONNECTION_INIT and keep alive is on and id is sent`() {
+    fun `Return GQL_CONNECTION_ACK + GQL_CONNECTION_KEEP_ALIVE when sending GQL_CONNECTION_INIT and keep alive is on`() {
         val config: GraphQLConfigurationProperties = mockk {
             every { subscriptions } returns mockk {
                 every { keepAliveInterval } returns 500
@@ -146,11 +124,11 @@ class ApolloSubscriptionProtocolHandlerTest {
         val handler = ApolloSubscriptionProtocolHandler(config, nullContextFactory, subscriptionHandler, objectMapper, subscriptionHooks)
 
         val initFlux = handler.handle(operationMessage, session)
-
-        val message = initFlux.blockFirst(Duration.ofSeconds(2))
-        assertNotNull(message)
-        assertEquals(expected = GQL_CONNECTION_ACK.type, actual = message.type)
-        initFlux.subscribe().dispose()
+        StepVerifier.create(initFlux)
+            .expectNextMatches { it.type == GQL_CONNECTION_ACK.type }
+            .expectNextMatches { it.type == GQL_CONNECTION_KEEP_ALIVE.type }
+            .thenCancel()
+            .verify()
     }
 
     @Test
@@ -235,8 +213,8 @@ class ApolloSubscriptionProtocolHandlerTest {
 
         StepVerifier.create(initFlux)
             .expectSubscription()
-            .expectNextMatches { it.type == "connection_ack" }
-            .expectNextMatches { it.type == "ka" }
+            .expectNextMatches { it.type == GQL_CONNECTION_ACK.type }
+            .expectNextMatches { it.type == GQL_CONNECTION_KEEP_ALIVE.type }
             .thenCancel()
             .verify()
 
@@ -276,10 +254,12 @@ class ApolloSubscriptionProtocolHandlerTest {
         handler.handle(simpleInitMessage.toJson(), session)
         val flux = handler.handle(operationMessage, session)
 
-        val message = flux.blockFirst(Duration.ofSeconds(2))
-        assertNotNull(message)
-        assertEquals(expected = GQL_CONNECTION_ERROR.type, actual = message.type)
-        assertEquals(expected = "abc", actual = message.id)
+        StepVerifier.create(flux)
+            .expectNextMatches {
+                it.type == GQL_CONNECTION_ERROR.type && it.payload == null
+            }
+            .expectComplete()
+            .verify()
     }
 
     @Test
@@ -465,12 +445,12 @@ class ApolloSubscriptionProtocolHandlerTest {
         }
         val subscriptionHandler: SpringGraphQLSubscriptionHandler = mockk()
         val subscriptionHooks: ApolloSubscriptionHooks = mockk {
-            every { onConnect(any(), any(), any()) } returns mono { null }
+            every { onConnect(any(), any(), any()) } returns null
         }
         val handler = ApolloSubscriptionProtocolHandler(config, nullContextFactory, subscriptionHandler, objectMapper, subscriptionHooks)
         val flux = handler.handle(simpleInitMessage.toJson(), session)
-        flux.blockFirst(Duration.ofSeconds(2))
-        verify(exactly = 1) { subscriptionHooks.onConnect(any(), session, any()) }
+        flux.subscribe().dispose()
+        verify(exactly = 1) { subscriptionHooks.onConnect(any(), any(), any()) }
     }
 
     @Test
@@ -487,11 +467,11 @@ class ApolloSubscriptionProtocolHandlerTest {
         }
         val subscriptionHandler: SpringGraphQLSubscriptionHandler = mockk()
         val subscriptionHooks: ApolloSubscriptionHooks = mockk {
-            every { onConnect(any(), any(), any()) } returns mono { null }
+            every { onConnect(any(), any(), any()) } returns null
         }
         val handler = ApolloSubscriptionProtocolHandler(config, nullContextFactory, subscriptionHandler, objectMapper, subscriptionHooks)
         val flux = handler.handle(operationMessage, session)
-        flux.blockFirst(Duration.ofSeconds(2))
+        flux.subscribe().dispose()
         verify(exactly = 1) { subscriptionHooks.onConnect(payload, session, any()) }
     }
 
@@ -512,12 +492,21 @@ class ApolloSubscriptionProtocolHandlerTest {
             every { executeSubscription(eq(graphQLRequest), any()) } returns Flux.just(expectedResponse)
         }
         val subscriptionHooks: ApolloSubscriptionHooks = mockk {
-            every { onConnect(any(), any(), any()) } returns Mono.just(DefaultGraphQLContext())
+            every { onConnect(any(), any(), any()) } returns null
             every { onOperation(any(), any(), any()) } returns Unit
             every { onOperationComplete(any()) } returns Unit
         }
         val handler = ApolloSubscriptionProtocolHandler(config, nullContextFactory, subscriptionHandler, objectMapper, subscriptionHooks)
-        handler.handle(simpleInitMessage.toJson(), session)
+        val initFlux = handler.handle(simpleInitMessage.toJson(), session)
+        StepVerifier.create(initFlux)
+            .expectNextCount(1)
+            .expectComplete()
+            .verify()
+
+        verify(exactly = 1) {
+            subscriptionHooks.onConnect(any(), any(), any())
+        }
+
         val startFlux = handler.handle(startMessage, session)
         StepVerifier.create(startFlux)
             .expectNextMatches { it.type == GQL_DATA.type && it.payload == expectedResponse }
@@ -526,9 +515,9 @@ class ApolloSubscriptionProtocolHandlerTest {
             .verify()
 
         verify(exactly = 1) {
-            subscriptionHooks.onConnect(any(), any(), any())
             subscriptionHooks.onOperation(any(), any(), any())
         }
+
         verifyOrder {
             subscriptionHooks.onConnect(any(), any(), any())
             subscriptionHooks.onOperation(any(), any(), any())
@@ -536,14 +525,13 @@ class ApolloSubscriptionProtocolHandlerTest {
     }
 
     @Test
-    fun `Do not send any messages when onConnect throws error`() {
+    fun `Return GQL_CONNECTION_ERROR when onConnect throws error`() {
         val config: GraphQLConfigurationProperties = mockk {
             every { subscriptions } returns mockk {
                 every { keepAliveInterval } returns null
             }
         }
         val initMessage = simpleInitMessage.toJson()
-        val startMessage = SubscriptionOperationMessage(GQL_START.type).toJson()
         val session: WebSocketSession = mockk {
             every { id } returns "123"
         }
@@ -554,9 +542,8 @@ class ApolloSubscriptionProtocolHandlerTest {
         }
         val handler = ApolloSubscriptionProtocolHandler(config, nullContextFactory, subscriptionHandler, objectMapper, subscriptionHooks)
         val initFlux = handler.handle(initMessage, session)
-        val startFlux = handler.handle(startMessage, session)
-        initFlux.blockFirst(Duration.ofSeconds(2))
-        StepVerifier.create(startFlux)
+
+        StepVerifier.create(initFlux)
             .expectNextMatches {
                 it.type == GQL_CONNECTION_ERROR.type && it.payload == null
             }
@@ -581,7 +568,7 @@ class ApolloSubscriptionProtocolHandlerTest {
         }
         val handler = ApolloSubscriptionProtocolHandler(config, nullContextFactory, subscriptionHandler, objectMapper, subscriptionHooks)
         val flux = handler.handle(operationMessage, session)
-        flux.blockFirst(Duration.ofSeconds(2))
+        flux.subscribe().dispose()
         verify(exactly = 1) { subscriptionHooks.onOperationComplete(session) }
     }
 
diff --git a/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/subscriptions/SubscriptionWebSocketHandlerIT.kt b/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/subscriptions/SubscriptionWebSocketHandlerIT.kt
index 83411bce30..17e37bc918 100644
--- a/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/subscriptions/SubscriptionWebSocketHandlerIT.kt
+++ b/servers/graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/server/spring/subscriptions/SubscriptionWebSocketHandlerIT.kt
@@ -61,12 +61,11 @@ class SubscriptionWebSocketHandlerIT(
     fun `verify subscription`() {
         val request = GraphQLRequest("subscription { characters }")
         val messageId = "1"
-        val initMessage = getInitMessage(messageId)
         val startMessage = getStartMessage(messageId, request)
         val dataOutput = TestPublisher.create<String>()
 
         val response = client.execute(uri) { session ->
-            executeSubsciption(session, initMessage, startMessage, dataOutput)
+            executeSubsciption(session, startMessage, dataOutput)
         }.subscribe()
 
         StepVerifier.create(dataOutput)
@@ -86,12 +85,11 @@ class SubscriptionWebSocketHandlerIT(
     fun `verify subscription to counter`() {
         val request = GraphQLRequest("subscription { counter }")
         val messageId = "2"
-        val initMessage = getInitMessage(messageId)
         val startMessage = getStartMessage(messageId, request)
         val dataOutput = TestPublisher.create<String>()
 
         val response = client.execute(uri) { session ->
-            executeSubsciption(session, initMessage, startMessage, dataOutput)
+            executeSubsciption(session, startMessage, dataOutput)
         }.subscribe()
 
         StepVerifier.create(dataOutput)
@@ -111,14 +109,13 @@ class SubscriptionWebSocketHandlerIT(
     fun `verify subscription with context`() {
         val request = GraphQLRequest("subscription { ticker }")
         val messageId = "3"
-        val initMessage = getInitMessage(messageId)
         val startMessage = getStartMessage(messageId, request)
         val dataOutput = TestPublisher.create<String>()
         val headers = HttpHeaders()
         headers.set("X-Custom-Header", "junit")
 
         val response = client.execute(uri, headers) { session ->
-            executeSubsciption(session, initMessage, startMessage, dataOutput)
+            executeSubsciption(session, startMessage, dataOutput)
         }.subscribe()
 
         StepVerifier.create(dataOutput)
@@ -131,11 +128,10 @@ class SubscriptionWebSocketHandlerIT(
 
     private fun executeSubsciption(
         session: WebSocketSession,
-        initMessage: String,
         startMessage: String,
         dataOutput: TestPublisher<String>
     ): Mono<Void> {
-        val firstMessage = session.textMessage(initMessage).toMono()
+        val firstMessage = session.textMessage(basicInitMessage).toMono()
             .concatWith(session.textMessage(startMessage).toMono())
 
         return session.send(firstMessage)
@@ -164,11 +160,7 @@ class SubscriptionWebSocketHandlerIT(
         fun subscription(): Subscription = SimpleSubscription()
 
         @Bean
-        fun customContextFactory(): SpringSubscriptionGraphQLContextFactory<SubscriptionContext> = object : SpringSubscriptionGraphQLContextFactory<SubscriptionContext>() {
-            override fun generateContext(request: WebSocketSession): SubscriptionContext = SubscriptionContext(
-                value = request.handshakeInfo.headers.getFirst("X-Custom-Header") ?: "default"
-            )
-        }
+        fun customContextFactory(): SpringSubscriptionGraphQLContextFactory<SubscriptionContext> = CustomContextFactory()
     }
 
     // GraphQL spec requires at least single query to be present as Query type is needed to run introspection queries
@@ -194,7 +186,13 @@ class SubscriptionWebSocketHandlerIT(
 
     data class SubscriptionContext(val value: String) : GraphQLContext
 
+    class CustomContextFactory : SpringSubscriptionGraphQLContextFactory<SubscriptionContext>() {
+        override suspend fun generateContext(request: WebSocketSession): SubscriptionContext = SubscriptionContext(
+            value = request.handshakeInfo.headers.getFirst("X-Custom-Header") ?: "default"
+        )
+    }
+
     private fun SubscriptionOperationMessage.toJson() = objectMapper.writeValueAsString(this)
-    private fun getInitMessage(id: String) = SubscriptionOperationMessage(ClientMessages.GQL_CONNECTION_INIT.type, id).toJson()
+    private val basicInitMessage = SubscriptionOperationMessage(ClientMessages.GQL_CONNECTION_INIT.type).toJson()
     private fun getStartMessage(id: String, payload: Any) = SubscriptionOperationMessage(ClientMessages.GQL_START.type, id, payload).toJson()
 }
diff --git a/website/docs/schema-generator/execution/contextual-data.md b/website/docs/schema-generator/execution/contextual-data.md
index 365bac6c24..36f1c53571 100644
--- a/website/docs/schema-generator/execution/contextual-data.md
+++ b/website/docs/schema-generator/execution/contextual-data.md
@@ -2,6 +2,7 @@
 id: contextual-data
 title: Contextual Data
 ---
+
 All GraphQL servers have a concept of a "context". A GraphQL context contains metadata that is useful to the GraphQL
 server, but shouldn't necessarily be part of the GraphQL schema. A prime example of something that is appropriate
 for the GraphQL context would be trace headers for an OpenTracing system such as
@@ -10,9 +11,8 @@ its function, but the server needs the information to ensure observability.
 
 The contents of the GraphQL context vary across applications and it is up to the GraphQL server developers to decide
 what it should contain. `graphql-kotlin-server` provides a simple mechanism to
-build context per query execution through
-[GraphQLContextFactory](../../server/graphql-context-factory.md).
-Once a context factory is available it will then be used to populate GraphQL context based on the incoming request and make it available during query execution.
+build a context per operation with the [GraphQLContextFactory](../../server/graphql-context-factory.md).
+If a custom factory is defined, it will then be used to populate GraphQL context based on the incoming request and make it available during execution.
 
 ## GraphQLContext Interface
 
@@ -20,40 +20,47 @@ The easiest way to specify a context class is to use the `GraphQLContext` marker
 it is just used to inform the schema generator that this is the class that should be used as the context for every request.
 
 ```kotlin
-
 class MyGraphQLContext(val customValue: String) : GraphQLContext
-
 ```
 
 Then you can just use the class as an argument and it will be automatically injected during execution time.
 
 ```kotlin
-
-class ContextualQuery {
+class ContextualQuery : Query {
     fun contextualQuery(
         context: MyGraphQLContext,
         value: Int
     ): String = "The custom value was ${context.customValue} and the value was $value"
 }
-
 ```
 
 The above query would produce the following GraphQL schema:
 
 ```graphql
-
-schema {
-  query: Query
-}
-
 type Query {
   contextualQuery(value: Int!): String!
 }
-
 ```
 
 Note that the argument that implements `GraphQLContext` is not reflected in the GraphQL schema.
 
+## Handling Context Errors
+
+The [GraphQLContextFactory](../../server/graphql-context-factory.md) may return `null`. If your factory implementation never returns `null`, then there is no need to change your schema.
+If the factory could return `null`, then the context arugments in your schema should be nullable so a runtime exception is not thrown.
+
+```kotlin
+class ContextualQuery : Query {
+    fun contextualQuery(context: MyGraphQLContext?, value: Int): String {
+        if (context != null) {
+            return "The custom value was ${context.customValue} and the value was $value"
+        }
+
+        return "The context was null and the value was $value"
+    }
+}
+```
+
 ## Injection Customization
 
 The context is injected into the execution through the `FunctionDataFetcher` class.
diff --git a/website/docs/server/graphql-context-factory.md b/website/docs/server/graphql-context-factory.md
index 3f9070a9dc..a220faf263 100644
--- a/website/docs/server/graphql-context-factory.md
+++ b/website/docs/server/graphql-context-factory.md
@@ -2,31 +2,67 @@
 id: graphql-context-factory
 title: GraphQLContextFactory
 ---
-Similar to the [GraphQLRequestParser](./graphql-request-parser.md), `GraphQLContextFactory` has a generic method for handling the `Request` and the `GraphQLContext`.
+
+`GraphQLContextFactory` is a generic method for generating a `GraphQLContext` for each request.
 
 ```kotlin
 
 interface GraphQLContextFactory<out Context : GraphQLContext, Request> {
-    fun generateContext(request: Request): Context?
+    suspend fun generateContext(request: Request): Context?
 }
 
 ```
 
-Given the server request, the interface should create the custom `GraphQLContext` class to be used for every new operation.
+Given the generic server request, the interface should create a `GraphQLContext` class to be used for every new operation.
 The context must implement the `GraphQLContext` interface from `graphql-kotlin-schema-generator`.
 See [execution context](../schema-generator/execution/contextual-data.md) for more info on how the context can be used in the schema functions.
 
-A specific `graphql-kotlin-*-server` library may provide an abstract class on top of this interface so users only have to be concerned with the context.
+## Nullable Context
 
-For example the `graphql-kotlin-spring-server` provides the following class, which sets the request type:
+The factory may return `null` if a context is not required for execution. This allows the library to have a default factory that just returns `null`.
+If your custom factory never returns `null`, then there is no need to use nullable arguments.
+However, if your custom factory may return `null`, you must define the context argument as nullable in the schema functions or a runtime exception will be thrown.
 
 ```kotlin
+data class CustomContext(val value: String) : GraphQLContext
 
-abstract class SpringGraphQLContextFactory<out T : GraphQLContext> : GraphQLContextFactory<T, ServerRequest>
+class CustomFactory : GraphQLContextFactory<CustomContext, ServerRequest> {
+    suspend fun generateContext(request: Request): Context? {
+        if (isSpecialRequest(request)) {
+            return null
+        }
+
+        val value = callSomeSuspendableService(request)
+        return CustomContext(value)
+    }
+}
+
+class MyQuery : Query {
+
+    fun getResults(context: CustomContext?, input: String): String {
+        if (context != null) {
+            return getDataWithContext(input, context)
+        }
+
+        return getBasicData(input)
+    }
+}
+```
+
+## Suspendable Factory
+The interface is marked as a `suspend` function to allow the asynchronous fetching of context data.
+This may be helpful if you need to call some other services to calculate a context value.
+
+## Server-Specific Abstractions
 
+A specific `graphql-kotlin-*-server` library may provide an abstract class on top of this interface so users only have to be concerned with the context class and not the server class type.
+For example the `graphql-kotlin-spring-server` provides the following class, which sets the request type:
+
+```kotlin
+abstract class SpringGraphQLContextFactory<out T : GraphQLContext> : GraphQLContextFactory<T, ServerRequest>
 ```
 
-## HTTP Headers
+## HTTP Headers and Cookies
 
-For common use cases around authorization, authentication, or tracing you may need to read HTTP headers.
+For common use cases around authorization, authentication, or tracing you may need to read HTTP headers and cookies.
 This should be done in the `GraphQLContextFactory` and relevant data should be added to the context to be accessible during schema exectuion.
diff --git a/website/docs/server/graphql-request-handler.md b/website/docs/server/graphql-request-handler.md
index 84ca8bc626..8434f96860 100644
--- a/website/docs/server/graphql-request-handler.md
+++ b/website/docs/server/graphql-request-handler.md
@@ -2,7 +2,10 @@
 id: graphql-request-handler
 title: GraphQLRequestHandler
 ---
-The `GraphQLRequestHandler` is an open and extendable class that contains the basic logic to get a `GraphQLResponse` from `graphql-kotlin-types`.
-It accepts a `GraphQLRequest`, an optional [GraphQLContext](./graphql-context-factory.md) and sends that to the GraphQL schema along with the [DataLoaderRegistry](data-loaders.md).
+The `GraphQLRequestHandler` is an open and extendable class that contains the basic logic to get a `GraphQLResponse`.
+
+It requires a `GraphQLSchema` and a [DataLoaderRegistryFactory](data-loaders.md) in the constructor.
+For each request, it accepts a `GraphQLRequest` and an optional [GraphQLContext](./graphql-context-factory.md), and calls the `DataLoaderRegistryFactory` to generate a new `DataLoaderRegistry`.
+Then all of these objects are sent to the schema for execution and the result is mapped to a `GraphQLResponse`.
 
 There shouldn't be much need to change this class but if you wanted to add custom logic or logging it is possible to override it or just create your own.
diff --git a/website/docs/server/graphql-request-parser.md b/website/docs/server/graphql-request-parser.md
index 00c1c2aa1d..be945a5a95 100644
--- a/website/docs/server/graphql-request-parser.md
+++ b/website/docs/server/graphql-request-parser.md
@@ -2,34 +2,30 @@
 id: graphql-request-parser
 title: GraphQLRequestParser
 ---
-The `GraphQLRequestParser` interface is requrired to parse the library-specific HTTP request object into the common `GraphQLRequest` class from `graphql-kotlin-types`.
+The `GraphQLRequestParser` interface is requrired to parse the library-specific HTTP request object into the common `GraphQLRequest` class.
 
 ```kotlin
-
 interface GraphQLRequestParser<Request> {
     suspend fun parseRequest(request: Request): GraphQLServerRequest<*>?
 }
-
 ```
 
 While not officially part of the spec, there is a standard format used by most GraphQL clients and servers for [serving GraphQL over HTTP](https://graphql.org/learn/serving-over-http/).
 Following the above convention, GraphQL clients should generally use HTTP POST requests with the following body structure
 
 ```json
-
 {
   "query": "...",
   "operationName": "...",
   "variables": { "myVariable": "someValue" }
 }
-
 ```
 
 where
 
--   `query` is a required field and contains operation (query, mutation or subscription) that specify their selection set to be executed
--   `operationName` is an optional operation name, only required if multiple operations are specified in `query` string
--   `variables` is an optional field that holds an arbitrary JSON objects that are referenced as input arguments from `query` string
+- `query` is a required field and contains the operation (query, mutation, or subscription) to be executed
+- `operationName` is an optional string, only required if multiple operations are specified in the `query` string.
+- `variables` is an optional map of JSON objects that are referenced as input arguments in the `query` string
 
 GraphQL Kotlin server supports both single and batch GraphQL requests. Batch requests are represented as a list of individual
 GraphQL requests. When processing batch requests, same context will be used for processing all requests and server will respond
diff --git a/website/docusaurus.config.js b/website/docusaurus.config.js
index 1dbeb75aea..d8113f3fd3 100644
--- a/website/docusaurus.config.js
+++ b/website/docusaurus.config.js
@@ -30,8 +30,14 @@ module.exports = {
   plugins: [],
   themeConfig: {
     image: "img/undraw_online.svg",
+    colorMode: {
+      defaultMode: 'dark',
+    },
     prism: {
-      additionalLanguages: ['kotlin'],
+      defaultLanguage: 'kotlin',
+      additionalLanguages: ['kotlin', 'groovy'],
+      theme: require('prism-react-renderer/themes/github'),
+      darkTheme: require('prism-react-renderer/themes/dracula')
     },
     navbar: {
       title: "GraphQL Kotlin",