Fix for kotlin await

Author: slinkydeveloper

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

@@ -27,7 +27,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,33 @@
 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].
+ *
+ * The state machine unwinds a suspended or closed invocation by completing the awaited future
+ * exceptionally with [AbortedExecutionException] — a bare [Throwable] the Java SDK sneaky-throws.
+ * Letting that raw throwable surface into coroutine code is not idiomatic: it is not a
+ * [CancellationException], so structured concurrency treats it as a genuine failure and a user
+ * `catch (e: Throwable)` would silently swallow the suspension instead of propagating it. Wrapping
+ * it as [CancellationException] keeps invocation suspension behaving like ordinary coroutine
+ * cancellation, as it did before the shared-core state machine contract.
+ *
+ * All SDK await sites go through this wrapper (instead of `kotlinx.coroutines.future.await`).
+ */
+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> {