Fix custom serde factory usage in kotlin (#477)

Author: slinkydeveloper

sdk-api-kotlin-gen/src/main/kotlin/dev/restate/sdk/kotlin/gen/KElementConverter.kt

@@ -8,6 +8,7 @@
 // https://github.com/restatedev/sdk-java/blob/main/LICENSE
 package dev.restate.sdk.kotlin.gen
 
+import com.google.devtools.ksp.KSTypeNotPresentException
 import com.google.devtools.ksp.KspExperimental
 import com.google.devtools.ksp.getAnnotationsByType
 import com.google.devtools.ksp.getVisibility
@@ -104,7 +105,7 @@ class KElementConverter(
     val customSerdeFactory: CustomSerdeFactory? =
         classDeclaration.getAnnotationsByType(CustomSerdeFactory::class).firstOrNull()
     if (customSerdeFactory != null) {
-      serdeFactoryDecl = "new " + customSerdeFactory.value + "()"
+      serdeFactoryDecl = parseAnnotationClassParameter { customSerdeFactory.value } + "()"
     }
     data.withSerdeFactoryDecl(serdeFactoryDecl)
   }
@@ -335,4 +336,18 @@ class KElementConverter(
 
     return typeName
   }
+
+  @OptIn(KspExperimental::class)
+  private fun parseAnnotationClassParameter(block: () -> KClass<*>): String? {
+    return try { // KSTypeNotPresentException will be thrown
+      block.invoke().qualifiedName
+    } catch (e: KSTypeNotPresentException) {
+      var res: String? = null
+      val declaration = e.ksType.declaration
+      if (declaration is KSClassDeclaration) {
+        declaration.qualifiedName?.asString()?.let { res = it }
+      }
+      res
+    }
+  }
 }

sdk-api-kotlin/src/main/kotlin/dev/restate/sdk/kotlin/serialization/KotlinSerializationSerdeFactory.kt

@@ -42,9 +42,9 @@ constructor(private val json: Json = Json.Default) : SerdeFactory {
 
   @PublishedApi
   internal class KtTypeTag<T>(
-      internal val type: KClass<*>,
+      val type: KClass<*>,
       /** Reified type */
-      internal val kotlinType: KType?
+      val kotlinType: KType?
   ) : TypeTag<T>
 
   override fun <T : Any?> create(typeTag: TypeTag<T>): Serde<T> {

sdk-core/src/test/kotlin/dev/restate/sdk/core/kotlinapi/CodegenTest.kt

@@ -19,6 +19,9 @@ import dev.restate.sdk.core.statemachine.ProtoUtils.*
 import dev.restate.sdk.kotlin.*
 import dev.restate.sdk.kotlin.serialization.*
 import dev.restate.serde.Serde
+import dev.restate.serde.SerdeFactory
+import dev.restate.serde.TypeRef
+import dev.restate.serde.TypeTag
 import java.util.stream.Stream
 import kotlinx.serialization.Serializable
 
@@ -212,6 +215,33 @@ class CodegenTest : TestDefinitions.TestSuite {
     }
   }
 
+  class MyCustomSerdeFactory : SerdeFactory {
+    override fun <T : Any?> create(typeTag: TypeTag<T?>): Serde<T?> {
+      check(typeTag is KotlinSerializationSerdeFactory.KtTypeTag)
+      check(typeTag.type == Byte::class)
+      return Serde.using<Byte>({ b -> byteArrayOf(b) }, { it[0] }) as Serde<T?>
+    }
+
+    override fun <T : Any?> create(typeRef: TypeRef<T?>): Serde<T?> {
+      check(typeRef.type == Byte::class)
+      return Serde.using<Byte>({ b -> byteArrayOf(b) }, { it[0] }) as Serde<T?>
+    }
+
+    override fun <T : Any?> create(clazz: Class<T?>?): Serde<T?> {
+      check(clazz == Byte::class.java)
+      return Serde.using<Byte>({ b -> byteArrayOf(b) }, { it[0] }) as Serde<T?>
+    }
+  }
+
+  @CustomSerdeFactory(MyCustomSerdeFactory::class)
+  @Service(name = "CustomSerdeService")
+  class CustomSerdeService {
+    @Handler
+    suspend fun echo(context: Context, input: Byte): Byte {
+      return input
+    }
+  }
+
   override fun definitions(): Stream<TestDefinition> {
     return Stream.of(
         testInvocation({ ServiceGreeter() }, "greet")
@@ -358,6 +388,10 @@ class CodegenTest : TestDefinitions.TestSuite {
                     Slice.EMPTY),
                 outputCmd(),
                 END_MESSAGE),
+        testInvocation({ CustomSerdeService() }, "echo")
+            .withInput(startMessage(1), inputCmd(byteArrayOf(1)))
+            .onlyBidiStream()
+            .expectingOutput(outputCmd(byteArrayOf(1)), END_MESSAGE),
     )
   }
 }