@@ -33,6 +33,9 @@ import kotlinx.coroutines.runBlocking
 import org.apache.airflow.sdk.execution.CoordinatorComm
 import org.apache.airflow.sdk.execution.LogSender
 import org.apache.airflow.sdk.execution.Logger
+import org.apache.airflow.sdk.execution.comm.ErrorResponse
+import org.apache.airflow.sdk.execution.comm.StartupDetails
+import org.apache.airflow.sdk.execution.runTask
 import kotlin.text.substringAfterLast
 import kotlin.text.substringBeforeLast
 
@@ -147,10 +150,11 @@ class Server(
           aSocket(SelectorManager(Dispatchers.IO)).tcp().connect(comm).use { socket ->
             logger.debug("Connected comm", mapOf("addr" to comm))
             CoordinatorComm(
-              bundle,
               socket.openReadChannel(),
               socket.openWriteChannel(autoFlush = true),
-            ).startProcessing()
+            ).use { coordinator ->
+              dispatchTask(bundle, coordinator)
+            }
           }
         } finally {
           deferral.complete(Unit)
@@ -164,4 +168,25 @@ class Server(
         }
       }
     }
+
+  internal suspend fun dispatchTask(
+    bundle: Bundle,
+    coordinator: CoordinatorComm,
+  ) {
+    val frame = coordinator.readMessage()
+    when (val body = frame.body) {
+      is StartupDetails -> runTaskAndReport(bundle, body, coordinator)
+      is ErrorResponse -> throw ApiError("[${body.error}] ${body.detail}")
+      else -> throw ApiError("Unexpected initial frame (id=${frame.id})")
+    }
+  }
+
+  private suspend fun runTaskAndReport(
+    bundle: Bundle,
+    startup: StartupDetails,
+    coordinator: CoordinatorComm,
+  ) {
+    val result = runTask(bundle, startup, coordinator)
+    coordinator.communicate<Unit>(result)
+  }
 }
@@ -23,12 +23,18 @@ import io.ktor.utils.io.ByteReadChannel
 import io.ktor.utils.io.ByteWriteChannel
 import io.ktor.utils.io.readByteArray
 import io.ktor.utils.io.writeByteArray
+import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
 import org.apache.airflow.sdk.ApiError
-import org.apache.airflow.sdk.Bundle
 import org.apache.airflow.sdk.execution.comm.ErrorResponse
-import org.apache.airflow.sdk.execution.comm.StartupDetails
+import java.util.concurrent.ConcurrentHashMap
 import kotlin.concurrent.atomics.AtomicInt
 import kotlin.concurrent.atomics.ExperimentalAtomicApi
 
