Enable Json Schema generation in kotlin

Fix test where a bad type was being used for serialization/deserialization

Author: slinkydeveloper

sdk-api-kotlin-gen/src/test/kotlin/dev/restate/sdk/kotlin/CodegenTest.kt

@@ -57,8 +57,8 @@ class CodegenTest : TestDefinitions.TestSuite {
     @Exclusive
     suspend fun complexType(
         context: ObjectContext,
-        request: Map<Output, List<out Input>>
-    ): Map<Input, List<out Output>> {
+        request: Map<String, List<out Input>>
+    ): Map<String, List<out Output>> {
       return mapOf()
     }
   }

sdk-api-kotlin/build.gradle.kts

@@ -13,6 +13,8 @@ dependencies {
   implementation(kotlinLibs.kotlinx.serialization.core)
   implementation(kotlinLibs.kotlinx.serialization.json)
 
+  implementation("io.bkbn:kompendium-json-schema:4.0.0-alpha")
+
   implementation(coreLibs.log4j.api)
   implementation(platform(coreLibs.opentelemetry.bom))
   implementation(coreLibs.opentelemetry.kotlin)

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

@@ -9,8 +9,12 @@
 package dev.restate.sdk.kotlin
 
 import dev.restate.sdk.common.DurablePromiseKey
+import dev.restate.sdk.common.RichSerde
 import dev.restate.sdk.common.Serde
 import dev.restate.sdk.common.StateKey
+import io.bkbn.kompendium.json.schema.KotlinXSchemaConfigurator
+import io.bkbn.kompendium.json.schema.SchemaGenerator
+import io.bkbn.kompendium.json.schema.definition.JsonSchema
 import java.nio.ByteBuffer
 import java.nio.charset.StandardCharsets
 import kotlin.reflect.typeOf
@@ -70,12 +74,13 @@ object KtSerdes {
       }
 
   /** Creates a [Serde] implementation using the `kotlinx.serialization` json module. */
-  fun <T : Any?> json(serializer: KSerializer<T>): Serde<T> {
-    return object : Serde<T> {
+  inline fun <reified T : Any?> json(serializer: KSerializer<T>): Serde<T> {
+    return object : RichSerde<T> {
       override fun serialize(value: T?): ByteArray {
         if (value == null) {
           return Json.encodeToString(JsonNull.serializer(), JsonNull).encodeToByteArray()
         }
+
         return Json.encodeToString(serializer, value).encodeToByteArray()
       }
 
@@ -86,6 +91,17 @@ object KtSerdes {
       override fun contentType(): String {
         return "application/json"
       }
+
+      override fun jsonSchema(): String {
+        val schema =
+            SchemaGenerator.fromTypeToSchema(
+                type = typeOf<T>(),
+                cache = mutableMapOf(),
+                schemaConfigurator = KotlinXSchemaConfigurator(),
+            )
+
+        return Json.encodeToString(JsonSchema.serializer(), schema)
+      }
     }
   }
 }

sdk-common/src/main/java/dev/restate/sdk/common/RichSerde.java

@@ -14,12 +14,15 @@
 /**
  * Richer version of {@link Serde} containing schema information.
  *
+ * <p>This API should be considered unstable to implement.
+ *
  * <p>You can create one using {@link #withSchema(Object, Serde)}.
  */
 public interface RichSerde<T extends @Nullable Object> extends Serde<T> {
 
   /**
-   * @return a Draft 2020-12 Json Schema
+   * @return a Draft 2020-12 Json Schema. It should be self-contained, and MUST not contain refs to
+   *     files. If the schema shouldn't be serialized with Jackson, return a {@link String}
    */
   Object jsonSchema();
 

sdk-core/src/main/java/dev/restate/sdk/core/EndpointManifest.java

@@ -10,6 +10,7 @@
 
 import static dev.restate.sdk.core.ServiceProtocol.*;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
 import dev.restate.sdk.common.HandlerType;
 import dev.restate.sdk.common.RichSerde;
 import dev.restate.sdk.common.ServiceType;
@@ -108,8 +109,17 @@ private static Input convertHandlerInput(HandlerSpecification<?, ?> spec) {
             : new Input().withRequired(true).withContentType(acceptContentType);
 
     if (spec.getRequestSerde() instanceof RichSerde) {
-      input.setJsonSchema(
-          Objects.requireNonNull(((RichSerde<?>) spec.getRequestSerde()).jsonSchema()));
+      Object jsonSchema =
+          Objects.requireNonNull(((RichSerde<?>) spec.getRequestSerde()).jsonSchema());
+      if (jsonSchema instanceof String) {
+        // We need to convert it to databind JSON value
+        try {
+          jsonSchema = MANIFEST_OBJECT_MAPPER.readTree((String) jsonSchema);
+        } catch (JsonProcessingException e) {
+          throw new RuntimeException("The schema generated by RichSerde is not a valid JSON", e);
+        }
+      }
+      input.setJsonSchema(jsonSchema);
     }
     return input;
   }
@@ -123,8 +133,17 @@ private static Output convertHandlerOutput(HandlerSpecification<?, ?> spec) {
                 .withSetContentTypeIfEmpty(false);
 
     if (spec.getResponseSerde() instanceof RichSerde) {
-      output.setJsonSchema(
-          Objects.requireNonNull(((RichSerde<?>) spec.getResponseSerde()).jsonSchema()));
+      Object jsonSchema =
+          Objects.requireNonNull(((RichSerde<?>) spec.getResponseSerde()).jsonSchema());
+      if (jsonSchema instanceof String) {
+        // We need to convert it to databind JSON value
+        try {
+          jsonSchema = MANIFEST_OBJECT_MAPPER.readTree((String) jsonSchema);
+        } catch (JsonProcessingException e) {
+          throw new RuntimeException("The schema generated by RichSerde is not a valid JSON", e);
+        }
+      }
+      output.setJsonSchema(jsonSchema);
     }
 
     return output;

sdk-core/src/main/java/dev/restate/sdk/core/ServiceProtocol.java

@@ -135,7 +135,7 @@ static String serviceDiscoveryProtocolVersionToHeaderValue(
             "Service discovery protocol version '%s' has no header value", version.getNumber()));
   }
 
-  private static final ObjectMapper MANIFEST_OBJECT_MAPPER = new ObjectMapper();
+  static final ObjectMapper MANIFEST_OBJECT_MAPPER = new ObjectMapper();
 
   @JsonFilter("V2FieldsFilter")
   interface V2Mixin {}