Fix for kotlin await

Author: slinkydeveloper

sdk-api-kotlin/src/main/kotlin/dev/restate/sdk/kotlin/ContextImpl.kt

@@ -29,7 +29,6 @@ import kotlin.jvm.optionals.getOrNull
 import kotlin.time.Duration
 import kotlin.time.toJavaDuration
 import kotlinx.coroutines.*
-import kotlinx.coroutines.future.await
 
 internal class ContextImpl
 internal constructor(

sdk-api-kotlin/src/main/kotlin/dev/restate/sdk/kotlin/Util.kt

@@ -9,9 +9,23 @@
 package dev.restate.sdk.kotlin
 
 import dev.restate.common.Slice
+import dev.restate.sdk.common.AbortedExecutionException
 import dev.restate.sdk.endpoint.definition.HandlerContext
 import dev.restate.serde.Serde
+import java.util.concurrent.CompletionStage
 import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.future.await as kotlinxAwait
+
+/**
+ * Awaits [this], translating the SDK-internal [AbortedExecutionException] control-flow signal into
+ * a coroutine [CancellationException] (coroutine idiomatic way of initiating an abort).
+ */
+internal suspend fun <T> CompletionStage<T>.await(): T =
+    try {
+      this.kotlinxAwait()
+    } catch (e: AbortedExecutionException) {
+      throw CancellationException("Restate invocation suspended or closed").apply { initCause(e) }
+    }
 
 internal fun <T : Any?> Serde<T>.serializeWrappingException(
     handlerContext: HandlerContext,

sdk-api-kotlin/src/main/kotlin/dev/restate/sdk/kotlin/futures.kt

@@ -23,7 +23,6 @@ import kotlin.time.Duration
 import kotlin.time.toJavaDuration
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.currentCoroutineContext
-import kotlinx.coroutines.future.await
 import kotlinx.coroutines.launch
 
 internal abstract class BaseDurableFutureImpl<T : Any?> : DurableFuture<T> {