diff --git a/internal/grpcwrap/provider.go b/internal/grpcwrap/provider.go
index dd2c34dcdbf6..44ed3f8d6961 100644
--- a/internal/grpcwrap/provider.go
+++ b/internal/grpcwrap/provider.go
@@ -558,11 +558,11 @@ func (p *provider) GetResourceIdentitySchemas(_ context.Context, req *tfplugin5.
 
 func (p *provider) UpgradeResourceIdentity(_ context.Context, req *tfplugin5.UpgradeResourceIdentity_Request) (*tfplugin5.UpgradeResourceIdentity_Response, error) {
 	resp := &tfplugin5.UpgradeResourceIdentity_Response{}
-	resource, ok := p.schema.ResourceTypes[req.TypeName]
+	resource, ok := p.identitySchemas.IdentityTypes[req.TypeName]
 	if !ok {
 		return nil, fmt.Errorf("resource identity schema not found for type %q", req.TypeName)
 	}
-	ty := resource.Identity.ImpliedType()
+	ty := resource.Body.ImpliedType()
 	upgradeResp := p.provider.UpgradeResourceIdentity(providers.UpgradeResourceIdentityRequest{
 		TypeName:        req.TypeName,
 		Version:         req.Version,
diff --git a/internal/grpcwrap/provider6.go b/internal/grpcwrap/provider6.go
index 77cf65e731b0..457334dbcd56 100644
--- a/internal/grpcwrap/provider6.go
+++ b/internal/grpcwrap/provider6.go
@@ -608,11 +608,11 @@ func (p *provider6) GetResourceIdentitySchemas(_ context.Context, req *tfplugin6
 
 func (p *provider6) UpgradeResourceIdentity(_ context.Context, req *tfplugin6.UpgradeResourceIdentity_Request) (*tfplugin6.UpgradeResourceIdentity_Response, error) {
 	resp := &tfplugin6.UpgradeResourceIdentity_Response{}
-	resource, ok := p.schema.ResourceTypes[req.TypeName]
+	resource, ok := p.identitySchemas.IdentityTypes[req.TypeName]
 	if !ok {
 		return nil, fmt.Errorf("resource identity schema not found for type %q", req.TypeName)
 	}
-	ty := resource.Identity.ImpliedType()
+	ty := resource.Body.ImpliedType()
 
 	upgradeResp := p.provider.UpgradeResourceIdentity(providers.UpgradeResourceIdentityRequest{
 		TypeName:        req.TypeName,