Propagate transport coroutine context

krpc/krpc-server/src/commonMain/kotlin/kotlinx/rpc/krpc/server/KrpcServer.kt

@@ -47,7 +47,7 @@ public abstract class KrpcServer(
     transport: KrpcTransport,
 ) : RpcServer, KrpcEndpoint {
     // we make a child here, so we can send cancellation messages before closing the connection
-    final override val coroutineContext: CoroutineContext = SupervisorJob(transport.coroutineContext.job)
+    final override val coroutineContext: CoroutineContext = transport.coroutineContext + SupervisorJob(transport.coroutineContext.job)
 
     private val logger = CommonLogger.logger(objectId())
 
@@ -130,7 +130,7 @@ public abstract class KrpcServer(
         descriptor: RpcServiceDescriptor<Service>,
         serviceFactory: (CoroutineContext) -> Service,
     ): KrpcServerService<Service> {
-        val serviceInstanceContext = SupervisorJob(coroutineContext.job)
+        val serviceInstanceContext = coroutineContext + SupervisorJob(coroutineContext.job)
 
         return KrpcServerService(
             service = serviceFactory(serviceInstanceContext),

krpc/krpc-test/src/commonTest/kotlin/kotlinx/rpc/krpc/test/CoroutineContextPropagationTest.kt

@@ -0,0 +1,57 @@
+/*
+ * Copyright 2023-2025 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.rpc.krpc.test
+
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.currentCoroutineContext
+import kotlinx.coroutines.test.runTest
+import kotlinx.rpc.krpc.rpcClientConfig
+import kotlinx.rpc.krpc.rpcServerConfig
+import kotlinx.rpc.krpc.serialization.json.json
+import kotlinx.rpc.withService
+import kotlin.coroutines.CoroutineContext
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+class CoroutineContextPropagationTest {
+    private val rpcServerConfig = rpcServerConfig {
+        serialization {
+            json()
+        }
+    }
+    private val rpcClientConfig = rpcClientConfig {
+        serialization {
+            json {
+                ignoreUnknownKeys = true
+            }
+        }
+    }
+
+    object CoroutineElement : CoroutineContext.Element {
+        object Key : CoroutineContext.Key<CoroutineElement>
+
+        override val key: CoroutineContext.Key<*> = Key
+    }
+
+    @Test
+    fun test() = runTest {
+        var actualContext: CoroutineElement? = null
+        val transport = LocalTransport(transportContext = CoroutineElement)
+        val server = KrpcTestServer(rpcServerConfig, transport.server)
+        val client = KrpcTestClient(rpcClientConfig, transport.client)
+        server.registerService(Echo::class) {
+            object : Echo {
+                override suspend fun echo(message: String): String = run {
+                    actualContext = currentCoroutineContext().get(CoroutineElement.Key)
+                    "response"
+                }
+
+                override val coroutineContext: CoroutineContext = it
+            }
+        }
+        client.withService(Echo::class).echo("request")
+        assertEquals(actualContext, CoroutineElement)
+    }
+}

krpc/krpc-test/src/commonTest/kotlin/kotlinx/rpc/krpc/test/LocalTransport.kt

@@ -13,15 +13,22 @@ import kotlinx.rpc.krpc.KrpcTransport
 import kotlinx.rpc.krpc.KrpcTransportMessage
 import kotlin.coroutines.CoroutineContext
 
-class LocalTransport(parentScope: CoroutineScope? = null) : CoroutineScope {
-    override val coroutineContext = parentScope?.run { SupervisorJob(coroutineContext.job) }
+class LocalTransport(
+    parentScope: CoroutineScope? = null,
+    transportContext: CoroutineContext? = null,
+) : CoroutineScope {
+    override val coroutineContext = parentScope
+        ?.run { SupervisorJob(coroutineContext.job) }
         ?: SupervisorJob()
 
     private val clientIncoming = Channel<KrpcTransportMessage>()
     private val serverIncoming = Channel<KrpcTransportMessage>()
 
     val client: KrpcTransport = object : KrpcTransport {
-        override val coroutineContext: CoroutineContext = Job([email protected])
+        override val coroutineContext: CoroutineContext = Job([email protected]).let {
+            if(transportContext != null) transportContext + it
+            else it
+        }
 
         override suspend fun send(message: KrpcTransportMessage) {
             serverIncoming.send(message)
@@ -33,7 +40,10 @@ class LocalTransport(parentScope: CoroutineScope? = null) : CoroutineScope {
     }
 
     val server: KrpcTransport = object : KrpcTransport {
-        override val coroutineContext: CoroutineContext = Job([email protected])
+        override val coroutineContext: CoroutineContext = Job([email protected]).let {
+            if(transportContext != null) transportContext + it
+            else it
+        }
 
         override suspend fun send(message: KrpcTransportMessage) {
             clientIncoming.send(message)