Example using custom json schema factory

build.gradle.kts

plugins {
  application
  kotlin("jvm") version "2.1.20"
  kotlin("plugin.serialization") version "2.1.20"
  id("com.google.devtools.ksp") version "2.1.20-1.0.32"
}

repositories {
  mavenCentral()
}

val restateVersion = "2.1.0"
val schemaKenerator = "2.1.2"

dependencies {
  // Annotation processor
  ksp("dev.restate:sdk-api-kotlin-gen:$restateVersion")

  // Restate SDK
  implementation("dev.restate:sdk-kotlin-http:$restateVersion")

  // Schema kenerator deps
  implementation("io.github.smiley4:schema-kenerator-core:$schemaKenerator")
  implementation("io.github.smiley4:schema-kenerator-serialization:$schemaKenerator")
  implementation("io.github.smiley4:schema-kenerator-jsonschema:${schemaKenerator}")

  // Logging
  implementation("org.apache.logging.log4j:log4j-api:2.24.3")
}

kotlin {
  jvmToolchain(21)
}

// Configure main class
application {
  mainClass.set("my.example.GreeterKt")
}

custom_json_support.kt

// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
//
// This file is part of the Restate Java SDK,
// which is released under the MIT license.
//
// You can find a copy of the license in file LICENSE in the root
// directory of this repository or package, or at
// https://github.com/restatedev/sdk-java/blob/main/LICENSE
package my.example

import dev.restate.serde.Serde
import dev.restate.serde.kotlinx.KotlinSerializationSerdeFactory
import io.github.smiley4.schemakenerator.jsonschema.JsonSchemaSteps
import io.github.smiley4.schemakenerator.jsonschema.JsonSchemaSteps.compileReferencing
import io.github.smiley4.schemakenerator.jsonschema.JsonSchemaSteps.generateJsonSchema
import io.github.smiley4.schemakenerator.jsonschema.JsonSchemaSteps.handleCoreAnnotations
import io.github.smiley4.schemakenerator.jsonschema.TitleBuilder
import io.github.smiley4.schemakenerator.jsonschema.data.IntermediateJsonSchemaData
import io.github.smiley4.schemakenerator.jsonschema.data.RefType
import io.github.smiley4.schemakenerator.jsonschema.jsonDsl.*
import io.github.smiley4.schemakenerator.serialization.SerializationSteps.analyzeTypeUsingKotlinxSerialization
import io.github.smiley4.schemakenerator.serialization.SerializationSteps.initial
import io.github.smiley4.schemakenerator.serialization.SerializationSteps.renameMembers
import kotlinx.serialization.ExperimentalSerializationApi
import kotlinx.serialization.KSerializer
import kotlinx.serialization.json.Json

class CustomJsonSerdeFactory : KotlinSerializationSerdeFactory(json = Json {}, jsonSchemaFactory = CustomJsonSchemaFactory)

// More or less copied from DefaultJsonSchemaFactory, just adding handleCoreAnnotations()
object CustomJsonSchemaFactory : KotlinSerializationSerdeFactory.JsonSchemaFactory {
  @OptIn(ExperimentalSerializationApi::class)
  override fun generateSchema(json: Json, serializer: KSerializer<*>) =
    runCatching {
      // This is the configuration of schema-kenerator library
      var initialStep =
        initial(serializer.descriptor).analyzeTypeUsingKotlinxSerialization {
          serializersModule = json.serializersModule
        }

      if (json.configuration.namingStrategy != null) {
        initialStep = initialStep.renameMembers(json.configuration.namingStrategy!!)
      }

      val intermediateStep =
        initialStep.generateJsonSchema {
          optionalHandling = JsonSchemaSteps.OptionalHandling.NON_REQUIRED
        }
          // Make sure we read schema kenerator annotations
          .handleCoreAnnotations()
      intermediateStep.writeTitles()
      intermediateStep.compileAndSanitize(serializer)
    }.getOrDefault(JsonObject(mutableMapOf())).let { Serde.StringifiedJsonSchema(it.prettyPrint()) }

