diff --git a/src/EFCore.Cosmos/Extensions/CosmosPropertyExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosPropertyExtensions.cs
index f01bde23c05..aaec1e162b9 100644
--- a/src/EFCore.Cosmos/Extensions/CosmosPropertyExtensions.cs
+++ b/src/EFCore.Cosmos/Extensions/CosmosPropertyExtensions.cs
@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System;
 using System.Linq;
 using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal;
 using Microsoft.EntityFrameworkCore.Metadata;
@@ -17,6 +18,9 @@ namespace Microsoft.EntityFrameworkCore
     /// </remarks>
     public static class CosmosPropertyExtensions
     {
+        private static readonly bool _useOldBehavior31664 =
+            AppContext.TryGetSwitch("Microsoft.EntityFrameworkCore.Issue31664", out var enabled31664) && enabled31664;
+
         /// <summary>
         ///     Returns the property name that the property is mapped to when targeting Cosmos.
         /// </summary>
@@ -37,6 +41,7 @@ private static string GetDefaultJsonPropertyName(IReadOnlyProperty property)
                 var pk = property.FindContainingPrimaryKey();
                 if (pk != null
                     && (property.ClrType == typeof(int) || ownership.Properties.Contains(property))
+                    && (property.IsShadowProperty() || _useOldBehavior31664)
                     && pk.Properties.Count == ownership.Properties.Count + (ownership.IsUnique ? 0 : 1)
                     && ownership.Properties.All(fkProperty => pk.Properties.Contains(fkProperty)))
                 {
diff --git a/src/EFCore.Cosmos/Metadata/Conventions/CosmosValueGenerationConvention.cs b/src/EFCore.Cosmos/Metadata/Conventions/CosmosValueGenerationConvention.cs
index c21238118c9..29d1b94feb1 100644
--- a/src/EFCore.Cosmos/Metadata/Conventions/CosmosValueGenerationConvention.cs
+++ b/src/EFCore.Cosmos/Metadata/Conventions/CosmosValueGenerationConvention.cs
@@ -21,6 +21,9 @@ public class CosmosValueGenerationConvention :
         ValueGenerationConvention,
         IEntityTypeAnnotationChangedConvention
     {
+        private static readonly bool _useOldBehavior31664 =
+            AppContext.TryGetSwitch("Microsoft.EntityFrameworkCore.Issue31664", out var enabled31664) && enabled31664;
+
         /// <summary>
         ///     Creates a new instance of <see cref="CosmosValueGenerationConvention" />.
         /// </summary>
@@ -82,8 +85,9 @@ public virtual void ProcessEntityTypeAnnotationChanged(
                 {
                     var pk = property.FindContainingPrimaryKey();
                     if (pk != null
-                        && !ownership.Properties.Contains(property)
+                        && !property.IsForeignKey()
                         && pk.Properties.Count == ownership.Properties.Count + 1
+                        && (property.IsShadowProperty() || _useOldBehavior31664)
                         && ownership.Properties.All(fkProperty => pk.Properties.Contains(fkProperty)))
                     {
                         return base.GetValueGenerated(property);
diff --git a/src/EFCore.Cosmos/Metadata/Internal/CosmosPropertyExtensions.cs b/src/EFCore.Cosmos/Metadata/Internal/CosmosPropertyExtensions.cs
index ecaf3357681..acf365b4e08 100644
--- a/src/EFCore.Cosmos/Metadata/Internal/CosmosPropertyExtensions.cs
+++ b/src/EFCore.Cosmos/Metadata/Internal/CosmosPropertyExtensions.cs
@@ -24,13 +24,13 @@ public static bool IsOrdinalKeyProperty(this IReadOnlyProperty property)
         {
             Check.DebugAssert(
                 property.DeclaringEntityType.IsOwned(), $"Expected {property.DeclaringEntityType.DisplayName()} to be owned.");
-            Check.DebugAssert(property.GetJsonPropertyName().Length == 0, $"Expected {property.Name} to be non-persisted.");
 
             return property.FindContainingPrimaryKey() is IReadOnlyKey key
                 && key.Properties.Count > 1
                 && !property.IsForeignKey()
                 && property.ClrType == typeof(int)
-                && property.ValueGenerated == ValueGenerated.OnAdd;
+                && property.ValueGenerated == ValueGenerated.OnAdd
+                && property.GetJsonPropertyName().Length == 0;
         }
     }
 }
diff --git a/test/EFCore.SqlServer.FunctionalTests/SqlServerEndToEndTest.cs b/test/EFCore.SqlServer.FunctionalTests/SqlServerEndToEndTest.cs
index ee455da0fab..36fbe325949 100644
--- a/test/EFCore.SqlServer.FunctionalTests/SqlServerEndToEndTest.cs
+++ b/test/EFCore.SqlServer.FunctionalTests/SqlServerEndToEndTest.cs
@@ -8,7 +8,6 @@
 using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Threading.Tasks;
-using Microsoft.EntityFrameworkCore.Metadata;
 using Microsoft.EntityFrameworkCore.TestUtilities;
 using Xunit;