@@ -44,10 +50,9 @@ data class OutgoingFrame(
 
 @OptIn(ExperimentalAtomicApi::class)
 class CoordinatorComm(
-  private val bundle: Bundle,
   private val reader: ByteReadChannel,
   private val writer: ByteWriteChannel,
-) {
+) : AutoCloseable {
   internal companion object {
     private val logger = Logger(CoordinatorComm::class)
 
@@ -57,77 +62,63 @@ class CoordinatorComm(
   }
 
   private val nextId = AtomicInt(0)
-  private var shutDownRequested = false
-  private val commMutex = Mutex()
+  private val writeMutex = Mutex()
+  private val stateMutex = Mutex()
+  private val pending = ConcurrentHashMap<Int, CompletableDeferred<IncomingFrame>>()
 
-  suspend fun startProcessing() {
-    while (!shutDownRequested) {
-      processOnce(::handleIncoming)
-    }
-    logger.debug("Goodbye")
-  }
+  @Volatile
+  private var readError: ApiError? = null
+  private var dispatcherStarted = false
 
-  private suspend fun processOnce(handle: suspend (IncomingFrame) -> Unit) {
-    val prefix = reader.readByteArray(4) // First 4 bytes as length.
-    if (prefix.size != 4) { // Something is terribly wrong. Let's bail.
-      logger.error("Need 4 prefix bytes", mapOf("actual" to prefix.size))
-      shutDownRequested = true
-      return
-    }
+  private val dispatcherScope = CoroutineScope(Dispatchers.IO + SupervisorJob())
 
-    val payloadLength = Frame.parseLengthPrefix(prefix)
-    val payload = reader.readByteArray(payloadLength)
-    if (payload.size != payloadLength) { // Something is terribly wrong. Let's bail.
-      logger.error(
-        "Payload length not right",
-        mapOf("expect" to payloadLength, "receive" to payload.size),
-      )
-      shutDownRequested = true
-      return
+  suspend fun readMessage(): IncomingFrame =
+    stateMutex.withLock {
+      check(!dispatcherStarted) { "readMessage cannot be used after the dispatcher has started" }
+      readFrame()
     }
-    val frame = decode(payload)
-    logger.debug("Handling", mapOf("id" to frame.id))
-    handle(frame)
-  }
 
   private suspend fun sendMessage(
     id: Int,
     body: Any,
   ) {
     val data = encode(OutgoingFrame(id, body))
     logger.debug("Sending", mapOf("id" to id, "body" to body))
-    writer.writeByteArray(Frame.lengthPrefix(data.size))
-    writer.writeByteArray(data)
-  }
-
-  suspend fun handleIncoming(frame: IncomingFrame) {
-    when (val request = frame.body) {
-      null -> {}
-      is ErrorResponse -> throw ApiError("[${request.error}] ${request.detail}")
-      is StartupDetails -> {
-        communicate<Unit>(runTask(bundle, request, this))
-        shutDownRequested = true
-      }
+    writeMutex.withLock {
+      writer.writeByteArray(Frame.lengthPrefix(data.size) + data)
     }
   }
 
   @Throws(ApiError::class)
   suspend fun communicateImpl(body: Any): Any {
     val requestId = nextId.fetchAndAdd(1)
-    return commMutex.withLock {
-      var frame: IncomingFrame? = null
+    val waiter = CompletableDeferred<IncomingFrame>()
 
-      suspend fun handle(f: IncomingFrame) {
-        frame = f
+    stateMutex.withLock {
+      readError?.let { throw it }
+      if (!dispatcherStarted) {
+        dispatcherStarted = true
+        dispatcherScope.launch { readLoop() }
       }
+      pending[requestId] = waiter
+    }
+
+    // A close() concurrent with the block above either drains this waiter (it
+    // was registered first) or is visible here, so no caller is left awaiting a
+    // comm that is already gone.
+    readError?.let {
+      pending.remove(requestId)
+      throw it
+    }
+
+    try {
       sendMessage(requestId, body)
-      processOnce(::handle)
-      val received = frame ?: throw ApiError("No response received")
-      if (received.id != requestId) {
-        throw ApiError("response id ${received.id} does not match request id $requestId")
-      }
-      received.body ?: Unit
+    } catch (e: Throwable) {
+      pending.remove(requestId)
+      throw e
     }
+
+    return waiter.await().body ?: Unit
   }
 
   @Throws(ApiError::class)
@@ -138,4 +129,73 @@ class CoordinatorComm(
       else -> throw ApiError("Unexpected response type ${response::class.java}")
     }
   }
+
+  /**
+   * Stop the dispatcher and fail anything still awaiting a response.
+   */
+  override fun close() {
+    dispatcherScope.cancel()
+    failAllWaiters(readError ?: ApiError("Coordinator comm closed"))
+  }
+
+  private suspend fun readPayload(): ByteArray {
+    val prefix = reader.readByteArray(4) // First 4 bytes as length.
+    if (prefix.size != 4) {
+      throw ApiError("Coordinator socket closed while reading frame length")
+    }
+    val payloadLength = Frame.parseLengthPrefix(prefix)
+    val payload = reader.readByteArray(payloadLength)
+    if (payload.size != payloadLength) {
+      throw ApiError("Coordinator socket closed while reading frame payload")
+    }
+    return payload
+  }
+
+  private suspend fun readFrame(): IncomingFrame {
+    val frame = decode(readPayload())
+    logger.debug("Received", mapOf("id" to frame.id))
+    return frame
+  }
+
+  private suspend fun readLoop() {
+    while (true) {
+      val raw =
+        try {
+          Frame.decodeRaw(readPayload())
+        } catch (e: CancellationException) {
+          // Coroutine cancellation is delivered by throwing this exception, and
+          // cooperative cancellation requires rethrowing it so it propagates.
+          throw e
+        } catch (e: Throwable) {
+          failAllWaiters(e)
+          return
+        }
+      logger.debug("Received", mapOf("id" to raw.id))
+
+      val waiter = pending.remove(raw.id)
+      if (waiter == null) {
+        logger.warning("Discarding response with no matching request", mapOf("id" to raw.id))
+        continue
+      }
+
+      try {
+        waiter.complete(IncomingFrame(raw.id, Frame.decodeBody(raw)))
+      } catch (e: CancellationException) {
+        throw e
+      } catch (e: Throwable) {
+        waiter.completeExceptionally(
+          ApiError("Cannot decode response for request ${raw.id}: ${e.message}")
+            .apply { initCause(e) },
+        )
+      }
+    }
+  }
+
+  private fun failAllWaiters(cause: Throwable) {
+    val error =
+      cause as? ApiError
+        ?: ApiError("Coordinator comm closed: ${cause.message}").apply { initCause(cause) }
+    readError = error
+    pending.keys.forEach { id -> pending.remove(id)?.completeExceptionally(error) }
+  }
 }
@@ -28,6 +28,12 @@ import org.apache.airflow.sdk.execution.comm.Discriminator
 import org.msgpack.core.MessagePack
 import java.io.ByteArrayOutputStream
 
+data class RawFrame(
+  val id: Int,
+  val rawBody: Any?,
+  val rawError: Any?,
+)
+
 object Frame {
   private val mapper =
     ObjectMapper().apply {
@@ -43,7 +49,7 @@ object Frame {
     body: Any,
   ): ByteArray = encodeFrame(id, body)
 
-  fun decode(bytes: ByteArray): IncomingFrame {
+  fun decodeRaw(bytes: ByteArray): RawFrame {
     val unpacker = MessagePack.newDefaultUnpacker(bytes)
     val headerSize = unpacker.unpackArrayHeader()
     check(headerSize >= 1) { "Unexpected Task SDK frame arity $headerSize" }
@@ -53,8 +59,14 @@ object Frame {
     val rawError = if (headerSize >= 3) unpacker.unpackAny() else null
     unpacker.close()
 
-    val body = decodeMessage(rawError) ?: decodeMessage(rawBody)
-    return IncomingFrame(id, body)
+    return RawFrame(id, rawBody, rawError)
+  }
+
+  fun decodeBody(raw: RawFrame): Any? = decodeMessage(raw.rawError) ?: decodeMessage(raw.rawBody)
+
+  fun decode(bytes: ByteArray): IncomingFrame {
+    val raw = decodeRaw(bytes)
+    return IncomingFrame(raw.id, decodeBody(raw))
   }
 
   fun lengthPrefix(length: Int) =
 
@@ -198,15 +198,15 @@ internal class Logger(
     arguments: Map<String, Any> = emptyMap(),
   ) = log(Level.DEBUG, message, arguments)
 
-  fun error(
+  fun warning(
     message: String,
     arguments: Map<String, Any> = emptyMap(),
-  ) = log(Level.ERROR, message, arguments)
+  ) = log(Level.WARNING, message, arguments)
 
-  fun warning(
+  fun error(
     message: String,
     arguments: Map<String, Any> = emptyMap(),
-  ) = log(Level.WARNING, message, arguments)
+  ) = log(Level.ERROR, message, arguments)
 
   private fun log(
     level: Level,