diff --git a/src/setup_payload/java/SetupPayloadParser-JNI.cpp b/src/setup_payload/java/SetupPayloadParser-JNI.cpp
index ffb4bc7..dd19771 100644
--- a/src/setup_payload/java/SetupPayloadParser-JNI.cpp
+++ b/src/setup_payload/java/SetupPayloadParser-JNI.cpp
@@ -113,6 +113,7 @@
     jfieldID productId             = env->GetFieldID(setupPayloadClass, "productId", "I");
     jfieldID commissioningFlow     = env->GetFieldID(setupPayloadClass, "commissioningFlow", "I");
     jfieldID discriminator         = env->GetFieldID(setupPayloadClass, "discriminator", "I");
+    jfieldID hasShortDiscriminator = env->GetFieldID(setupPayloadClass, "hasShortDiscriminator", "Z");
     jfieldID setUpPinCode          = env->GetFieldID(setupPayloadClass, "setupPinCode", "J");
     jfieldID discoveryCapabilities = env->GetFieldID(setupPayloadClass, "discoveryCapabilities", "Ljava/util/Set;");
 
@@ -120,11 +121,9 @@
     env->SetIntField(setupPayload, vendorId, payload.vendorID);
     env->SetIntField(setupPayload, productId, payload.productID);
     env->SetIntField(setupPayload, commissioningFlow, static_cast<int>(payload.commissioningFlow));
-    // TODO: The API we have here does not handle short discriminators in any
-    // sane way.  Just do what we used to do, which is pretend that a short
-    // discriminator is actually a long discriminator with the low bits all 0.
     uint16_t discriminatorValue;
-    if (payload.discriminator.IsShortDiscriminator())
+    bool isShortDiscriminator = payload.discriminator.IsShortDiscriminator();
+    if (isShortDiscriminator)
     {
         discriminatorValue = static_cast<uint16_t>(payload.discriminator.GetShortValue())
             << (SetupDiscriminator::kLongBits - SetupDiscriminator::kShortBits);
@@ -134,6 +133,7 @@
         discriminatorValue = payload.discriminator.GetLongValue();
     }
     env->SetIntField(setupPayload, discriminator, discriminatorValue);