  @OptIn(ExperimentalSerializationApi::class)
  private fun IntermediateJsonSchemaData.compileAndSanitize(serializer: KSerializer<*>): JsonNode {
    val compiledSchema = this.compileReferencing(RefType.SIMPLE)

    // In case of nested schemas, compileReferencing also contains self schema...
    val rootSchemaName =
      TitleBuilder.BUILDER_SIMPLE(
        compiledSchema.typeData, this.typeDataById
      )

    // If schema is not json object, then it's boolean, so we're good no need for
    // additional manipulation
    if (compiledSchema.json !is JsonObject) {
      return compiledSchema.json
    }

    // Assemble the final schema now
    val rootNode = compiledSchema.json as JsonObject
    // Add $schema
    rootNode.properties.put(
      "\$schema", JsonTextValue("https://json-schema.org/draft/2020-12/schema")
    )
    // Add $defs
    val definitions =
      compiledSchema.definitions.filter { it.key != rootSchemaName }.toMutableMap()
    if (definitions.isNotEmpty()) {
      rootNode.properties.put("\$defs", JsonObject(definitions))
    }
    // Replace all $refs
    rootNode.fixRefsPrefix("#/definitions/$rootSchemaName")
    // If the root type is nullable, it should be in the schema too
    if (serializer.descriptor.isNullable) {
      val oldTypeProperty = rootNode.properties["type"]
      if (oldTypeProperty is JsonTextValue) {
        rootNode.properties["type"] = array {
          item(oldTypeProperty.value)
          item(JsonTextValue("null"))
        }
      } else if (oldTypeProperty is JsonArray) {
        oldTypeProperty.items.add(JsonTextValue("null"))
      }
    }

    return rootNode
  }

  private fun IntermediateJsonSchemaData.writeTitles() {
    this.entries.forEach { schema ->
      if (schema.json is JsonObject) {
        if ((schema.typeData.isMap ||
              schema.typeData.isCollection ||
              schema.typeData.isEnum ||
              schema.typeData.isInlineValue ||
              schema.typeData.typeParameters.isNotEmpty() ||
              schema.typeData.members.isNotEmpty()) &&
          (schema.json as JsonObject).properties["title"] == null
        ) {
          (schema.json as JsonObject).properties["title"] =
            JsonTextValue(TitleBuilder.BUILDER_SIMPLE(schema.typeData, this.typeDataById))
        }
      }
    }
  }

  private fun JsonNode.fixRefsPrefix(rootDefinition: String) {
    when (this) {
      is JsonArray -> this.items.forEach { it.fixRefsPrefix(rootDefinition) }
      is JsonObject -> this.fixRefsPrefix(rootDefinition)
      else -> {}
    }
  }

  private fun JsonObject.fixRefsPrefix(rootDefinition: String) {
    this.properties.computeIfPresent("\$ref") { key, node ->
      if (node is JsonTextValue) {
        if (node.value.startsWith(rootDefinition)) {
          JsonTextValue("#/" + node.value.removePrefix(rootDefinition))
        } else {
          JsonTextValue("#/\$defs/" + node.value.removePrefix("#/definitions/"))
        }
      } else {
        node
      }
    }
    this.properties.values.forEach { it.fixRefsPrefix(rootDefinition) }
  }
}

Greeter.kt

package my.example

import dev.restate.sdk.annotation.CustomSerdeFactory
import dev.restate.sdk.annotation.Handler
import dev.restate.sdk.annotation.Service
import dev.restate.sdk.http.vertx.RestateHttpServer
import dev.restate.sdk.kotlin.*
import dev.restate.sdk.kotlin.endpoint.endpoint
import dev.restate.serde.kotlinx.KotlinSerializationSerdeFactory
import io.github.smiley4.schemakenerator.core.annotations.Default
import io.github.smiley4.schemakenerator.core.annotations.Description
import io.github.smiley4.schemakenerator.core.annotations.Example
import io.github.smiley4.schemakenerator.core.annotations.Title
import kotlinx.serialization.Serializable
import kotlinx.serialization.json.Json
import kotlin.time.Duration.Companion.seconds

@Description("some description")
@Serializable
data class Greeting(
  @Description("String field description")
  @Default("A default String value")
  @Example("An example of a String value")
  val name: String)

@Service
@CustomSerdeFactory(CustomJsonSerdeFactory::class)
class Greeter {

  @Handler
  suspend fun greet(ctx: Context, greeting: Greeting): String {
    // Respond to caller
    return "You said hi to ${greeting.name}!";
  }
}

fun main() {
  RestateHttpServer.listen(endpoint {
    bind(Greeter())
  })
}