From f74f6f7a970e31d0d4313ddbe06d0ae29c602a3b Mon Sep 17 00:00:00 2001
From: Mo Sobhy <mo@sobhy.dev>
Date: Sun, 31 May 2026 23:16:19 +0100
Subject: [PATCH] [Bugfix] Adjust `$ref` handling logic in
 `DefaultJsonSchemaFactory` to avoid replacing nested type name if it shares
 the prefix of the root class name

When you have a type that shares the name of it's parent type (eg: Task & TaskType), with the way the current `fixRefsPrefix` works (recurses through all subtypes to handle nested types) you end up replacing the def path of the subtype and it becomes unserialisable and crashes at runtime, and you have to use custom serialisers or rename your class. This checks that the root definition is followed by a / (to indicate nesting) before replacing it).
---
 .../serde/kotlinx/DefaultJsonSchemaFactory.kt | 11 ++--
 .../restate/serde/kotlinx/KotlinxSerdeTest.kt | 65 +++++++++++++++++++
 2 files changed, 71 insertions(+), 5 deletions(-)

diff --git a/sdk-serde-kotlinx/src/main/kotlin/dev/restate/serde/kotlinx/DefaultJsonSchemaFactory.kt b/sdk-serde-kotlinx/src/main/kotlin/dev/restate/serde/kotlinx/DefaultJsonSchemaFactory.kt
index f1b1c88b..a14368ac 100644
--- a/sdk-serde-kotlinx/src/main/kotlin/dev/restate/serde/kotlinx/DefaultJsonSchemaFactory.kt
+++ b/sdk-serde-kotlinx/src/main/kotlin/dev/restate/serde/kotlinx/DefaultJsonSchemaFactory.kt
@@ -23,7 +23,6 @@ import io.github.smiley4.schemakenerator.jsonschema.jsonDsl.array
 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 kotlin.collections.set
 import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.KSerializer
 import kotlinx.serialization.json.Json
@@ -67,7 +66,7 @@ object DefaultJsonSchemaFactory : KotlinSerializationSerdeFactory.JsonSchemaFact
                 // Add $schema
                 rootNode.properties.put(
                     "\$schema",
-                    JsonTextValue("https://json-schema.org/draft/2020-12/schema"),
+                  JsonTextValue("https://json-schema.org/draft/2020-12/schema"),
                 )
                 // Add $defs
                 val definitions =
@@ -109,7 +108,7 @@ object DefaultJsonSchemaFactory : KotlinSerializationSerdeFactory.JsonSchemaFact
                 (schema.json as JsonObject).properties["title"] == null
         ) {
           (schema.json as JsonObject).properties["title"] =
-              JsonTextValue(TitleBuilder.BUILDER_SIMPLE(schema.typeData, this.typeDataById))
+            JsonTextValue(TitleBuilder.BUILDER_SIMPLE(schema.typeData, this.typeDataById))
         }
       }
     }
@@ -126,8 +125,10 @@ object DefaultJsonSchemaFactory : KotlinSerializationSerdeFactory.JsonSchemaFact
   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))
+        if (node.value == rootDefinition) {
+          JsonTextValue("#/")
+        } else if (node.value.startsWith("$rootDefinition/")) {
+          JsonTextValue("#/" + node.value.removePrefix("$rootDefinition/"))
         } else {
           JsonTextValue("#/\$defs/" + node.value.removePrefix("#/definitions/"))
         }
diff --git a/sdk-serde-kotlinx/src/test/kotlin/dev/restate/serde/kotlinx/KotlinxSerdeTest.kt b/sdk-serde-kotlinx/src/test/kotlin/dev/restate/serde/kotlinx/KotlinxSerdeTest.kt
index 4583a389..a5de8386 100644
--- a/sdk-serde-kotlinx/src/test/kotlin/dev/restate/serde/kotlinx/KotlinxSerdeTest.kt
+++ b/sdk-serde-kotlinx/src/test/kotlin/dev/restate/serde/kotlinx/KotlinxSerdeTest.kt
@@ -176,6 +176,71 @@ class KotlinxSerdeTest {
     )
   }
 
+  @Serializable
+  enum class TaskStatus {
+    TODO,
+    IN_PROGRESS,
+    DONE,
+  }
+
+  @Serializable
+  enum class PriorityOrder {
+    HIGH,
+    MID,
+    LOW,
+  }
+
+  @Serializable
+  data class Task(val title: String, val status: TaskStatus, val priority: PriorityOrder)
+
+  @Test
+  fun schemaGenWithExternalEnum() {
+    testSchemaGen<Task>(
+        $$"""
+        {
+                 "type": "object",
+                 "required": [
+                    "title",
+                    "status",
+                    "priority"
+                 ],
+                 "properties": {
+                    "title": {
+                       "type": "string"
+                    },
+                    "priority": {
+                       "$ref": "#/$defs/PriorityOrder"
+                    },
+                    "status": {
+                       "$ref": "#/$defs/TaskStatus"
+                    }
+                 },
+                 "title": "Task",
+                 "$schema": "https://json-schema.org/draft/2020-12/schema",
+                 "$defs": {
+                    "TaskStatus": {
+                       "enum": [
+                          "TODO",
+                          "IN_PROGRESS",
+                          "DONE"
+                       ],
+                       "title": "TaskStatus"
+                    },
+                    "PriorityOrder": {
+                       "enum": [
+                          "HIGH",
+                          "MID",
+                          "LOW"
+                       ],
+                       "title": "PriorityOrder"
+                    }
+                 }
+              }
+        """
+            .trimIndent()
+    )
+  }
+
   inline fun <reified T : Any?> testSchemaGen(expectedSchema: String) {
     val expectedJsonElement = Json.decodeFromString<JsonElement>(expectedSchema)
     val actualSchema =