Properly use UndispatchedCoroutine context for context-specific elements (#3413)

Fixes #3411

Author: qwwdfsad

kotlinx-coroutines-core/common/src/Builders.common.kt

@@ -163,7 +163,7 @@ public suspend fun <T> withContext(
         if (newContext[ContinuationInterceptor] == oldContext[ContinuationInterceptor]) {
             val coroutine = UndispatchedCoroutine(newContext, uCont)
             // There are changes in the context, so this thread needs to be updated
-            withCoroutineContext(newContext, null) {
+            withCoroutineContext(coroutine.context, null) {
                 return@sc coroutine.startUndispatchedOrReturn(coroutine, block)
             }
         }

kotlinx-coroutines-core/jvm/test/ThreadContextElementTest.kt

@@ -156,6 +156,41 @@ class ThreadContextElementTest : TestBase() {
             }
         }
     }
+
+    class JobCaptor(val capturees: ArrayList<Job> = ArrayList()) : ThreadContextElement<Unit> {
+
+        companion object Key : CoroutineContext.Key<MyElement>
+
+        override val key: CoroutineContext.Key<*> get() = Key
+
+        override fun updateThreadContext(context: CoroutineContext) {
+            capturees.add(context.job)
+        }
+
+        override fun restoreThreadContext(context: CoroutineContext, oldState: Unit) {
+        }
+    }
+
+    @Test
+    fun testWithContextJobAccess() = runTest {
+        val captor = JobCaptor()
+        val manuallyCaptured = ArrayList<Job>()
+        runBlocking(captor) {
+            manuallyCaptured += coroutineContext.job
+            withContext(CoroutineName("undispatched")) {
+                manuallyCaptured += coroutineContext.job
+                withContext(Dispatchers.IO) {
+                    manuallyCaptured += coroutineContext.job
+                }
+                // Context restored, captured again
+                manuallyCaptured += coroutineContext.job
+            }
+            // Context restored, captured again
+            manuallyCaptured += coroutineContext.job
+        }
+
+        assertEquals(manuallyCaptured, captor.capturees)
+    }
 }
 
 class MyData