+    env->SetBooleanField(setupPayload, hasShortDiscriminator, isShortDiscriminator);
     env->SetLongField(setupPayload, setUpPinCode, payload.setUpPINCode);
 
     env->SetObjectField(setupPayload, discoveryCapabilities,
@@ -275,19 +275,28 @@
 {
     jclass setupPayloadClass = env->FindClass("chip/setuppayload/SetupPayload");
 
-    jfieldID version               = env->GetFieldID(setupPayloadClass, "version", "I");
-    jfieldID vendorId              = env->GetFieldID(setupPayloadClass, "vendorId", "I");
-    jfieldID productId             = env->GetFieldID(setupPayloadClass, "productId", "I");
-    jfieldID commissioningFlow     = env->GetFieldID(setupPayloadClass, "commissioningFlow", "I");
-    jfieldID discriminator         = env->GetFieldID(setupPayloadClass, "discriminator", "I");
-    jfieldID setUpPinCode          = env->GetFieldID(setupPayloadClass, "setupPinCode", "J");
-    jfieldID discoveryCapabilities = env->GetFieldID(setupPayloadClass, "discoveryCapabilities", "Ljava/util/Set;");
+    jfieldID version                      = env->GetFieldID(setupPayloadClass, "version", "I");
+    jfieldID vendorId                     = env->GetFieldID(setupPayloadClass, "vendorId", "I");
+    jfieldID productId                    = env->GetFieldID(setupPayloadClass, "productId", "I");
+    jfieldID commissioningFlow            = env->GetFieldID(setupPayloadClass, "commissioningFlow", "I");
+    jfieldID discriminator                = env->GetFieldID(setupPayloadClass, "discriminator", "I");
+    jfieldID hasShortDiscriminatorFieldId = env->GetFieldID(setupPayloadClass, "hasShortDiscriminator", "Z");
+    jfieldID setUpPinCode                 = env->GetFieldID(setupPayloadClass, "setupPinCode", "J");
+    jfieldID discoveryCapabilities        = env->GetFieldID(setupPayloadClass, "discoveryCapabilities", "Ljava/util/Set;");
 
-    payload.version           = env->GetIntField(jPayload, version);
-    payload.vendorID          = env->GetIntField(jPayload, vendorId);
-    payload.productID         = env->GetIntField(jPayload, productId);
-    payload.commissioningFlow = static_cast<CommissioningFlow>(env->GetIntField(jPayload, commissioningFlow));
-    payload.discriminator.SetLongValue(env->GetIntField(jPayload, discriminator));
+    payload.version                = env->GetIntField(jPayload, version);
+    payload.vendorID               = env->GetIntField(jPayload, vendorId);
+    payload.productID              = env->GetIntField(jPayload, productId);
+    payload.commissioningFlow      = static_cast<CommissioningFlow>(env->GetIntField(jPayload, commissioningFlow));
+    jboolean hasShortDiscriminator = env->GetBooleanField(jPayload, hasShortDiscriminatorFieldId);
+    if (hasShortDiscriminator)
+    {
+        payload.discriminator.SetShortValue(env->GetShortField(jPayload, discriminator));
+    }
+    else
+    {
+        payload.discriminator.SetLongValue(env->GetIntField(jPayload, discriminator));
+    }
     payload.setUpPINCode = static_cast<uint32_t>(env->GetLongField(jPayload, setUpPinCode));
 
     jobject discoveryCapabilitiesObj = env->GetObjectField(jPayload, discoveryCapabilities);
diff --git a/src/setup_payload/java/src/chip/setuppayload/SetupPayload.java b/src/setup_payload/java/src/chip/setuppayload/SetupPayload.java
index 6c3f4ab..18723cc 100644
--- a/src/setup_payload/java/src/chip/setuppayload/SetupPayload.java
+++ b/src/setup_payload/java/src/chip/setuppayload/SetupPayload.java
@@ -21,6 +21,12 @@
   /** The CHIP device discriminator: */
   public int discriminator;
   /**
+   * If hasShortDiscriminator is true, the discriminator value contains just the high 4 bits of the
+   * full discriminator. For example, if hasShortDiscriminator is true and discriminator is 0xA,
+   * then the full discriminator can be anything in the range 0xA00 t0 0xAFF.
+   */
+  public boolean hasShortDiscriminator;
+  /**
    * The CHIP device setup PIN code: setupPINCode SHALL be greater than 0. Also invalid setupPINCode
    * is {000000000, 11111111, 22222222, 33333333, 44444444, 55555555, 66666666, 77777777, 88888888,
    * 99999999, 12345678, 87654321}.
@@ -41,12 +47,33 @@
       Set<DiscoveryCapability> discoveryCapabilities,
       int discriminator,
       long setupPinCode) {
+    this(
+        version,
+        vendorId,
+        productId,
+        commissioningFlow,
+        discoveryCapabilities,
+        discriminator,
+        false,
+        setupPinCode);
+  }
+
+  public SetupPayload(
+      int version,
+      int vendorId,
+      int productId,
+      int commissioningFlow,
+      Set<DiscoveryCapability> discoveryCapabilities,
+      int discriminator,
+      boolean hasShortDiscriminator,
+      long setupPinCode) {
     this.version = version;
     this.vendorId = vendorId;
     this.productId = productId;
     this.commissioningFlow = commissioningFlow;
     this.discoveryCapabilities = discoveryCapabilities;
     this.discriminator = discriminator;
+    this.hasShortDiscriminator = hasShortDiscriminator;
     this.setupPinCode = setupPinCode;
     this.optionalQRCodeInfo = new HashMap<Integer, OptionalQRCodeInfo>();
   }