Repository: victools/jsonschema-generator (java)

Issue:
Redundant definition splitting and unwanted inlining in self-referencing type hierarchies
## Given

### Types

```java
@JsonTypeInfo(use = JsonTypeInfo.Id.NAME)
@JsonSubTypes({
    @JsonSubTypes.Type(value = OrCondition.class, name = "or"),
    @JsonSubTypes.Type(value = EqualsCondition.class, name = "equals")
})
public sealed interface SearchCondition permits OrCondition, EqualsCondition {}

public record OrCondition(List<SearchCondition> conditions) implements SearchCondition {}

public record EqualsCondition(Object value) implements SearchCondition {}
```

### Config

```java
var config = new SchemaGeneratorConfigBuilder(SchemaVersion.DRAFT_2020_12, OptionPreset.PLAIN_JSON)
    .with(new JacksonModule())
    .build();
var jsonSchema = new SchemaGenerator(config).generateSchema(SearchCondition.class);

System.out.println(jsonSchema.toPrettyString());
```

## Problems
- definitions are split in two parts (though technically correct, it's pretty confusing to read);
- base class definition is inlined (again, it's correct, but not optimal);
- there's no way to avoid definition splitting with current api (`SchemaGenerationContext#createStandardDefinitionReference` is the culprit, but the only available alternative - `SchemaGenerationContext#createStandardDefinition` - understandably leads to infinite recursion).

Using a custom definition provider, I was able to avoid base class definition inlining, but not the definition splitting.

## Possible solutions
- simple: make `SchemaGenerationContextImpl#traverseGenericType` (or it's overload) `public` - this will allow to cache definition `ObjectNode` without falling into infinite recursion;
- proper: incorporate definition reference caching into library itself by adding something like `SchemaGenerationContext#getOrCreateStandardDefinitionReference`.

## Schemas

### Generated
```json
{
    "$schema": "https://json-schema.org/draft/2020-12/schema",
    "$defs": {
        "EqualsCondition-1": {
            "type": "object",
            "properties": {
                "value": {}
            }
        },
        "EqualsCondition-2": {
            "$ref": "#/$defs/EqualsCondition-1",
            "type": "object",
            "properties": {
                "@type": {
                    "const": "equals"
                }
            },
            "required": [
                "@type"
            ]
        },
        "OrCondition-1": {
            "type": "object",
            "properties": {
                "conditions": {
                    "type": "array",
                    "items": {
                        "anyOf": [
                            {
                                "$ref": "#/$defs/OrCondition-2"
                            },
                            {
                                "$ref": "#/$defs/EqualsCondition-2"
                            }
                        ]
                    }
                }
            }
        },
        "OrCondition-2": {
            "$ref": "#/$defs/OrCondition-1",
            "type": "object",
            "properties": {
                "@type": {
                    "const": "or"
                }
            },
            "required": [
                "@type"
            ]
        }
    },
    "anyOf": [
        {
            "$ref": "#/$defs/OrCondition-2"
        },
        {
            "$ref": "#/$defs/EqualsCondition-2"
        }
    ]
}
```

### Ideal

```json
{
    "$schema": "https://json-schema.org/draft/2020-12/schema",
    "$defs": {
        "SearchCondition": {
            "oneOf": [
                {
                    "$ref": "#/$defs/EqualsCondition"
                },
                {
                    "$ref": "#/$defs/OrCondition"
                }
            ]
        },
        "EqualsCondition": {
            "type": "object",
            "properties": {
                "@type": {
                    "const": "equals"
                },
                "value": {}
            },
            "required": [
                "@type",
                "value"
            ]
        },
        "OrCondition": {
            "type": "object",
            "properties": {
                "@type": {
                    "const": "or"
                },
                "conditions": {
                    "type": "array",
                    "items": {
                        "$ref": "#/$defs/SearchCondition"
                    }
                }
            },
            "required": [
                "@type",
                "conditions"
            ]
        }
    },
    "$ref": "#/$defs/SearchCondition"
}
```

Pull request:
feat: Enable Custom Definition never to be inlined
Fixes #280.

In main generator: introducing the option to never inline a custom definition even when only referenced once without having to globally enable `Option.DEFINITIONS_FOR_ALL_OBJECTS`.

In jackson module: introducing `JacksonOption.ALWAYS_REF_SUBTYPES` in order to apply the new behavior to subtype (wrapper) definitions.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 67fb790..6ee2148 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,10 +8,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### `jsonschema-generator`
 #### Added
 - new `Option.DEFINITIONS_FOR_MEMBER_SUPERTYPES` to disable the transparent member subtype resolution, i.e., enable inclusion of a supertype schema
+- new `DefinitionType.ALWAYS_REF` for custom definitions, to produce centralised definition even if just referenced once
+
+#### Fixed
+- under some circumstances, even after the general schema clean-up procedure there were unnecessary `allOf` wrappers containing just a single entry
 
 #### Changed
 - enable moving subtype schema into `$defs` and thereby reduce number of unnecessary `anyOf` wrappers
 
+### `jsonschema-module-jackson`
+#### Added
+- new `JacksonOption.ALWAYS_REF_SUBTYPES`, to produce centralised definition for each resolved subtype (wrapper) even if just referenced once
+
 ### `jsonschema-maven-plugin`
 #### Added
 - support `<classpath>` parameter in order to also consider compile dependencies for the schema generation (and/or ignoring runtime dependencies)
diff --git a/jsonschema-generator/src/main/java/com/github/victools/jsonschema/generator/CustomDefinition.java b/jsonschema-generator/src/main/java/com/github/victools/jsonschema/generator/CustomDefinition.java
index 61307ad..06aad9d 100644
--- a/jsonschema-generator/src/main/java/com/github/victools/jsonschema/generator/CustomDefinition.java
+++ b/jsonschema-generator/src/main/java/com/github/victools/jsonschema/generator/CustomDefinition.java
@@ -99,6 +99,18 @@ public class CustomDefinition {
         return this.definitionType == DefinitionType.INLINE;
     }
 
+    /**
+     * Getter for the flag indicating whether this custom definition should always be included in the {@link SchemaKeyword#TAG_DEFINITIONS} even if it
+     * occurs only once.
+     *
+     * @return whether an occurrence of this custom definition should always be replaced by a reference
+     *
+     * @since 4.27.0
+     */
+    public final boolean shouldAlwaysProduceDefinition() {
+        return this.definitionType == DefinitionType.ALWAYS_REF;
+    }
+
     /**
      * Getter for the associated extent of additional attributes being collected and applied.
      *
@@ -122,9 +134,16 @@ public class CustomDefinition {
         INLINE,
         /**
          * Follow the standard behaviour of being moved into the {@link SchemaKeyword#TAG_DEFINITIONS} if it is being referenced multiple times or
-         * {@link Option#DEFINITIONS_FOR_ALL_OBJECTS} is enabled.
+         * {@link Option#DEFINITIONS_FOR_ALL_OBJECTS} is enabled. Unless {@link Option#INLINE_ALL_SCHEMAS} is enabled, forcing this to be inlined.
+         */
+        STANDARD,
+        /**
+         * Always include this entry in the {@link SchemaKeyword#TAG_DEFINITIONS} even if it is just referenced once. This bypasses even
+         * {@link Option#INLINE_ALL_SCHEMAS}.
+         *
+         * @since 4.27.0
          */
-        STANDARD;
+        ALWAYS_REF;
     }
 
     /**
diff --git a/jsonschema-generator/src/main/java/com/github/victools/jsonschema/generator/Option.java b/jsonschema-generator/src/main/java/com/github/victools/jsonschema/generator/Option.java
index 8f17757..e446509 100644
--- a/jsonschema-generator/src/main/java/com/github/victools/jsonschema/generator/Option.java
+++ b/jsonschema-generator/src/main/java/com/github/victools/jsonschema/generator/Option.java
@@ -261,6 +261,9 @@ public enum Option {
      * Whether a member (field/method), having a declared type for which subtypes are being detected, should be included as standalone definition with
      * any collected member attributes assigned directly – and the subtypes only being handled as generic types – or each of its subtypes should be
      * treated as alternative sub-schema for this member (field/method) including any attributes derived from that member.
+     * <br>
+     * Warning: this should only be enabled if there are no relevant annotations (e.g. like jackson {@code @JsonTypeInfo}) on individual properties,
+     * that override the general behavior exclusively for single subtypes.
      *
      * @since 4.27.0
      */
diff --git a/jsonschema-generator/src/main/java/com/github/victools/jsonschema/generator/SchemaBuilder.java b/jsonschema-generator/src/main/java/com/github/victools/jsonschema/generator/SchemaBuilder.java
index 8d47b0d..82cbcf6 100644
--- a/jsonschema-generator/src/main/java/com/github/victools/jsonschema/generator/SchemaBuilder.java
+++ b/jsonschema-generator/src/main/java/com/github/victools/jsonschema/generator/SchemaBuilder.java
@@ -202,20 +202,8 @@ public class SchemaBuilder {
         final boolean inlineAllSchemas = this.config.shouldInlineAllSchemas();
 
         final AtomicBoolean considerOnlyDirectReferences = new AtomicBoolean(false);
-        Predicate<DefinitionKey> shouldProduceDefinition = definitionKey -> {
-            if (inlineAllSchemas) {
-                return false;
-            }
-            if (definitionKey.equals(mainSchemaKey)) {
-                return true;
-            }
-            List<ObjectNode> references = generationContext.getReferences(definitionKey);
-            if (considerOnlyDirectReferences.get() && references.isEmpty()) {
-                return false;
-            }
-            List<ObjectNode> nullableReferences = generationContext.getNullableReferences(definitionKey);
-            return createDefinitionsForAll || (references.size() + nullableReferences.size()) > 1;
-        };
+        Predicate<DefinitionKey> shouldProduceDefinition = this.getShouldProduceDefinitionCheck(mainSchemaKey, considerOnlyDirectReferences,
+                createDefinitionsForAll, inlineAllSchemas);
         Map<DefinitionKey, String> baseReferenceKeys = this.getReferenceKeys(mainSchemaKey, shouldProduceDefinition, generationContext);
         considerOnlyDirectReferences.set(true);
         final boolean createDefinitionForMainSchema = this.config.shouldCreateDefinitionForMainSchema();
@@ -250,7 +238,8 @@ public class SchemaBuilder {
                     definition = this.config.createObjectNode().put(this.config.getKeyword(SchemaKeyword.TAG_REF), referenceKey);
                 }
                 generationContext.makeNullable(definition);
-                if (!inlineAllSchemas && (createDefinitionsForAll || nullableReferences.size() > 1)) {
+                if (generationContext.shouldNeverInlineDefinition(definitionKey)
+                        || (!inlineAllSchemas && (createDefinitionsForAll || nullableReferences.size() > 1))) {
                     String nullableDefinitionName = this.definitionNamingStrategy
                             .adjustNullableName(definitionKey, definitionName, generationContext);
                     definitionsNode.set(nullableDefinitionName, definition);
@@ -265,6 +254,41 @@ public class SchemaBuilder {
         return definitionsNode;
     }
 
+    /**
+     * Produce reusable predicate for checking whether a given type should produce an entry in the {@link SchemaKeyword#TAG_DEFINITIONS} or not.
+     *
+     * @param mainSchemaKey main type to consider
+     * @param considerOnlyDirectReferences whether to ignore nullable references when determing about definition vs. inlining
+     * @param createDefinitionsForAll whether to produce definitions for all schemas by default (unless explicitly stated otherwise)
+     * @param inlineAllSchemas whether to inline all schemas by default (unless explicitly stated otherwise)
+     * @return reusable predicate
+     */
+    private Predicate<DefinitionKey> getShouldProduceDefinitionCheck(DefinitionKey mainSchemaKey, AtomicBoolean considerOnlyDirectReferences,
+            boolean createDefinitionsForAll, boolean inlineAllSchemas) {
+        return definitionKey -> {
+            if (generationContext.shouldNeverInlineDefinition(definitionKey)) {
+                // e.g. custom definition explicitly marked to always produce a definition
+                return true;
+            }
+            if (inlineAllSchemas) {
+                // global setting: always inline schemas by default
+                return false;
+            }
+            if (definitionKey.equals(mainSchemaKey)) {
+                return true;
+            }
+            List<ObjectNode> references = generationContext.getReferences(definitionKey);
+            if (considerOnlyDirectReferences.get() && references.isEmpty()) {
+                return false;
+            }
+            if (createDefinitionsForAll || references.size() > 1) {
+                return true;
+            }
+            List<ObjectNode> nullableReferences = generationContext.getNullableReferences(definitionKey);
+            return (references.size() + nullableReferences.size()) > 1;
+        };
+    }
+
     /**
      * Derive the applicable keys for the collected entries for the {@link SchemaKeyword#TAG_DEFINITIONS} in the given context.
      *
diff --git a/jsonschema-generator/src/main/java/com/github/victools/jsonschema/generator/impl/SchemaGenerationContextImpl.java b/jsonschema-generator/src/main/java/com/github/victools/jsonschema/generator/impl/SchemaGenerationContextImpl.java
index 9c6d99c..07c765c 100644
--- a/jsonschema-generator/src/main/java/com/github/victools/jsonschema/generator/impl/SchemaGenerationContextImpl.java
+++ b/jsonschema-generator/src/main/java/com/github/victools/jsonschema/generator/impl/SchemaGenerationContextImpl.java
@@ -64,6 +64,7 @@ public class SchemaGenerationContextImpl implements SchemaGenerationContext {
     private final Map<DefinitionKey, ObjectNode> definitions = new LinkedHashMap<>();
     private final Map<DefinitionKey, List<ObjectNode>> references = new HashMap<>();
     private final Map<DefinitionKey, List<ObjectNode>> nullableReferences = new HashMap<>();
+    private final Set<DefinitionKey> neverInlinedDefinitions  = new HashSet<>();
 
     /**
      * Constructor initialising type resolution context.
@@ -111,6 +112,25 @@ public class SchemaGenerationContextImpl implements SchemaGenerationContext {
         return this;
     }
 
+    /**
+     * Based on the given custom definition for the given type, potentially mark it as never to be inlined, i.e., that it should always be included in
+     * the {@link SchemaKeyword#TAG_DEFINITIONS}.
+     *
+     * @param customDefinition custom definition to potentially mark as never to be inlined
+     * @param javaType type to which the definition belongs
+     * @param ignoredDefinitionProvider first custom definition provider that was ignored when creating the definition (is null in most cases)
+     * @return this context (for chaining)
+     *
+     * @since 4.27.0
+     */
+    SchemaGenerationContextImpl markDefinitionAsNeverInlinedIfRequired(CustomDefinition customDefinition, ResolvedType javaType,
+            CustomDefinitionProviderV2 ignoredDefinitionProvider) {
+        if (customDefinition.shouldAlwaysProduceDefinition()) {
+            this.neverInlinedDefinitions.add(new DefinitionKey(javaType, ignoredDefinitionProvider));
+        }
+        return this;
+    }
+
     /**
      * Whether this context (already) contains a definition for the specified type, considering custom definition providers after the specified one.
      *
@@ -180,6 +200,19 @@ public class SchemaGenerationContextImpl implements SchemaGenerationContext {
         return Collections.unmodifiableList(this.nullableReferences.getOrDefault(key, Collections.emptyList()));
     }
 
+    /**
+     * Determine whether the definition for the given type should always be included in the {@link SchemaKeyword#TAG_DEFINITIONS}, even if only
+     * occurring once.
+     *
+     * @param key definition key to determine desired definition behaviour for
+     * @return whether to always produce a referenced definition for the given type
+     *
+     * @since 4.27.0
+     */
+    public boolean shouldNeverInlineDefinition(DefinitionKey key) {
+        return this.neverInlinedDefinitions.contains(key);
+    }
+
     /* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
      * Here comes the logic for traversing types and populating this context *
      * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */
@@ -293,6 +326,7 @@ public class SchemaGenerationContextImpl implements SchemaGenerationContext {
                 }
             }
             if (customDefinition != null) {
+                this.markDefinitionAsNeverInlinedIfRequired(customDefinition, targetType, ignoredDefinitionProvider);
                 logger.debug("applying configured custom definition for {}", targetType);
                 definition.setAll(customDefinition.getValue());
                 includeTypeAttributes = customDefinition.shouldIncludeAttributes();
diff --git a/jsonschema-module-jackson/src/main/java/com/github/victools/jsonschema/module/jackson/JacksonModule.java b/jsonschema-module-jackson/src/main/java/com/github/victools/jsonschema/module/jackson/JacksonModule.java
index c22b80a..6486c69 100644
--- a/jsonschema-module-jackson/src/main/java/com/github/victools/jsonschema/module/jackson/JacksonModule.java
+++ b/jsonschema-module-jackson/src/main/java/com/github/victools/jsonschema/module/jackson/JacksonModule.java
@@ -108,7 +108,7 @@ public class JacksonModule implements Module {
         boolean lookUpSubtypes = !this.options.contains(JacksonOption.SKIP_SUBTYPE_LOOKUP);
         boolean includeTypeInfoTransform = !this.options.contains(JacksonOption.IGNORE_TYPE_INFO_TRANSFORM);
         if (lookUpSubtypes || includeTypeInfoTransform) {
-            JsonSubTypesResolver subtypeResolver = new JsonSubTypesResolver();
+            JsonSubTypesResolver subtypeResolver = new JsonSubTypesResolver(this.options);
             if (lookUpSubtypes) {
                 generalConfigPart.withSubtypeResolver(subtypeResolver);
                 fieldConfigPart.withTargetTypeOverridesResolver(subtypeResolver::findTargetTypeOverrides);
diff --git a/jsonschema-module-jackson/src/main/java/com/github/victools/jsonschema/module/jackson/JacksonOption.java b/jsonschema-module-jackson/src/main/java/com/github/victools/jsonschema/module/jackson/JacksonOption.java
index 2355fc9..dbf0811 100644
--- a/jsonschema-module-jackson/src/main/java/com/github/victools/jsonschema/module/jackson/JacksonOption.java
+++ b/jsonschema-module-jackson/src/main/java/com/github/victools/jsonschema/module/jackson/JacksonOption.java
@@ -65,6 +65,13 @@ public enum JacksonOption {
      * Use this option to skip property name changes according to {@link com.fasterxml.jackson.databind.annotation.JsonNaming JsonNaming} annotations.
      */
     IGNORE_PROPERTY_NAMING_STRATEGY,
+    /**
+     * Use this option to ensure all looked-up subtypes according to {@code @JsonSubTypes} annotations are referenced via the central "$defs". This
+     * applies to the wrapping schema produced based on {@code @JsonTypeInfo} annotations, e.g., with wrapping object/array or additional property.
+     *
+     * @since 4.27.0
+     */
+    ALWAYS_REF_SUBTYPES,
     /**
      * Use this option to skip the automatic look-up of subtypes according to {@code @JsonSubTypes} annotations.
      */
diff --git a/jsonschema-module-jackson/src/main/java/com/github/victools/jsonschema/module/jackson/JsonSubTypesResolver.java b/jsonschema-module-jackson/src/main/java/com/github/victools/jsonschema/module/jackson/JsonSubTypesResolver.java
index 3194a35..ec19f64 100644
--- a/jsonschema-module-jackson/src/main/java/com/github/victools/jsonschema/module/jackson/JsonSubTypesResolver.java
+++ b/jsonschema-module-jackson/src/main/java/com/github/victools/jsonschema/module/jackson/JsonSubTypesResolver.java
@@ -33,6 +33,8 @@ import com.github.victools.jsonschema.generator.SchemaKeyword;
 import com.github.victools.jsonschema.generator.SubtypeResolver;
 import com.github.victools.jsonschema.generator.TypeContext;
 import com.github.victools.jsonschema.generator.impl.AttributeCollector;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.Collectors;
@@ -43,6 +45,29 @@ import java.util.stream.Stream;
  */
 public class JsonSubTypesResolver implements SubtypeResolver, CustomDefinitionProviderV2 {
 
+    private final CustomDefinition.DefinitionType subtypeDefinitionType;
+
+    /**
+     * Default constructor equivalent to calling {@code new JsonSubTypesResolver(Collections.emptyList())}.
+     */
+    public JsonSubTypesResolver() {
+        this(Collections.emptyList());
+        // default constructor for backward compatibility, in case this was extended by someone
+    }
+
+    /**
+     * Constructor expecting list of enabled module options.
+     * <br>
+     * Currently, only the {@link JacksonOption#ALWAYS_REF_SUBTYPES} is considered here. Other relevant options are handled by the module class.
+     *
+     * @param options module options to derive differing behavior from
+     */
+    public JsonSubTypesResolver(Collection<JacksonOption> options) {
+        this.subtypeDefinitionType = options.contains(JacksonOption.ALWAYS_REF_SUBTYPES)
+                ? CustomDefinition.DefinitionType.ALWAYS_REF
+                : CustomDefinition.DefinitionType.STANDARD;
+    }
+
     /*
      * Looking-up declared subtypes for encountered supertype in general.
      */
@@ -116,7 +141,7 @@ public class JsonSubTypesResolver implements SubtypeResolver, CustomDefinitionPr
         if (definition == null) {
             return null;
         }
-        return new CustomDefinition(definition, CustomDefinition.DefinitionType.STANDARD, CustomDefinition.AttributeInclusion.NO);
+        return new CustomDefinition(definition, this.subtypeDefinitionType, CustomDefinition.AttributeInclusion.NO);
     }
 
     /**
diff --git a/slate-docs/source/includes/_jackson-module.md b/slate-docs/source/includes/_jackson-module.md
index 68170b7..dcdf68f 100644
--- a/slate-docs/source/includes/_jackson-module.md
+++ b/slate-docs/source/includes/_jackson-module.md
@@ -34,3 +34,7 @@ SchemaGeneratorConfigBuilder configBuilder = new SchemaGeneratorConfigBuilder(Sc
 Schema attributes derived from annotations on getter methods are also applied to their associated fields.
 
 To use it, just pass a module instance into `SchemaGeneratorConfigBuilder.with(Module)`, optionally providing `JacksonOption` values in the module's constructor.
+
+----
+
+For a description of the available `JacksonOption` values and their meaning (apart from their mentionings in above enumeration), please refer to the `JacksonOption` class or its JavaDoc.
diff --git a/slate-docs/source/includes/_main-generator.md b/slate-docs/source/includes/_main-generator.md
index af2543b..0a2746d 100644
--- a/slate-docs/source/includes/_main-generator.md
+++ b/slate-docs/source/includes/_main-generator.md
@@ -287,7 +287,7 @@ configBuilder.without(
       <td colspan="2"><code>Option.DEFINITIONS_FOR_MEMBER_SUPERTYPES</code></td>
     </tr>
     <tr>
-      <td>For a member (field/method), having a declared type for which subtypes are being detected, include a single definition with any collected member attributes assigned directly. Any subtypes are only being handled as generic types, i.e., outside of the member context.</td>
+      <td>For a member (field/method), having a declared type for which subtypes are being detected, include a single definition with any collected member attributes assigned directly. Any subtypes are only being handled as generic types, i.e., outside of the member context. That means, certain relevent annotations may be ignored (e.g. a jackson <code>@JsonTypeInfo</code> override on a single member would not be correctly reflected in the produced schema).</td>
       <td>For a member (field/method), having a declared type for which subtypes are being detected, include a list of definittions: one for each subtype in the given member's context. This allows independently interpreting contextual information (e.g., member annotations) for each subtype.</td>
     </tr>
     <tr>
@@ -1057,7 +1057,7 @@ Other useful methods available in the context of a custom definition provider ar
 * `SchemaGenerationContext.getTypeContext().resolve()` allowing you to produce `ResolvedType` instances which are expected by various other methods.
 
 <aside class="notice">
-    On the <code>CustomDefinition</code>'s constructor, you are able to decide whether it should be "inlined" and
+    On the <code>CustomDefinition</code>'s constructor, you are able to decide whether it should be "inlined" or always result in a referenced/central definition and
     whether or not the attributes collected through the various other <a href="#generator-individual-configurations">Individual Configurations</a> shall be added.
 </aside>
 
@@ -1074,7 +1074,7 @@ CustomPropertyDefinitionProvider provider = (member, context) -> Optional
                 return null;
             }
         })
-        .map(CustomDefinition::new)
+        .map(CustomPropertyDefinition::new)
         .orElse(null);
 // if you don't rely on specific field/method functionality,
 // you can reuse the same provider for both of them


Test patch:
diff --git a/jsonschema-generator/src/test/java/com/github/victools/jsonschema/generator/SchemaGeneratorComplexTypesTest.java b/jsonschema-generator/src/test/java/com/github/victools/jsonschema/generator/SchemaGeneratorComplexTypesTest.java
index d22135f..82be232 100644
--- a/jsonschema-generator/src/test/java/com/github/victools/jsonschema/generator/SchemaGeneratorComplexTypesTest.java
+++ b/jsonschema-generator/src/test/java/com/github/victools/jsonschema/generator/SchemaGeneratorComplexTypesTest.java
@@ -100,12 +100,13 @@ public class SchemaGeneratorComplexTypesTest {
     }
 
     static Stream<Arguments> parametersForTestGenerateSchema() {
-        Module neutralModule = configBuilder -> configBuilder.forTypesInGeneral().withCustomDefinitionProvider((javaType, _context) -> {
+        Module neutralModule = configBuilder -> configBuilder.forTypesInGeneral().withCustomDefinitionProvider((javaType, context) -> {
             if (Integer.class == javaType.getErasedType()) {
                 ObjectNode customNode = configBuilder.getObjectMapper()
                         .createObjectNode()
+                        .put(context.getKeyword(SchemaKeyword.TAG_TYPE), context.getKeyword(SchemaKeyword.TAG_TYPE_INTEGER))
                         .put("$comment", "custom definition for Integer.class");
-                return new CustomDefinition(customNode, false);
+                return new CustomDefinition(customNode, CustomDefinition.DefinitionType.ALWAYS_REF, CustomDefinition.AttributeInclusion.YES);
             }
             return null;
         });
diff --git a/jsonschema-generator/src/test/resources/com/github/victools/jsonschema/generator/testclass1-FULL_DOCUMENTATION.json b/jsonschema-generator/src/test/resources/com/github/victools/jsonschema/generator/testclass1-FULL_DOCUMENTATION.json
index 63fbb03..fdd380a 100644
--- a/jsonschema-generator/src/test/resources/com/github/victools/jsonschema/generator/testclass1-FULL_DOCUMENTATION.json
+++ b/jsonschema-generator/src/test/resources/com/github/victools/jsonschema/generator/testclass1-FULL_DOCUMENTATION.json
@@ -1,7 +1,15 @@
 {
     "definitions": {
-        "Integer-nullable": {
+        "Integer": {
+            "type": "integer",
             "$comment": "custom definition for Integer.class"
+        },
+        "Integer-nullable": {
+            "anyOf": [{
+                    "type": "null"
+                }, {
+                    "$ref": "#/definitions/Integer"
+                }]
         }
     },
     "type": "object",
diff --git a/jsonschema-module-jackson/src/test/java/com/github/victools/jsonschema/module/jackson/SubtypeResolutionSimpleIntegrationTest.java b/jsonschema-module-jackson/src/test/java/com/github/victools/jsonschema/module/jackson/SubtypeResolutionSimpleIntegrationTest.java
index 212eb98..f749704 100644
--- a/jsonschema-module-jackson/src/test/java/com/github/victools/jsonschema/module/jackson/SubtypeResolutionSimpleIntegrationTest.java
+++ b/jsonschema-module-jackson/src/test/java/com/github/victools/jsonschema/module/jackson/SubtypeResolutionSimpleIntegrationTest.java
@@ -50,7 +50,7 @@ public class SubtypeResolutionSimpleIntegrationTest {
     public void testIntegration() throws Exception {
         SchemaGeneratorConfig config = new SchemaGeneratorConfigBuilder(SchemaVersion.DRAFT_2019_09, OptionPreset.PLAIN_JSON)
                 .with(Option.DEFINITIONS_FOR_MEMBER_SUPERTYPES)
-                .with(new JacksonModule())
+                .with(new JacksonModule(JacksonOption.ALWAYS_REF_SUBTYPES))
                 .build();
         SchemaGenerator generator = new SchemaGenerator(config);
         JsonNode result = generator.generateSchema(TestClassForSubtypeResolution.class);
diff --git a/jsonschema-module-jackson/src/test/resources/com/github/victools/jsonschema/module/jackson/subtype-simple-integration-test-result.json b/jsonschema-module-jackson/src/test/resources/com/github/victools/jsonschema/module/jackson/subtype-simple-integration-test-result.json
index 099ba4d..492c7f5 100644
--- a/jsonschema-module-jackson/src/test/resources/com/github/victools/jsonschema/module/jackson/subtype-simple-integration-test-result.json
+++ b/jsonschema-module-jackson/src/test/resources/com/github/victools/jsonschema/module/jackson/subtype-simple-integration-test-result.json
@@ -1,41 +1,47 @@
 {
     "$schema": "https://json-schema.org/draft/2019-09/schema",
     "$defs": {
+        "TestSubClassA": {
+            "allOf": [{
+                    "type": "object",
+                    "properties": {
+                        "aProperty": {
+                            "type": "string"
+                        }
+                    }
+                }, {
+                    "type": "object",
+                    "properties": {
+                        "@type": {
+                            "const": "SubClassA"
+                        }
+                    },
+                    "required": ["@type"]
+                }]
+        },
+        "TestSubClassB": {
+            "allOf": [{
+                    "type": "object",
+                    "properties": {
+                        "bProperty": {
+                            "type": "integer"
+                        }
+                    }
+                }, {
+                    "type": "object",
+                    "properties": {
+                        "@type": {
+                            "const": "SubClassB"
+                        }
+                    },
+                    "required": ["@type"]
+                }]
+        },
         "TestSuperClass": {
             "anyOf": [{
-                    "allOf": [{
-                            "type": "object",
-                            "properties": {
-                                "aProperty": {
-                                    "type": "string"
-                                }
-                            }
-                        }, {
-                            "type": "object",
-                            "properties": {
-                                "@type": {
-                                    "const": "SubClassA"
-                                }
-                            },
-                            "required": ["@type"]
-                        }]
+                    "$ref": "#/$defs/TestSubClassA"
                 }, {
-                    "allOf": [{
-                            "type": "object",
-                            "properties": {
-                                "bProperty": {
-                                    "type": "integer"
-                                }
-                            }
-                        }, {
-                            "type": "object",
-                            "properties": {
-                                "@type": {
-                                    "const": "SubClassB"
-                                }
-                            },
-                            "required": ["@type"]
-                        }]
+                    "$ref": "#/$defs/TestSubClassB"
                 }]
         }
     },